# AUTOGENERATED! DO NOT EDIT! File to edit: ../02_annots.ipynb.

# %% auto 0
__all__ = [
    "get_species_name",
    "get_annots",
    "is_protein_coding",
    "get_ts",
    "get_es",
    "get_cs",
    "get_ranges",
    "intersect",
    "map_back_feat_ids",
    "intersect_with_seqs",
    "intersect_with_feats",
    "get_ts_data",
]

# %% ../02_annots.ipynb 6
## helper functions
import logging
import roux.lib.df as rd  # noqa
import pandas as pd


from roux.stat.paired import get_diff_sorted

from roux.lib.io import read_table, to_table
from pathlib import Path

try:
    from chrov.core import get_cache_dir_path
except:
    from dups.core import get_cache_dir_path


def get_species_name(species):
    if species.lower() == "hsapiens":
        species = "homo sapiens"
    return species


def get_annots(
    ensembl_release,
    species,
):
    from pyensembl import find_species_by_name
    from pyensembl import EnsemblRelease

    return EnsemblRelease(
        release=ensembl_release,
        species=find_species_by_name(get_species_name(species)),
    )


def is_protein_coding(
    t,
    strict=False,
):
    if t.biotype != "protein_coding":
        return False
    else:
        if not strict:
            return True
        else:
            return t.complete and t.contains_start_codon and t.contains_stop_codon


def get_ts(
    gene_id,
    protein_coding=True,
    strict=False,
    **kws_annots,
):
    annots = get_annots(
        **kws_annots,
    )
    g = annots.gene_by_id(gene_id)
    ## get transcripts
    ts = [t for t in g.transcripts]
    if protein_coding:
        ts = [t for t in ts if is_protein_coding(t, strict=strict)]
    # if tn is not None:
    #     ts=ts[:tn]
    return ts


def get_es(ts):
    return (
        pd.DataFrame(
            [
                {
                    "t.id": t.id,
                    "t.length": t.length,
                    "t.start": t.start,
                    "t.end": t.end,
                    "t.strand": t.strand,
                    "e.start": e.start,
                    "e.end": e.end,
                    "e.id": e.id,
                }
                for t in ts
                for e in t.exons
            ]
        )
        ## longest first
        .sort_values(["e.start"], ascending=[True])
        .assign(
            **{
                "e.length": lambda df: df.apply(
                    lambda x: get_diff_sorted(x["e.start"], x["e.end"]), axis=1
                ),
            },
        )
    )


def get_cs(
    ts,
    strict=False,
):
    return (
        pd.DataFrame(
            [
                {
                    "t.id": t.id,
                    "t.cds": t.coding_sequence_position_ranges
                    if is_protein_coding(t, strict=strict)
                    else None,
                }
                for t in ts
                for e in t.exons
            ]
        )
        .explode("t.cds")
        .set_index("t.id")["t.cds"]
        .apply(pd.Series)
        .reset_index()
        .rename(
            columns={
                0: "c.start",
                1: "c.end",
            },
            errors="raise",
        )
        .sort_values(["c.start"], ascending=[True])
        .log.drop_duplicates()
        .assign(
            **{
                "c.id": lambda df: df.apply(
                    lambda x: f"{x['c.start']}-{x['c.end']}", axis=1
                ),
            }
        )
        .log.dropna()
    )


def get_ranges(
    data: pd.DataFrame,
    col_start="start",
    col_end="end",
    convert=True,  # convert to 0-based which is used by pyranges
    # convert=False, # convert to 0-based which is used by pyranges
):
    import pyranges as pr

    df = (
        data.rename(
            columns={col_start: "Start", col_end: "End"},
            errors="raise",
        )
        .astype({"Start": int, "End": int})
        .assign(
            Chromosome="tmp",
        )
    )
    if convert:
        df = df.assign(
            Start=lambda df: df["Start"] - 1,
        )

    # Create a PyRanges object
    return pr.PyRanges(df)


def intersect(
    df1,
    df2,
    feat1_id="t.id",
    feat2_id="e.id",
    feat1_start="e.start",
    feat2_start="c.start",
    feat1_end="e.end",
    feat2_end="c.end",
    common_ids=True,
    convert=True,
    **kws_get_ranges,
):
    # print(df1,df2)
    if common_ids:
        cols1_ids = [feat1_id, feat2_id, feat1_start, feat1_end]
        cols2_ids = [feat1_id, feat2_start, feat2_end]
        df1 = df1.loc[:, cols1_ids].drop_duplicates()
        df2 = df2.loc[:, cols2_ids]
    else:
        cols1_ids = [feat1_id, feat1_start, feat1_end]
        cols2_ids = [feat2_id, feat2_start, feat2_end]
        df1 = df1.loc[:, cols1_ids].drop_duplicates()
        df2 = df2.dropna().loc[:, cols2_ids].drop_duplicates()

    gr1 = get_ranges(
        df1,
        col_start=feat1_start,
        col_end=feat1_end,
        convert=convert,
        **kws_get_ranges,
    )
    # gr1.head(1)

    gr2 = get_ranges(
        df2,
        col_start=feat2_start,
        col_end=feat2_end,
        convert=convert,
        **kws_get_ranges,
    )
    # gr2.head(1)
    df3 = gr1.intersect(
        gr2,
    ).as_df()
    if len(df3) == 0:
        return
    else:
        if convert:
            df3 = df3.assign(
                Start=lambda df: df["Start"] + 1,
            )
        return df3.astype({"Start": int, "End": int}).drop(["Chromosome"], axis=1)


def map_back_feat_ids(
    df1,
    df2,
    feat1_id="t.id",
    feat1_start="e.start",
    feat1_end="e.end",
    feat2_id="e.id",
    feat2_start="c.start",
    feat2_end="c.end",
    **kws_get_ranges,
):
    # gr1 = get_ranges(
    #     df1.loc[:, [feat1_id, feat1_start, feat1_end]].drop_duplicates(),
    #     col_start=feat1_start,
    #     col_end=feat1_end,
    #     **kws_get_ranges,
    # )
    # # gr1.head(1)

    # gr2 = get_ranges(
    #     df2.dropna().loc[:, [feat2_id, feat2_start, feat2_end]].drop_duplicates(),
    #     col_start=feat2_start,
    #     col_end=feat2_end,
    #     **kws_get_ranges,
    # )

    # gri=(
    #     gr1
    #     .intersect(
    #         gr2,
    #     )
    #     .as_df()
    # )

    gri = intersect(
        df1,
        df2,
        feat1_id=feat1_id,
        feat1_start=feat1_start,
        feat1_end=feat1_end,
        feat2_id=feat2_id,
        feat2_start=feat2_start,
        feat2_end=feat2_end,
        common_ids=False,
        **kws_get_ranges,
    )
    assert len(gri) > 0, "expected intersection, but not found"

    # print(gri)
    ## map
    df3 = gri.assign(
        **{feat2_id: lambda df: df.apply(lambda x: f"{x['Start']}-{x['End']}", axis=1)}
    ).loc[:, [feat1_id, feat2_id]]
    return df2.log.merge(
        right=df3,
        on=feat2_id,
        how="left",
        # validate="m:1",
    )


def intersect_with_seqs(
    df1,
    df2,
    seq_id="t.id",
    feat1_id="e.id",
    feat1_start="e.start",
    feat1_end="e.end",
    feat1_prefix=None,
    feat2_id="c.id",
    feat2_start="c.start",
    feat2_end="c.end",
    feat2_prefix=None,
    **kws_get_ranges,
):
    if feat1_prefix is None:
        feat1_prefix = feat1_id.split(".")[0]
    if feat2_prefix is None:
        feat2_prefix = feat2_id.split(".")[0]
    mapped_feat_prefix = f"{feat2_prefix}{feat1_prefix}"
    mapped_feat_id = f"{mapped_feat_prefix}.id"
    mapped_feat_start = f"{mapped_feat_prefix}.start"
    mapped_feat_end = f"{mapped_feat_prefix}.end"
    mapped_feat_length = f"{mapped_feat_prefix}.length"

    df1_ = (
        df1.groupby(
            seq_id,
            as_index=False,
        )
        .apply(
            lambda df: intersect(
                df,
                df2.query(expr=f"`{seq_id}` == '{df.name}'"),
                feat1_id=seq_id,
                feat1_start=feat1_start,
                feat1_end=feat1_end,
                feat2_id=feat1_id,
                feat2_start=feat2_start,
                feat2_end=feat2_end,
                **kws_get_ranges,
            )
        )
        .reset_index(drop=True)
    )
    # print(df1_)
    if len(df1_) == 0:
        return
    else:
        return df1_.rename(
            columns={
                "Start": mapped_feat_start,
                "End": mapped_feat_end,
            },
            errors="raise",
        ).assign(
            **{
                mapped_feat_id: lambda df: df.apply(
                    lambda x: f"{x[mapped_feat_start]}-{x[mapped_feat_end]}", axis=1
                ),
                mapped_feat_length: lambda df: df.apply(
                    lambda x: get_diff_sorted(x[mapped_feat_start], x[mapped_feat_end]),
                    axis=1,
                ),
            },
        )


def intersect_with_feats(
    df1,
    df2,
    seq_id="t.id",
    feat1_id="e.id",
    feat1_start="e.start",
    feat1_end="e.end",
    feat1_prefix=None,
    feat2_id="c.id",
    feat2_start="c.start",
    feat2_end="c.end",
    feat2_prefix=None,
):
    df2 = df2.dropna(
        subset=[
            feat2_id,
            feat2_start,
            feat2_end,
        ],
    )
    if feat1_prefix is None:
        feat1_prefix = feat1_id.split(".")[0]
    if feat2_prefix is None:
        feat2_prefix = feat2_id.split(".")[0]

    mapped_feat_prefix = f"{feat2_prefix}{feat1_prefix}"
    mapped_feat_id = f"{mapped_feat_prefix}.id"
    mapped_feat_start = f"{mapped_feat_prefix}.start"
    mapped_feat_end = f"{mapped_feat_prefix}.end"
    mapped_feat_length = f"{mapped_feat_prefix}.length"

    ## map to tids
    df3 = intersect_with_seqs(
        df1,
        df2,
        seq_id=seq_id,
        feat1_id=feat1_id,
        feat1_start=feat1_start,
        feat1_end=feat1_end,
        feat1_prefix=feat1_prefix,
        feat2_id=feat2_id,
        feat2_start=feat2_start,
        feat2_end=feat2_end,
        feat2_prefix=feat2_prefix,
    )
    if df3 is None:
        logging.error("no overlap found")
        return
    # df3.head(1)
    ## map back the feat ids
    df3_ = map_back_feat_ids(
        df2,
        df3,
        feat1_id=feat2_id,
        feat1_start=feat2_start,
        feat1_end=feat2_end,
        feat2_id=mapped_feat_id,
        feat2_start=mapped_feat_start,
        feat2_end=mapped_feat_end,
    )
    assert df3_[feat2_id].isnull().sum() == df3_[mapped_feat_start].isnull().sum(), (
        "mapping back to feats failed"
    )

    df4 = df1.log.merge(
        right=df3_,
        how="outer",
        on=[seq_id, feat1_id],
        # validate="1:m",
    )
    return df4


## cached
def get_ts_data(
    gene_id,
    ensembl_release,
    species,
    protein_coding,
    cds_prefix,
    force=False,
    **kws_get_cache,
):
    ## cache
    cache_dir_path = get_cache_dir_path(
        species=get_species_name(species),
        ensembl_release=ensembl_release,
        source="www.ensembl.org/pyensembl/",
        **kws_get_cache,
    )
    cachep = f"{cache_dir_path}/{gene_id}/ts/protein_coding={protein_coding}.pqt"

    if (not Path(cachep).exists()) or force:
        logging.info("fetching data")
        ts = get_ts(
            gene_id=gene_id,
            ensembl_release=ensembl_release,
            species=species,
            protein_coding=protein_coding,
        )
        df1 = get_es(
            ts=ts,
        )
        ## cds
        df2 = get_cs(ts)
        ## mapping cds
        df3 = intersect_with_feats(
            df1,
            df2,
            seq_id="t.id",
            feat1_id="e.id",
            feat1_start="e.start",
            feat1_end="e.end",
            feat1_prefix=None,
            feat2_id="c.id",
            feat2_start="c.start",
            feat2_end="c.end",
            feat2_prefix=cds_prefix,
        )
        # if df3 is None:
        #     return
        to_table(
            df3,
            cachep,
        )
    else:
        logging.info(f"read ts from cache: {cachep}")
        df3 = read_table(cachep)
    return df3
