# AUTOGENERATED! DO NOT EDIT! File to edit: ../03_isoforms.ipynb.

# %% auto 0
__all__ = [
    "get_blocks",
    "map_feats_to_blocks",
    "plot_seq_feats",
    "get_plot_data",
    "set_plot_data",
    "plot_isoforms",
]

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

import matplotlib.pyplot as plt

from roux.stat.paired import get_diff_sorted
from roux.stat.transform import rescale

try:
    from chrov.annots import get_ranges, intersect_with_feats, get_ts_data
except:
    from dups.annots import get_ranges, intersect_with_feats, get_ts_data


def get_blocks(
    df1: pd.DataFrame,
    col_start="e.start",
    col_end="e.end",
    col_block_start="eb.start",
    col_block_end="eb.end",
    # col_strand='t.strand',
) -> pd.DataFrame:
    if df1[col_start].isnull().all() or df1[col_end].isnull().all():
        logging.warning("missing feats data")
        # df1=df1.assign(
        #     **{
        #         col_block_start : col_start,
        #         col_block_end : col_end,
        #     }
        # )
        # logging.warning("feats set to 1 block")
        return df1.drop(
            [
                col_start,
                col_end,
            ],
            axis=1,
        )

    ## no features
    df2_ = df1.query(expr=f"`{col_start}`.isnull()")

    gr = (
        df1.dropna(
            subset=[col_start],
        )
        .loc[
            :,
            [
                col_start,
                col_end,
                # col_strand,
            ],
        ]
        .log.drop_duplicates()
        .pipe(
            get_ranges,
            col_start=col_start,
            col_end=col_end,
            convert=False,
        )
    )
    # gr.head(1)

    gr_clu = gr.cluster(
        # count=True
    )
    # gr_clu.head(1)

    gr_mer = gr_clu.merge(by="Cluster")
    # gr_mer.head(1)

    df2 = (
        gr.as_df()
        .log.merge(
            right=gr_clu.as_df(),
            how="inner",
            on=[
                "Start",
                "End",
                # 'Strand',
                "Chromosome",
            ],
            validate="1:1",
        )
        .drop(["Chromosome"], axis=1)
        .log.merge(
            right=(
                gr_mer.as_df()
                .loc[:, ["Start", "End", "Cluster"]]
                .rename(
                    columns={
                        "Start": "b.start",
                        "End": "b.end",
                    },
                    errors="raise",
                )
            ),
            how="inner",
            on="Cluster",
            validate="m:1",
        )
        .rename(
            columns={
                "Cluster": "b.id",
                "Start": col_start,
                "End": col_end,
                # 'Strand':col_strand,
            },
            errors="raise",
        )
        .merge(
            right=df1,
            on=[
                # col_strand,
                col_start,
                col_end,
            ],
            how="inner",
            validate="1:m",
        )
        .sort_index(axis=1)
        .assign(
            **{
                "b.length": lambda df: df.apply(
                    lambda x: get_diff_sorted(x["b.start"], x["b.end"]), axis=1
                ),
                "eb.overlap": lambda df: df.apply(
                    lambda x: min(
                        [1, get_diff_sorted(x[col_start], x[col_end]) / x["b.length"]]
                    ),
                    axis=1,
                ),
                col_block_start: lambda df: df.apply(
                    lambda x: (
                        x["b.id"]
                        + (get_diff_sorted(x[col_start], x["b.start"]) / x["b.length"])
                    ),
                    axis=1,
                ),
                f"{col_block_end} raw": lambda df: df.apply(
                    lambda x: (
                        (x["b.id"] + 1)
                        - (get_diff_sorted(x[col_end], x["b.end"]) / x["b.length"])
                    ),
                    axis=1,
                ),
                col_block_end: lambda df: df.apply(
                    lambda x: rescale(
                        [x[col_block_start], x[f"{col_block_end} raw"]],
                        range1=[x["b.id"], x["b.id"] + 1],
                        range2=[x["b.id"], x["b.id"] + 0.9],
                    )[1],
                    axis=1,
                ),
                "eb.id": lambda df: df.apply(
                    lambda x: f"{x[col_block_start]}-{x[col_block_end]}", axis=1
                ),
                "eb.length": lambda df: df.apply(
                    lambda x: get_diff_sorted(x[col_block_start], x[col_block_end]),
                    axis=1,
                ),
            }
        )
    )
    return pd.concat([df2, df2_], axis=0).reset_index(drop=True)


def map_feats_to_blocks(
    data,
    feat_start="c.start",
    feat_end="c.end",
    prefix="c",
    range1_prefix="e",
    range2_prefix="eb",
):
    """
    Rescale the coordinates of the features to the exon.


    """
    # if not f"{range1_prefix}.start" in data or not f"{range1_prefix}.end" in data:

    #     return

    return (
        data
        # .dropna(subset=['c.start','c.end'],axis=0)
        .assign(
            **{
                f"{prefix}b.id": lambda df: df.apply(
                    lambda x: rescale(
                        [x[feat_start], x[feat_end]],
                        range1=[x[f"{range1_prefix}.start"], x[f"{range1_prefix}.end"]],
                        range2=[x[f"{range2_prefix}.start"], x[f"{range2_prefix}.end"]],
                    ),
                    axis=1,
                ),
                f"{prefix}b.start": lambda df: df.apply(
                    lambda x: x[f"{prefix}b.id"][0], axis=1
                ),
                f"{prefix}b.end": lambda df: df.apply(
                    lambda x: x[f"{prefix}b.id"][1], axis=1
                ),
            }
        ).assign(
            **{
                f"{prefix}b.id": lambda df: df[f"{prefix}b.id"].apply(
                    lambda x: f"{x[0]}-{x[1]}"
                )
            }
        )
    )


def plot_seq_feats(
    df1,
    col_id,  ## sequence id
    col_start="t.start",
    col_end="t.end",
    col_feat_id=None,
    col_feat_start="e.start",
    col_feat_end="e.end",
    order=None,  # list|{cols:asceding}
    sort: dict = None,  # list|{cols:asceding}
    order_top: list = None,  # show on top
    kind: str = "joined",
    show_lines: bool = True,
    color_line: str = "gray",
    test: bool = False,
    ax=None,
    **kws_ranges,
):
    """
    Plot sequence features.
    """
    ax = ax if ax is not None else plt.gca()

    if len(df1) == 0:
        logging.error("empty data found. exiting")
        return ax

    if isinstance(sort, dict):
        order_sorted = (
            df1.sort_values(
                list(sort.keys()),
                ascending=list(sort.values()),
            )[col_id]
            .drop_duplicates()
            .tolist()[::-1]
        )  # since y axis is not inverted
    else:
        order_sorted = sorted(df1[col_id].unique())

    if order is None:
        order = order_sorted
    else:
        if isinstance(sort, dict):
            ## filter by order
            order = [s for s in order_sorted if s in order]
        else:
            ## order by the provided order
            order = [s for s in order if s in order_sorted]

    if order_top is not None:
        order = [s for s in order if s not in order_top] + [
            s for s in order if s in order_top
        ]
    logging.info(f"order={order}")

    if "y" in df1:
        logging.warning(
            "y column found, would be overwritten for ordering the sequences."
        )

    df1 = df1.assign(y=lambda df: df[col_id].map({s: i for i, s in enumerate(order)}))
    # print(data.columns)
    # print(show_lines)
    if show_lines:
        # todo: use plot ranges with low width of the line, to label the sequences without any features
        col_start_lines = "eb.start" if "eb.start" in df1 else col_start
        col_end_lines = "eb.end" if "eb.end" in df1 else col_end
        data_ = (
            df1.groupby(col_id)
            .agg(
                {
                    col_start_lines: "min",
                    col_end_lines: "max",
                    "y": "min",
                }
            )
            # .loc[:, [col_id, col_start_lines, col_end_lines, "y"]]
            .reset_index(0)
            .drop_duplicates()
        )
        # print(data_.iloc[0,:])
        df_ = data_.apply(
            lambda x: ax.hlines(
                y=x["y"],
                xmin=x[col_start_lines],  # =data[col_start].min()
                xmax=x[col_end_lines],
                color=color_line,
                zorder=0,
                lw=0.5,
            ),
            axis=1,
        )

    feats = True
    if col_feat_start not in df1 or col_feat_end not in df1:
        feats = False
    else:
        if df1[col_feat_start].isnull().all() or df1[col_feat_end].isnull().all():
            feats = False

    if not feats:
        logging.warning("missing feats data")
        ax.set(
            yticks=range(len(order)),
            yticklabels=order,
        )
        # ax.axes.xaxis.set_visible(test)
        # ax.axes.yaxis.set_visible(test)
        _ = plt.setp(
            ax.spines.values(),
            zorder=0,
            color="none",
        )
        if "strand" in kws_ranges:
            if kws_ranges["strand"] == "-":
                ax.invert_xaxis()
        return ax

    kws_ranges = {
        **dict(
            col_id=col_id,
            col_label=col_feat_id,
            col_start=col_feat_start,
            col_end=col_feat_end,
            kind=kind,
            y="y",
        ),
        **kws_ranges,
    }

    ## pre-processing
    cols = list(
        set(
            [
                v
                for k, v in kws_ranges.items()
                if (k.startswith("col_") or k == "hue") and v is not None
            ]
        )
    ) + ["y"]

    data = df1.log(col_id).loc[:, cols].log.drop_duplicates().log.dropna().log(col_id)

    if len(data) == 0:
        print(df1.rd.check_na())

    from chrov.viz.ranges import plot_ranges

    # %run ../chrov/chrov/viz/ranges.py
    plot_ranges(data=data, ax=ax, **kws_ranges)

    # if df1[col_id].nunique() != data[col_id].nunique():
    #     return ax

    ax.set_ylim(
        # -1,
        # data[col_id].nunique(),
        -0.25,
        len(order) - 0.75,
    )

    return ax


def get_plot_data(
    data,
    feats,
    force,
):
    if force:
        data, feats = None, None
    else:
        if data is not None:
            if "data type" in data:
                # print(data['data type'].unique())
                feats_ = data.query("`data type`=='feats'").dropna(how="all", axis=1)
                data = data.query("`data type`=='data'").dropna(how="all", axis=1)
                if feats is None:
                    if len(feats_) == 0:
                        feats = None
                    else:
                        feats = feats_.copy()
                else:
                    logging.warning(
                        "feats: prioritizing the input over the contents of the plot data."
                    )
            else:
                data = None
    return data, feats


def set_plot_data(
    data,
    feats,
):
    if feats is not None:
        if "data type" in feats:
            feats = feats.drop(["data type"], axis=1)
    if "data type" in data:
        data = data.drop(["data type"], axis=1)

    return pd.concat(
        dict(
            data=data,
            feats=feats,
        ),
        axis=0,
        names=["data type"],
    ).reset_index(0)


## Wrapper
def plot_isoforms(
    gene_id,
    ensembl_release,
    species,
    layout=None,  #'blocks',
    protein_coding=True,
    ## constants
    cds_prefix="c",
    ## colors
    # color_exons="lightblue",  #'lightskyblue',
    # color_exons="gainsboro", #dcdcdc
    # color_exons="whitesmoke", #too light
    color_exons="#e9e9e9",
    # color_cds='skyblue',
    # color_cds="deepskyblue",
    color_cds="lightblue",
    ## features
    feat_id="d.id",
    feat_start="d.start",
    feat_end="d.end",
    # color_feats='deepskyblue',
    color_feats="dodgerblue",
    # feat_prefix=None,
    ## i/o
    feats=None,
    data: pd.DataFrame = None,
    return_data=False,
    force: bool = False,  # over-write cache
    verbose: bool = False,
    ax=None,
    kws_get_cache={},
    **kws_plot_seq_feats,
):
    cds_in_exon_prefix = f"{cds_prefix}e"

    if isinstance(feats, pd.DataFrame):
        if len(feats) == 0:
            feats = None

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

    if feats is not None:
        # if feat_prefix is None:
        feat_prefix = feat_id.split(".")[0]
        assert feat_id.startswith(f"{feat_prefix}.")
        assert feat_start.startswith(f"{feat_prefix}.")
        assert feat_end.startswith(f"{feat_prefix}.")
        feats_in_exon_prefix = f"{feat_prefix}e"
    # if not feats is None:
    # else:
    if layout == "blocks":
        suffix = "b"
    else:
        suffix = ""

    ## getting data
    if data is None:
        df3 = get_ts_data(
            gene_id=gene_id,
            ensembl_release=ensembl_release,
            species=species,
            protein_coding=protein_coding,
            cds_prefix="c",
            **kws_get_cache,
        )

        if feats is not None:
            # print(df3.columns)
            # df3_=df3.query(
            #     expr=f"`{feat_id}`.isnull()"
            # )
            # df3=df3.query(
            #     expr=f"~(`{feat_id}`.isnull())"
            # )

            df3 = intersect_with_feats(
                df3,
                feats,
                seq_id="t.id",
                feat1_id="e.id",
                feat1_start="e.start",
                feat1_end="e.end",
                feat2_id=feat_id,
                feat2_start=feat_start,
                feat2_end=feat_end,
            )

            if df3 is None:
                logging.error("intersect_with_feats: no overlap found")
                return
            else:
                df3.log("t.id")

            # df3=pd.concat(
            #     [
            #         df3,
            #         df3_
            #     ],
            #     axis=0,
            # ).reset_index(drop=True)

        if layout == "blocks":
            ## exon blocks
            df4 = get_blocks(
                df3,
                col_start="e.start",
                col_end="e.end",
                col_block_start="eb.start",
                col_block_end="eb.end",
            )
            if df4 is None:
                return
                logging.error("map_feats_to_blocks: no overlap")

            ## mapping cds
            df5 = map_feats_to_blocks(
                df4,
                feat_start=f"{cds_in_exon_prefix}.start",
                feat_end=f"{cds_in_exon_prefix}.end",
                prefix=cds_in_exon_prefix,
            )
            if df5 is None:
                return
                logging.error("map_feats_to_blocks: no overlap")
            if feats is not None:
                ## mapping feats
                df5 = map_feats_to_blocks(
                    df5,
                    # feat_start=feat_start,
                    # feat_end=feat_end,
                    # prefix=feat_prefix,
                    feat_start=f"{feats_in_exon_prefix}.start",
                    feat_end=f"{feats_in_exon_prefix}.end",
                    prefix=feats_in_exon_prefix,
                )
            data = df5.copy()
        else:
            data = df3.copy()
        if "t.strand" not in data:
            data["t.strand"] = data["t.strand"].unique()[0]

    # print(data.columns)
    # return data
    strand = data["t.strand"].unique()[0]
    # print(strand)
    ## plotting parameters
    kws_plot = dict(
        plot=dict(
            col_id="t.id",
            col_start="t.start",
            col_end="t.end",
            kind="joined",
            strand=strand,
            test=False,
        ),
        exons=dict(
            color=color_exons,
            col_feat_id=f"e{suffix}.id",
            col_feat_start=f"e{suffix}.start",
            col_feat_end=f"e{suffix}.end",
            show_lines=layout != "blocks",
        ),
        cds=dict(
            color=color_cds,
            col_feat_id=f"{cds_in_exon_prefix}{suffix}.id",
            col_feat_start=f"{cds_in_exon_prefix}{suffix}.start",
            col_feat_end=f"{cds_in_exon_prefix}{suffix}.end",
            show_labels=False,
            show_lines=False,
        ),
    )
    # print(kws_plot)

    if feats is not None:
        ## appending feats
        kws_plot = {
            **kws_plot,
            **dict(
                feats=dict(
                    color=color_feats if color_feats not in data else "none",
                    hue=None if color_feats not in data else color_feats,
                    col_feat_id=f"{feats_in_exon_prefix}{suffix}.id",
                    col_feat_start=f"{feats_in_exon_prefix}{suffix}.start",
                    col_feat_end=f"{feats_in_exon_prefix}{suffix}.end",
                    show_labels=False,
                    show_lines=False,
                ),
            ),
        }

    ## plot
    for k in kws_plot:
        if k != "plot":
            logging.info(f"plotting {k} ..")
            kws_plot_ = {
                **kws_plot["plot"],
                **kws_plot[k],
                **kws_plot_seq_feats,
            }
            if verbose:
                print(
                    kws_plot["plot"],
                    kws_plot[k],
                    kws_plot_seq_feats,
                    kws_plot_,
                )
            ax = plot_seq_feats(
                data,
                ax=ax,
                **kws_plot_,
            )
            if k == "exons":
                xlims = ax.get_xlim()

    _ = ax.set(
        xlim=xlims,
        # ylim=[-0.5, data["t.id"].nunique()],
    )
    if layout == "blocks":
        _ = ax.get_xaxis().set_visible(False)
    if return_data:
        return set_plot_data(
            data,
            feats,
        )
    else:
        return ax
