from __future__ import annotations

import numpy as np
import pandas as pd

from . import ops
from .core.specs import _get_default_colnames, _verify_columns

__all__ = [
    "make_chromarms",
    "binnify",
    "digest",
    "frac_mapped",
    "frac_gc",
    "seq_gc",
    "frac_gene_coverage",
    "pair_by_distance",
    "mark_runs",
    "merge_runs"
]


def make_chromarms(
    chromsizes,
    midpoints,
    cols_chroms=("chrom", "length"),
    cols_mids=("chrom", "mid"),
    suffixes=("_p", "_q"),
):
    """
    Split chromosomes into chromosome arms.

    Parameters
    ----------
    chromsizes : pandas.Dataframe or dict-like
        If dict or pandas.Series, a map from chromosomes to lengths in bp.
        If pandas.Dataframe, a dataframe with columns defined by cols_chroms.
        If cols_chroms is a triplet (e.g. 'chrom','start','end'), then
        values in chromsizes[cols_chroms[1]].values must all be zero.

    midpoints : pandas.Dataframe or dict-like
        Mapping of chromosomes to midpoint (aka centromere) locations.
        If dict or pandas.Series, a map from chromosomes to midpoints in bp.
        If pandas.Dataframe, a dataframe with columns defined by cols_mids.

    cols_chroms : (str, str) or (str, str, str)
        Two columns

    suffixes : tuple, optional
        Suffixes to name chromosome arms. Defaults to p and q.

    Returns
    -------
    df_chromarms
        4-column BED-like DataFrame (chrom, start, end, name).
        Arm names are chromosome names + suffix.
        Any chromosome not included in ``mids`` will be not be split.

    """
    columns_to_drop = ["index", "sub_index_"]
    if len(cols_chroms) == 2:
        ck1, sk1 = cols_chroms
    elif len(cols_chroms) == 3:
        ck1, sk1, ek1 = cols_chroms

    if isinstance(chromsizes, (pd.Series, dict)):
        chromsizes = dict(chromsizes)
        df_chroms = pd.DataFrame(
            {
                ck1: list(chromsizes.keys()),
                "length": list(chromsizes.values()),
            }
        )
    elif isinstance(chromsizes, pd.DataFrame):
        df_chroms = chromsizes.copy()
    else:
        raise ValueError("unknown input type for chromsizes")

    if len(cols_chroms) == 2:
        _verify_columns(df_chroms, [ck1, sk1])
        columns_to_drop += [sk1]
        df_chroms["end"] = df_chroms[sk1].values
        df_chroms["start"] = 0
        sk1, ek1 = "start", "end"
    elif len(cols_chroms) == 3:
        ck1, sk1, ek1 = cols_chroms
        _verify_columns(df_chroms, [ck1, sk1, ek1], unique_cols=True)
        if any(df_chroms[sk1].values != 0):
            raise ValueError("all values in starts column must be zero")
    else:
        raise ValueError("invalid number of cols_chroms")

    ck2, sk2 = cols_mids
    if isinstance(midpoints, (pd.Series, dict)):
        midpoints = dict(midpoints)
        df_mids = pd.DataFrame.from_dict(midpoints, orient="index", columns=[sk2])
        df_mids.reset_index(inplace=True)
        df_mids.rename(columns={"index": ck2}, inplace=True)
    elif isinstance(midpoints, pd.DataFrame):
        df_mids = midpoints.copy()
    else:
        raise ValueError("unknown input type for midpoints")
    _verify_columns(df_mids, [ck2, sk2])
    df_mids["start"] = df_mids[sk2]
    df_mids["end"] = df_mids[sk2]

    df_chromarms = ops.subtract(
        df_chroms,
        df_mids,
        cols1=(ck1, sk1, ek1),
        cols2=(ck2, "start", "end"),
        return_index=True,
    )
    if df_chromarms["sub_index_"].max() > 1:
        raise ValueError(
            "chromosome split into more than two arms, double-check midpoints"
        )
    df_chromarms["name"] = df_chromarms[ck1] + [
        suffixes[i] for i in df_chromarms["sub_index_"].values
    ]
    # df_chromarms.drop(columns=columns_to_drop, inplace=True)
    return df_chromarms[[ck1, sk1, ek1, "name"]]


def binnify(chromsizes, binsize, rel_ids=False):
    """
    Divide a genome into evenly sized bins.

    Parameters
    ----------
    chromsizes : Series
        pandas Series indexed by chromosome name with chromosome lengths in bp.

    binsize : int
        size of bins in bp

    Returns
    -------
    bintable : pandas.DataFrame with columns: 'chrom', 'start', 'end'.

    """

    if not isinstance(binsize, int):
        raise ValueError("binsize must be int")

    def _each(chrom):
        clen = chromsizes[chrom]
        n_bins = int(np.ceil(clen / binsize))
        binedges = np.arange(0, (n_bins + 1)) * binsize
        binedges[-1] = clen
        return pd.DataFrame(
            {"chrom": [chrom] * n_bins, "start": binedges[:-1], "end": binedges[1:]},
            columns=["chrom", "start", "end"],
        )

    bintable = pd.concat(map(_each, chromsizes.keys()), axis=0, ignore_index=True)

    if rel_ids:
        bintable["rel_id"] = bintable.groupby("chrom").cumcount()

    # if as_cat:
    #     bintable['chrom'] = pd.Categorical(
    #         bintable['chrom'],
    #         categories=list(chromsizes.keys()),
    #         ordered=True)

    return bintable


def digest(fasta_records, enzyme):
    """
    Divide a genome into restriction fragments.

    Parameters
    ----------
    fasta_records : OrderedDict
        Dictionary of chromosome names to sequence records.
        Created by: bioframe.load_fasta('/path/to/fasta.fa')

    enzyme: str
        Name of restriction enzyme.

    Returns
    -------
    Dataframe with columns: 'chrom', 'start', 'end'.

    """
    try:
        import Bio.Restriction as biorst
        import Bio.Seq as bioseq
    except ImportError:
        raise ImportError("Biopython is required to use digest") from None

    # http://biopython.org/DIST/docs/cookbook/Restriction.html#mozTocId447698
    if not isinstance(fasta_records, dict):
        raise ValueError(
            "fasta records must be provided as an OrderedDict, can be created "
            "by bioframe.load_fasta"
        )
    chroms = fasta_records.keys()
    try:
        cut_finder = getattr(biorst, enzyme).search
    except AttributeError as e:
        raise ValueError(f"Unknown enzyme name: {enzyme}") from e

    def _each(chrom):
        seq = bioseq.Seq(str(fasta_records[chrom][:]))
        cuts = np.r_[0, np.array(cut_finder(seq)) + 1, len(seq)].astype(np.int64)
        n_frags = len(cuts) - 1

        frags = pd.DataFrame(
            {"chrom": [chrom] * n_frags, "start": cuts[:-1], "end": cuts[1:]},
            columns=["chrom", "start", "end"],
        )
        return frags

    return pd.concat(map(_each, chroms), axis=0, ignore_index=True)


def frac_mapped(df, fasta_records, return_input=True):
    """
    Calculate the fraction of mapped base-pairs for each interval in a dataframe.

    Parameters
    ----------
    df : pandas.DataFrame
        A sets of genomic intervals stored as a DataFrame.

    fasta_records : OrderedDict
        Dictionary of chromosome names to sequence records.
        Created by: bioframe.load_fasta('/path/to/fasta.fa')

    return_input: bool
        if False, only return Series named frac_mapped.

    Returns
    -------
    df_mapped : pd.DataFrame
        Original dataframe with new column 'frac_mapped' appended.

    """

    if not set(df["chrom"].values).issubset(set(fasta_records.keys())):
        raise ValueError(
            "chrom from intervals not in fasta_records: "
            "double-check genome agreement"
        )
    if not isinstance(fasta_records, dict):
        raise ValueError(
            "fasta records must be provided as an OrderedDict, can be created "
            "by bioframe.load_fasta"
        )

    def _each(bin):
        s = str(fasta_records[bin.chrom][bin.start : bin.end])
        nbases = len(s)
        n = s.count("N")
        n += s.count("n")
        return (nbases - n) / nbases if nbases > 0 else 0

    if return_input:
        return pd.concat(
            [df, df.apply(_each, axis=1).rename("frac_mapped", inplace=True)],
            axis="columns",
        )
    else:
        return df.apply(_each, axis=1).rename("frac_mapped", inplace=True)


def frac_gc(df, fasta_records, mapped_only=True, return_input=True):
    """
    Calculate the fraction of GC basepairs for each interval in a dataframe.

    Parameters
    ----------
    df : pandas.DataFrame
        A sets of genomic intervals stored as a DataFrame.

    fasta_records : OrderedDict
        Dictionary of chromosome names to sequence records.
        Created by: bioframe.load_fasta('/path/to/fasta.fa')

    mapped_only: bool
        if True, ignore 'N' in the fasta_records for calculation.
        if True and there are no mapped base-pairs in an interval, return np.nan.

    return_input: bool
        if False, only return Series named frac_mapped.

    Returns
    -------
    df_mapped : pd.DataFrame
        Original dataframe with new column 'GC' appended.

    """
    if not set(df["chrom"].values).issubset(set(fasta_records.keys())):
        raise ValueError(
            "chrom from intervals not in fasta_records: double-check genome agreement"
        )
    if not isinstance(fasta_records, dict):
        raise ValueError(
            "fasta records must be provided as an OrderedDict, can be created "
            "by bioframe.load_fasta"
        )

    def _each(chrom_group):
        chrom = chrom_group.name
        seq = fasta_records[chrom]
        seq = str(seq[:])
        gc = []
        for _, bin in chrom_group.iterrows():
            s = seq[bin["start"] : bin["end"]]
            gc.append(seq_gc(s, mapped_only=mapped_only))
        return gc

    agg = df.groupby("chrom", sort=False)[["start", "end"]].apply(_each)
    out_col = pd.Series(data=np.concatenate(agg.values), index=df.index).rename("GC")

    if return_input:
        return pd.concat([df, out_col], axis="columns")
    else:
        return out_col


def seq_gc(seq, mapped_only=True):
    """
    Calculate the fraction of GC basepairs for a string of nucleotides.

    Parameters
    ----------
    seq : str
        Basepair input

    mapped_only: bool
        if True, ignore 'N' in the sequence for calculation.
        if True and there are no mapped base-pairs, return np.nan.

    Returns
    -------
    gc : float
        calculated gc content.

    """
    if not isinstance(seq, str):
        raise ValueError("reformat input sequence as a str")
    g = seq.count("G")
    g += seq.count("g")
    c = seq.count("C")
    c += seq.count("c")
    nbases = len(seq)
    if mapped_only:
        n = seq.count("N")
        n += seq.count("n")
        nbases -= n
    return (g + c) / nbases if nbases > 0 else np.nan


def frac_gene_coverage(df, ucsc_mrna):
    """
    Calculate number and fraction of overlaps by predicted and verified
    RNA isoforms for a set of intervals stored in a dataframe.

    Parameters
    ----------
    df : pd.DataFrame
        Set of genomic intervals stored as a dataframe.

    ucsc_mrna: str or DataFrame
        Name of UCSC genome or all_mrna.txt dataframe from UCSC or similar.

    Returns
    -------
    df_gene_coverage : pd.DataFrame

    """
    if isinstance(ucsc_mrna, str):
        from .io.resources import UCSCClient

        mrna = UCSCClient(ucsc_mrna).fetch_mrna()
    else:
        mrna = ucsc_mrna

    mrna = mrna.rename(columns={"tName": "chrom", "tStart": "start", "tEnd": "end"})
    df_gene_coverage = ops.coverage(df, mrna)
    df_gene_coverage = ops.count_overlaps(df_gene_coverage, mrna)

    return df_gene_coverage


def pair_by_distance(
    df,
    min_sep,
    max_sep,
    min_intervening=None,
    max_intervening=None,
    relative_to="midpoints",
    cols=None,
    return_index=False,
    keep_order=False,
    suffixes=("_1", "_2"),
):
    """
    From a dataframe of genomic intervals, find all unique pairs of intervals
    that are between ``min_sep`` and ``max_sep`` bp separated from each other.

    Parameters
    ----------
    df : pandas.DataFrame
        A BED-like dataframe.
    min_sep, max_sep : int
        Minimum and maximum separation between intervals in bp.
        Min > 0 and Max >= Min.
    min_intervening, max_intervening : int
        Minimum and maximum number of intervening intervals separating pairs.
        Min > 0 and Max >= Min.
    relative_to : str,
        Whether to calculate distances between interval "midpoints" or "endpoints".
        Default "midpoints".
    cols : (str, str, str) or None
        The names of columns containing the chromosome, start and end of the
        genomic intervals, provided separately for each set. The default
        values are 'chrom', 'start', 'end'.
    return_index : bool
        If True, return indicies of pairs as two new columns
        ('index'+suffixes[0] and 'index'+suffixes[1]). Default False.
    keep_order : bool, optional
        If True, sort the output dataframe to preserve the order
        of the intervals in df1. Default False.
        Note that it relies on sorting of index in the original dataframes,
        and will reorder the output by index.
    suffixes : (str, str), optional
        The column name suffixes for the two interval sets in the output.
        The first interval of each output pair is always upstream of the
        second.

    Returns
    -------
    pandas.DataFrame
        A BEDPE-like dataframe of paired intervals from ``df``.

    """

    # Create the copy of original dataset:
    df = df.copy()

    # Index column name
    index_col = return_index if isinstance(return_index, str) else "index"
    index_col_1 = index_col + suffixes[0]
    index_col_2 = index_col + suffixes[1]
    if return_index or keep_order:
        df.index.name = index_col

    # Get columns for pairing
    ck, sk, ek = _get_default_colnames() if cols is None else cols

    # Sort intervals by genomic coordinates
    df = df.sort_values([ck, sk, ek]).reset_index(drop=False)

    if min_sep >= max_sep:
        raise ValueError("min_sep must be less than max_sep")
    if min_sep < 0:
        raise ValueError("min_sep must be >=0")

    if min_intervening is None:
        min_intervening = 0
    if max_intervening is None:
        max_intervening = df.index.max()
    if min_intervening > max_intervening:
        raise ValueError("min_intervening must be less or equal to max_intervening")
    if min_intervening < 0:
        raise ValueError("min_intervening must be >=0")

    mids = (df[sk] + df[ek]) // 2

    # For each interval, generate a probe interval on its right
    if relative_to == "endpoints":
        print("endpoint")
        ref = df[ek]
    elif relative_to == "midpoints":
        ref = mids
    else:
        raise ValueError("relative_to must either specify 'midpoints' or 'endpoints' ")
    right_probe = (
        df[[ck, index_col]].copy() if (return_index or keep_order) else df[[ck]].copy()
    )
    right_probe[sk] = ref + min_sep // 2
    right_probe[ek] = ref + (max_sep + 1) // 2

    # For each interval, also generate a probe interval on its left
    if relative_to == "endpoints":
        ref = df[sk]
    elif relative_to == "midpoints":
        ref = mids
    else:
        raise ValueError("relative_to must either specify 'midpoints' or 'endpoints' ")
    left_probe = (
        df[[ck, index_col]].copy() if (return_index or keep_order) else df[[ck]].copy()
    )
    left_probe[sk] = ref - max_sep // 2
    left_probe[ek] = ref - (min_sep + 1) // 2

    # Intersect right-handed probes (from intervals on the left)
    # with left-handed probes (from intervals on the right)
    idxs = ops.overlap(
        right_probe,
        left_probe,
        suffixes=suffixes,
        how="inner",
        return_index=True,
        return_input=False,
    )

    # Select only the pairs that are separated by
    # at least min_intervening intervals and no more than max_intervening intervals
    idxs["intervening"] = (
        np.abs(idxs[f"index{suffixes[0]}"] - idxs[f"index{suffixes[1]}"]) - 1
    )
    idxs = idxs[
        (idxs["intervening"] <= max_intervening)
        & (idxs["intervening"] >= min_intervening)
    ]

    left_ivals = (
        df.iloc[idxs[f"index{suffixes[0]}"].values]
        .rename(columns=lambda x: x + suffixes[0])
        .reset_index(drop=True)
    )
    right_ivals = (
        df.iloc[idxs[f"index{suffixes[1]}"].values]
        .rename(columns=lambda x: x + suffixes[1])
        .reset_index(drop=True)
    )

    out_df = pd.concat([left_ivals, right_ivals], axis=1)

    if keep_order:
        out_df = out_df.sort_values([index_col_1, index_col_2])

    if not return_index:
        out_df = out_df.drop([index_col_1, index_col_2], axis=1)

    out_df.reset_index(drop=True, inplace=True)

    return out_df


def mark_runs(
    df: pd.DataFrame,
    col: str,
    *,
    allow_overlaps: bool = False,
    reset_counter: bool = True,
    run_col: str = 'run',
    cols: tuple[str, str, str] | None = None,
) -> pd.DataFrame:
    """
    Mark runs of spatially consecutive intervals sharing the same value of
    ``col``.

    Parameters
    ----------
    df : DataFrame
        A bioframe dataframe.
    col : str
        The column to mark runs of values for.
    allow_overlaps : bool, optional [default: False]
        If True, allow intervals in ``df`` to overlap. This may cause
        unexpected results.
    reset_counter : bool, optional [default: True]
        If True, reset the run counter for each chromosome.
    run_col : str, optional [default: 'run']
        The name of the column to store the run numbers in.

    Returns
    -------
    pandas.DataFrame
        A reordered copy the input dataframe with an additional column 'run'
        marking runs of values in the input column.

    Notes
    -----
    This is similar to :func:`cluster`, but only clusters intervals sharing
    the same value of ``col``.

    Examples
    --------

    >>> df = pd.DataFrame({
    ...     'chrom': ['chr1', 'chr1', 'chr1', 'chr1', 'chr1', 'chr1'],
    ...     'start': [0, 100, 200, 300, 400, 500],
    ...     'end': [100, 200, 300, 400, 500, 600],
    ...     'value': [1, 1, 1, 2, 2, 2],
    ... })

    >>> mark_runs(df, 'value')
        chrom  start  end  value  run
    0   chr1      0  100      1    0
    1   chr1    100  200      1    0
    2   chr1    200  300      1    0
    3   chr1    300  400      2    1
    4   chr1    400  500      2    1
    5   chr1    500  600      2    1

    See Also
    --------
    merge_runs
    cluster
    merge
    """
    ck, sk, ek = _get_default_colnames() if cols is None else cols

    if not allow_overlaps and len(ops.overlap(df, df)) > len(df):
        raise ValueError("Not a proper bedGraph: found overlapping intervals.")

    result = []
    n_runs = 0

    for _, group in df.groupby(ck, sort=False):
        group = group.sort_values([sk, ek])
        starts = group[sk].to_numpy()
        ends = group[ek].to_numpy()

        # Extend ends by running max
        ends = np.maximum.accumulate(ends)

        # Find borders of interval clusters and assign cluster ids
        is_cluster_border = np.r_[True, starts[1:] > ends[:-1], False]

        # Find borders of consecutive equal values
        values = group[col].to_numpy()
        if values.dtype.kind == 'f':
            is_value_border = np.r_[
                True,
                ~np.isclose(values[1:], values[:-1], equal_nan=True),
                False
            ]
        else:
            is_value_border = np.r_[True, values[1:] != values[:-1], False]

        # Find index extents of runs
        is_border = is_cluster_border | is_value_border
        sum_borders = np.cumsum(is_border)
        run_ids = sum_borders[:-1] - 1

        # Assign run numbers to intervals
        if reset_counter:
            n_runs = 0
        group[run_col] = n_runs + run_ids
        n_runs += sum_borders[-1]

        result.append(group)

    return pd.concat(result)


def merge_runs(
    df: pd.DataFrame,
    col: str,
    *,
    allow_overlaps: bool = False,
    agg: dict | None = None,
    cols: tuple[str, str, str] | None = None,
) -> pd.DataFrame:
    """
    Merge runs of spatially consecutive intervals sharing the same value of
    ``col``.

    Parameters
    ----------
    df : DataFrame
        A bioframe dataframe.
    col : str
        The column to compress runs of values for.
    allow_overlaps : bool, optional [default: False]
        If True, allow intervals in ``df`` to overlap. This may cause
        unexpected results.
    agg : dict, optional [default: None]
        A dictionary of additional column names and aggregation functions to
        apply to each run. Takes the format:
            {'agg_name': ('column_name', 'agg_func')}

    Returns
    -------
    pandas.DataFrame
        Dataframe with consecutive intervals in the same run merged.

    Notes
    -----
    This is similar to :func:`merge`, but only merges intervals sharing
    the same value of ``col``.

    Examples
    --------
    >>> df = pd.DataFrame({
    ...     'chrom': ['chr1', 'chr1', 'chr1', 'chr1', 'chr1', 'chr1'],
    ...     'start': [0, 100, 200, 300, 400, 500],
    ...     'end': [100, 200, 300, 400, 500, 600],
    ...     'value': [1, 1, 1, 2, 2, 2],
    ... })

    >>> merge_runs(df, 'value')
        chrom  start  end  value
    0   chr1      0  300      1
    1   chr1    300  600      2

    >>> merge_runs(df, 'value', agg={'sum': ('value', 'sum')})
        chrom  start  end  value  sum
    0   chr1      0  300      1    3
    1   chr1    300  600      2    6

    See Also
    --------
    mark_runs
    cluster
    merge
    """
    ck, sk, ek = _get_default_colnames() if cols is None else cols

    if agg is None:
        agg = {}

    df_runs = mark_runs(
        df,
        col,
        allow_overlaps=allow_overlaps,
        reset_counter=False,
        run_col='_run',
    )
    df_merged = (
        df_runs
        .groupby('_run')
        .agg(**{
            ck: (ck, 'first'),
            sk: (sk, 'min'),
            ek: (ek, 'max'),
            col: (col, 'first'),
            **agg
         })
    )
    return df_merged.reset_index(drop=True)