import numpy as np  # NOQA: E402
import pandas as pd
from sorted_nearest import k_nearest_next_nonoverlapping, k_nearest_previous_nonoverlapping  # type: ignore


def nearest_previous_idx(d1, d2, k, ties=None):
    d1s = d1.Start.sort_values()
    d2e = d2.End.sort_values()

    # ix where d1s could be inserted into d2e to preserve order of d2e
    # start = time()
    ix = np.searchsorted(d2e, d1s, side="right") - 1
    # end = time()
    # print("np previous", end - start)

    valid = ix >= 0
    ix = ix[valid]

    d1s = d1s.iloc[valid]

    lidx, ridx_pos, dist = k_nearest_previous_nonoverlapping(d1s.values, d2e.values, d1s.index.values, ix, k, ties)

    dist += 1

    ridx = d2e.iloc[ridx_pos].index

    return lidx, ridx, dist


# @profile
def nearest_next_idx(d1, d2, k, ties=None):
    d1e = d1.End.sort_values()
    d2s = d2.Start.sort_values()

    # ix where d1e could be inserted into d2s to preserve order of d2s
    # start = time()
    ix = np.searchsorted(d2s, d1e, side="left")
    # end = time()
    # print("np next", end - start)

    # - print(ix)
    valid = ix < len(d2s)
    ix = ix[valid]

    d1e = d1e.iloc[valid]

    lidx, ridx_pos, dist = k_nearest_next_nonoverlapping(d1e.values, d2s.values, d1e.index.values, ix, k, ties)

    dist += 1

    ridx = d2s.iloc[ridx_pos].index

    return lidx, ridx, dist


def nearest(d1, d2, **kwargs):
    suffix = kwargs.get("suffix", "_b")
    ties = kwargs.get("ties", None)

    plidx, pridx, pdist = nearest_previous_idx(d1, d2, d1.__k__.values, ties)
    nlidx, nridx, ndist = nearest_next_idx(d1, d2, d1.__k__.values, ties)

    pk = d1.__k__.reindex(plidx)
    nk = d1.__k__.reindex(nlidx)
    # px = d1.__IX__.reindex(plidx)
    # nx = d1.__IX__.reindex(nlidx)

    p = pd.DataFrame({"LX": plidx, "RX": pridx, "D": pdist, "PN": "P", "k": pk})
    n = pd.DataFrame({"LX": nlidx, "RX": nridx, "D": ndist, "PN": "N", "k": nk})

    df = pd.concat([p, n]).sort_values(["LX", "D"])
    # print(df)
    # raise

    if df.empty:
        return None

    xdf = df

    # row_indexer = xdf.PN == "P"
    # xdf.loc[row_indexer, "D"] = -xdf[row_indexer].D

    d1 = d1.reindex(xdf.LX)
    d2 = d2.reindex(xdf.RX)
    d1.index = range(len(d1))
    d2.index = range(len(d1))
    d2 = d2.drop("Chromosome", axis=1)
    df = d1.join(d2, rsuffix=suffix)
    df.insert(df.shape[1], "Distance", xdf.D.values)

    # to_drop = [c for c in df.columns if "__k__" in c]

    # df = df.drop(to_drop, axis=1)

    return df


def nearest_previous(d1, d2, **kwargs):
    suffix = kwargs.get("suffix", "_b")
    ties = kwargs.get("ties", None)

    lidx, ridx, dist = nearest_previous_idx(d1, d2, d1.__k__.values, ties)

    d1 = d1.reindex(lidx)
    d2 = d2.reindex(ridx)
    d1.index = range(len(d1))
    d2.index = range(len(d1))

    d2 = d2.drop("Chromosome", axis=1)
    df = d1.join(d2, rsuffix=suffix)
    df.insert(df.shape[1], "Distance", dist)

    # df = remove_duplicates_single(df, ties)

    # df.loc[:, "Distance"] = - df.Distance
    # to_drop = [c for c in df.columns if "__k__" in c]
    # df = df.drop(to_drop, axis=1)

    return df


def nearest_next(d1, d2, **kwargs):
    suffix = kwargs.get("suffix", "_b")
    ties = kwargs.get("ties", None)

    lidx, ridx, dist = nearest_next_idx(d1, d2, d1.__k__.values, ties)

    d1 = d1.reindex(lidx)
    d2 = d2.reindex(ridx)
    d2 = d2.drop("Chromosome", axis=1)
    d1.index = range(len(d1))
    d2.index = range(len(d1))
    df = d1.join(d2, rsuffix=suffix)
    df.insert(df.shape[1], "Distance", dist)

    return df


# @profile
def _nearest(d1, d2, **kwargs):
    if d1.empty or d2.empty:
        return None

    how = kwargs["how"]

    # if kwargs.get("strandedness") == "opposite":
    #     strand_dict = {"+": "-", "-": "+"}
    # else:
    #     strand_dict = {"+": "+", "-": "-"}

    if how in ["upstream", "downstream"] and kwargs["stranded"]:
        strand = d1.Strand.iloc[0]
        # strand = strand_dict[strand]

        __nearest = {
            ("+", "upstream"): nearest_previous,
            ("-", "upstream"): nearest_next,
            ("+", "downstream"): nearest_next,
            ("-", "downstream"): nearest_previous,
        }[strand, how]
    elif how in ["upstream", "downstream"] and not kwargs["stranded"]:
        __nearest = {"upstream": nearest_previous, "downstream": nearest_next}[how]
    else:
        __nearest = nearest

    df = __nearest(d1, d2, **kwargs)

    return df
