# AUTOGENERATED! DO NOT EDIT! File to edit: ../04_domains.ipynb.

# %% auto 0
__all__ = [
    "query_domains",
    "format_domains",
    "get_domains",
    "get_ploc",
    "get_p_bounds",
    "get_protein_cds_coords",
    "to_gpos",
    "get_ds_data",
    "plot_domains",
    "annot_feats",
]

# %% ../04_domains.ipynb 5
## helper functions
import logging

logging.basicConfig(level=logging.INFO)

import pandas as pd

import roux.lib.df as rd  # noqa
from roux.stat.paired import get_diff_sorted

import matplotlib.pyplot as plt

try:
    from dups.core import query_genes, get_cache_dir_path
    from dups.annots import is_protein_coding, get_annots
    from dups.isoforms import (
        get_ranges,
        get_blocks,
        get_plot_data,
        set_plot_data,
        get_ts,
        plot_seq_feats,
        plot_isoforms,
        intersect_with_feats,
        map_feats_to_blocks,
    )
except:
    from chrov.core import query_genes, get_cache_dir_path
    from chrov.annots import is_protein_coding, get_annots
    from chrov.isoforms import (
        get_ranges,
        get_blocks,
        get_plot_data,
        set_plot_data,
        get_ts,
        plot_seq_feats,
        plot_isoforms,
        intersect_with_feats,
        map_feats_to_blocks,
    )

from functools import partial

query_domains = partial(
    query_genes,
    cols=[
        "ensembl_gene_id",
        "ensembl_peptide_id",
        "interpro",  # Interpro ID
        "interpro_short_description",  # Interpro Short Description
        "interpro_description",  # Interpro Description
        "interpro_start",  # Interpro start
        "interpro_end",  # Interpro end
    ],
)


def format_domains(
    df0_doms,
    min_length=10,
):
    ## cleaning
    df1_doms = df0_doms.log.dropna(subset=["Interpro ID"]).log.query(
        expr="~(`Interpro Description`.str.contains('superfamily'))"
    )
    # df1_doms.head(1)

    ## formatting

    ## merge overlapping domains
    df2_doms = (
        get_ranges(
            df1_doms,
            col_start="Interpro start",
            col_end="Interpro end",
        )
        .boundaries(["protein id", "Interpro Short Description"])
        .as_df()
        .log()
        .drop(["Chromosome"], axis=1)
        .rename(
            columns={
                "protein id": "p.id",
                "Interpro Short Description": "d.id",
                "Start": "d.start",
                "End": "d.end",
            },
            errors="raise",
        )
        .assign(
            **{
                "d.id": lambda df: df["d.id"].str.replace("_", " "),
                "d.length": lambda df: df["d.end"] - df["d.start"],
            }
        )
        .log.query(expr=f"`d.length`>={min_length}")
        .loc[:, ["p.id", "d.id", "d.start", "d.end", "d.length"]]
        .rd.assert_dense()
    )
    # df2_doms.head(1)
    return df2_doms


def get_domains(
    species,
    ensembl_release,
    cache_dir_path=None,
    cols=["p.id", "d.id", "d.start", "d.end", "d.length"],
    force=False,
    **kws_get_cache,
):
    from pathlib import Path
    from roux.lib.io import read_table, to_table

    cache_dir_path = get_cache_dir_path(
        species=species,
        ensembl_release=ensembl_release,
        source="www.ensembl.org/biomart",
        **kws_get_cache,
    )

    from roux.lib.str import encode

    outp = f"{cache_dir_path}{encode(dict(cols=cols), short=True)}.pqt"

    if not Path(outp).exists() or force:
        logging.warning("fetching domains..")
        df0_doms = query_domains(
            species=species,
            release=ensembl_release,
            verbose=False,
        )
        df1 = format_domains(
            df0_doms,
            min_length=10,
        )
        to_table(df1, outp)
        logging.info(f"cache path: {outp}")
    else:
        logging.warning(f"read from cache: {outp}")
        df1 = read_table(outp)
    return df1


def get_ploc(pids, **kws_annots):
    """
    Get protein location.
    """

    annots = get_annots(
        **kws_annots,
    )
    ps = []
    for k in pids:
        t = annots.transcript_by_protein_id(k)
        ps.append(
            {
                "p.id": k,
                "p.start": 1,
                "p.end": len(
                    t.protein_sequence,
                ),
            }
        )
    return pd.DataFrame(ps)


def get_p_bounds(
    pos,
    pstart: bool,
    strand,
):
    """
    Get the genome coord. for a protein position depending on the strand.
    """
    if strand == "+":
        if pstart:
            return min(pos)
        else:
            return max(pos)
    elif strand == "-":
        if pstart:
            return max(pos)
        else:
            return min(pos)


# copy of from beditor.lib.get_mutations import get_protein_cds_coords
def get_protein_cds_coords(annots, protein_id: str, strict=True) -> pd.DataFrame:
    """Get protein CDS coordinates

    Args:
        annots: pyensembl annotations
        protein_id (str): protein ID

    Returns:
        pd.DataFrame: output table
    """
    t = annots.transcript_by_protein_id(protein_id)
    # g=annots.gene_by_protein_id(protein_id)

    if not is_protein_coding(
        t,
        strict=strict,
    ):  # (t.is_protein_coding and t.contains_start_codon and t.contains_stop_codon):
        logging.error(
            f"Excluded protein id {protein_id} because either it is not protein-coding or does not contain start and stop codons."
        )
        return  # None, None, None

    assert (
        len(t.protein_sequence) * 3
        == len(t.coding_sequence) - 3
        == sum([i[1] - i[0] + 1 for i in t.coding_sequence_position_ranges])
    ), "CDS length is not compatible with the protein sequence"

    from roux.lib.set import flatten

    return (
        pd.DataFrame(
            {
                "pos": flatten(
                    [
                        list(range(i[0], i[1] + 1))
                        for i in t.coding_sequence_position_ranges
                    ]
                ),
            }
        )
        .sort_values("pos", ascending=t.strand == "+")
        .assign(
            **{
                "chrom": t.contig,
                "strand": t.strand,
                "base": list(t.coding_sequence[:-3]),
                "aa": flatten([[s] * 3 for s in list(t.protein_sequence)]),
                "aa pos": flatten(
                    [[i + 1] * 3 for i, s in enumerate(t.protein_sequence)]
                ),
            }
        )
        .astype(
            {
                "chrom": str,
                "strand": str,
                "pos": int,
                "base": str,
                "aa": str,
                "aa pos": int,
            }
        )
    )


def to_gpos(
    df3,
    ensembl_release,
    species,
    cols_feats=["d.id", "d.start", "d.end"],
):
    """
    To the genome positions of the domains.
    """

    for c in cols_feats:
        if c not in df3:
            return None  # df3
        elif df3[c].isnull().all():
            return None  # df3

    df3 = df3.dropna(
        subset=cols_feats,
    )

    if len(df3) == 0:
        logging.warning("missing feats")
        return df3

    annots = get_annots(
        ensembl_release=ensembl_release,
        species=species,
    )

    ## map the aa pos to genome coord.s
    # from beditor.lib.get_mutations import get_protein_cds_coords
    df4_ = (
        df3.groupby("p.id")
        .apply(
            lambda df: get_protein_cds_coords(
                annots=annots,
                protein_id=df.name,
            )
        )
        .reset_index(0)
    )
    # print(df4_)

    df4 = (
        df4_.groupby(
            [
                "p.id",
                "aa pos",
            ]
        )
        .agg({"pos": list})
        .reset_index()
    )
    # print(df4)

    ## merging
    df5 = (
        df3.melt(
            id_vars=["t.id", "p.id", "d.id"],
            value_vars=["d.start", "d.end"],
            value_name="aa pos",
        )
        .merge(
            right=df4,
            on=["p.id", "aa pos"],
            how="inner",
            validate="m:1",
        )
        .assign(
            posi=lambda df: df.apply(
                lambda x: get_p_bounds(
                    pos=x["pos"],
                    pstart=x["variable"].endswith("start"),
                    strand=annots.transcript_by_protein_id(x["p.id"]).strand,
                ),
                axis=1,
            )
        )
    )

    # print(df5)
    df6 = (
        df5.pivot(
            index=["t.id", "p.id", "d.id"],
            columns="variable",
            values=["aa pos", "posi"],
        )
        .swaplevel(axis=1)
        .rd.flatten_columns()
        .rd.renameby_replace({" posi": ""})
        .reset_index()
        ## sort start and end
        .rd.sort_columns_by_values(
            subset=["d.start", "d.end"],
            suffixes=["start", "end"],
            clean=True,
        )
        .assign(
            **{
                "d.length": lambda x: get_diff_sorted(
                    x["d.end"],
                    x["d.start"],
                ),
            },
        )
    )
    # print(df6)

    assert not any(df6["d.length"] < 0)
    return pd.concat(
        [
            df6.sort_values("d.start"),
            df3,
        ],
        axis=0,
    ).reset_index(drop=True)


## cache
from pathlib import Path


def get_ds_data(
    gene_id,
    ensembl_release,
    species,
    layout,
    suffix,
    # flt=None,#'longest',
    force=False,
    **kws_get_cache,
):
    ## cache
    cache_dir_path = get_cache_dir_path(
        species=species,
        ensembl_release=ensembl_release,
        source="www.ensembl.org/pyensembl/",
        **kws_get_cache,
    )
    cachep = f"{cache_dir_path}/{gene_id}/ps/layout={layout}.pqt"

    if (not Path(cachep).exists()) or force:
        logging.info("fetching data for protein-coding transcripts")
        # if not protein_coding:
        #     # protein_coding=True
        #     logging.warning(f"protein_coding set to True because biotype={biotype}")
        logging.info("fetching domains data")

        ## domains
        df0_dom = get_domains(
            species=species,
            ensembl_release=ensembl_release,
        )

        ## Querying protein ids by gene id
        ts = get_ts(
            gene_id=gene_id,
            ensembl_release=ensembl_release,
            species=species,
            protein_coding=True,
        )

        # if 'order' in kws_plot_seq_feats:
        # if flt=='longest':
        #     # =df3.sort_values('t.length',ascending=False).head(1)['t.id'].tolist()
        #     ts=sorted(ts, key=lambda p: p.length)[-1:]

        df_ = pd.DataFrame(
            [(t.id, t.protein_id) for t in ts],
            columns=[
                "t.id",
                "p.id",
            ],
        )

        ## map to domains
        df1 = df_.merge(
            right=df0_dom,
            how="left",
            on="p.id",
            validate="1:m",
        )

        ### Mapping start and end of protein
        df2 = get_ploc(
            pids=df1["p.id"].unique(),
            species=species,
            ensembl_release=ensembl_release,
        )

        df3 = df2.merge(
            right=df1,
            on="p.id",
            how="inner",
            validate="1:m",
        ).sort_values(
            ["p.end", "d.length"],
            ascending=[True, False],
        )
        # df3.head(1)
        if layout == "blocks":
            df3 = get_blocks(
                df3.log.dropna(),
                col_start="d.start",
                col_end="d.end",
                col_block_start=f"d{suffix}.start",
                col_block_end=f"d{suffix}.end",
            )
    else:
        from roux.lib.io import read_table

        logging.info(f"read ts from cache: {cachep}")
        df3 = read_table(cachep)
    return df3


def plot_domains(
    gene_id,
    species,
    ensembl_release,
    biotype="p",
    layout=None,
    # flt=None,
    ## i/o
    feats=None,
    data: pd.DataFrame = None,
    return_data=False,
    force=False,
    verbose=False,
    test=False,
    ax=None,
    kws_get_cache={},
    **kws_ranges,
):
    protein_coding = True
    if "protein_coding" in kws_ranges:
        del kws_ranges["protein_coding"]

    if not verbose:
        logging.disable(logging.CRITICAL)

    data, feats = get_plot_data(
        data=data,
        feats=feats,
        force=force,
    )

    if layout == "blocks":
        suffix = "b"
    else:
        suffix = ""

    if biotype == "t" and not protein_coding:
        ## use the input data if available
        ## else fetch data inside plot_soforms and map the domain feats from below
        data_t = data  # .copy()

    if data is None:
        data = get_ds_data(
            gene_id=gene_id,
            ensembl_release=ensembl_release,
            species=species,
            layout=layout,
            # flt=flt,
            suffix=suffix,
            force=force,
            **kws_get_cache,
        )

    ## updating the kws param.s
    kws_ranges = {
        **kws_ranges,
        **dict(
            kws_legend=dict(
                title="Domains",
                ncol=2,
                handletextpad=0.5,
            )
        ),
    }

    if biotype.startswith("p"):
        kws_plot = dict(
            plot=dict(
                col_id="t.id",
                col_start="p.start",
                col_end="p.end",
                col_feat_id="d.id",
                hue="d.id",
                start=1,
                kind="joined",
                show_lines=layout
                != "blocks",  ## todo: if one transcript, line is not shown
            )
        )
        # print(kws_plot)
        ### with protein residue numbers
        ax = plot_seq_feats(
            data,
            col_feat_start=f"d{suffix}.start",
            col_feat_end=f"d{suffix}.end",
            **kws_plot["plot"],
            **kws_ranges,
            ax=ax,
        )

    elif biotype.startswith("t"):
        if feats is None:
            feats = to_gpos(
                data,
                ensembl_release=ensembl_release,
                species=species,
            )
            if feats is not None:
                feats = feats.loc[:, ["t.id", "d.id", "d.start", "d.end"]]
            # else:
            #     feats=None
        # print(f"feats is None {feats is None}")

        data = plot_isoforms(
            gene_id=gene_id,
            ensembl_release=ensembl_release,
            species=species,
            layout=layout,
            protein_coding=protein_coding,
            ## columns
            feat_id="d.id",
            feat_start="d.start",
            feat_end="d.end",
            color_feats="d.id",
            # i/o
            feats=feats,
            data=data if protein_coding else data_t,
            return_data=return_data,
            force=force,
            ax=ax,
            **kws_ranges,
        )
        ax = plt.gca()
    if biotype == "p":
        ax.set(
            xlabel=None,
        )
    if layout == "blocks":
        if not test:
            _ = ax.get_xaxis().set_visible(False)

    if not verbose:
        logging.disable(logging.NOTSET)

    if return_data:
        return set_plot_data(
            data,
            feats,
        )
    else:
        return ax


def annot_feats(
    feats,
    data,
    layout,
    feat_id,
    feat_start,
    feat_end,
    ## constants
    feat_prefix="f",
    ## plotting
    figsize=[5, 1],
    verbose=False,
    ax=None,
    kws_text={},
    **kws_annot_side_curved,
):
    if not verbose:
        logging.disable(logging.CRITICAL)

    ## remove features from the data
    cols = ["t.id", "t.start", "t.end", "e.id", "e.start", "e.end"] + (
        ["eb.id", "eb.start", "eb.end"] if layout == "blocks" else []
    )
    # print(cols)

    data_back = (
        data.query(expr="`data type`=='data'")
        .drop(["data type"], axis=1)
        .log.dropna(subset=["t.id", "t.start", "t.end"])
        .loc[:, cols]
        .astype(
            {
                "t.start": int,
                "t.end": int,
            }
        )
        .log.drop_duplicates()
    )
    ### feats
    ## get the y
    os = [o for o in ax.get_children() if isinstance(o, plt.Text)]
    seqid_to_y = {
        o.get_text().strip(): o.get_position()[1]
        for o in os
        if o.get_text().strip() != ""
    }

    ## get the x
    df1 = feats.astype(
        {
            feat_start: int,
            feat_end: int,
        },
    ).assign(
        **{
            f"{feat_prefix}.id": lambda df: df[feat_id].copy(),
            f"{feat_prefix}.start": lambda df: df[feat_start].copy(),
            f"{feat_prefix}.end": lambda df: df[feat_end].copy(),
            f"{feat_prefix}.y": lambda df: df["t.id"].map(seqid_to_y),
        }
    )
    # if any(df1[feat_start]==df1[feat_end]):
    #     logging.warning("expected 1-based coordinates, but start=end found. Therefore offsetting start to proceed further.")
    #     df1=(
    #         df1
    #         .assign(
    #             **{
    #                 feat_end: lambda df: df.apply(
    #                     lambda x: x[feat_start]-1 if x[feat_start]==x[feat_end] else x[feat_end],
    #                     axis=1,
    #                 )
    #             },
    #         )
    #     )
    # print(
    #     data_back.to_dict(orient='records'),
    #     df1.to_dict(orient='records'),
    #      )
    ### intersect with the sequences
    df3 = intersect_with_feats(
        (
            data_back.assign(**{"e.start": lambda df: df["e.start"] - 1})
        ),  ## fix untill switched to bedtools
        df1,
        seq_id="t.id",
        feat1_id="e.id",
        feat1_start="e.start",
        feat1_end="e.end",
        feat1_prefix=None,
        feat2_id=f"{feat_prefix}.id",
        feat2_start=f"{feat_prefix}.start",
        feat2_end=f"{feat_prefix}.end",
        feat2_prefix=feat_prefix,
    )
    if df3 is None:
        logging.error("intersect_with_feats: no overlap found")
        return
    if layout == "blocks":
        ## rescale
        df3 = map_feats_to_blocks(
            df3,
            # feat_start=feat_start,
            # feat_end=feat_end,
            # prefix=feat_prefix,
            feat_start=f"{feat_prefix}e.start",
            feat_end=f"{feat_prefix}e.end",
            prefix=feat_prefix + "e",
        )
        suffix = "eb"
    else:
        suffix = ""

    df4 = df3.log.merge(
        right=df1,
        how="inner",
        on=["t.id", "f.id"],
        # validate=""
    ).assign(
        **{
            f"{feat_prefix}{suffix}.center": lambda df: (
                df.loc[
                    :, [f"{feat_prefix}{suffix}.start", f"{feat_prefix}{suffix}.end"]
                ].mean(axis=1)
            ),
        }
    )
    from roux.viz.annot import annot_side_curved

    # %run ../../code/roux/roux/viz/annot.py
    _ = annot_side_curved(
        ax=ax,
        data=df4,
        colx=f"{feat_prefix}{suffix}.center",
        coly=f"{feat_prefix}.y",
        ## text
        col_label=feat_id,
        loc="top",
        **{
            **dict(
                off=1,
                # lim=ax.get_xlim(),
                kws_text={
                    **dict(
                        rotation=45 if df4[feat_id].nunique() > 1 else 0,
                        ha="left",
                        va="baseline",
                    ),
                    **kws_text,
                },
                ##line
                zorder=3,
            ),
            **kws_annot_side_curved,
        },
    )
    if not verbose:
        logging.disable(logging.NOTSET)

    return ax
