from __future__ import annotations

import numpy as np
import logging

import snapatac2
from snapatac2._snapatac2 import AnnData, AnnDataSet
from snapatac2.tools._misc import aggregate_X
from snapatac2._utils import find_elbow, is_anndata
from ._base import render_plot, heatmap, kde2d, scatter, scatter3d
from ._network import network_scores, network_edge_stat
import snapatac2._snapatac2 as internal

__all__ = [
    'tsse', 'frag_size_distr', 'umap', 'network_scores', 'spectral_eigenvalues',
    'regions', 'motif_enrichment', 'coverage'
]

def valid_cells(
    values,
    width: int = 500,
    height: int = 400,
    **kwargs,
):
    import plotly.graph_objects as go

    values = sorted(values, reverse=True)
    result = {}
    for x, y in enumerate(values):
        x = x + 1
        if y in result:
            x_, n = result[y]
            result[y] = (x_ + x, n + 1)
        else:
            result[y] = (x, 1)
    for y, (x, n) in result.items():
        result[y] = x / n
    y, x = zip(*result.items())

    fig = go.Figure()
    fig.add_trace(go.Scatter(x=x, y=y))
    fig.update_xaxes(type="log")
    fig.update_yaxes(type="log")
    fig.update_layout(
        xaxis_title="Barcodes",
        yaxis_title="Counts",
    )

    return render_plot(fig, width, height, **kwargs)

def tsse(
    adata: AnnData,
    min_fragment: int = 500,
    width: int = 500,
    height: int = 400,
    **kwargs,
) -> 'plotly.graph_objects.Figure' | None:
    """Plot the TSS enrichment vs. number of fragments density figure.

    Parameters
    ----------
    adata
        Annotated data matrix.
    min_fragment
        The cells' unique fragments lower than it should be removed
    width
        The width of the plot
    height
        The height of the plot
    kwargs        
        Additional arguments passed to :func:`~snapatac2.pl.render_plot` to
        control the final plot output. Please see :func:`~snapatac2.pl.render_plot`
        for details.

    Returns
    -------
    'plotly.graph_objects.Figure' | None
        If `show=False` and `out_file=None`, an `plotly.graph_objects.Figure` will be 
        returned, which can then be further customized using the plotly API.

    See Also
    --------
    render_plot

    Examples
    --------
    .. plotly::

        >>> import snapatac2 as snap
        >>> data = snap.read(snap.datasets.pbmc5k(type='h5ad'))
        >>> fig = snap.pl.tsse(data, show=False, out_file=None)
        >>> fig.show()
    """
    if "tsse" not in adata.obs:
        raise ValueError("TSS enrichment score is not computed, please run `metrics.tsse` first.")

    selected_cells = np.where(adata.obs["n_fragment"] >= min_fragment)[0]
    x = adata.obs["n_fragment"].to_numpy()[selected_cells]
    y = adata.obs["tsse"].to_numpy()[selected_cells]

    fig = kde2d(x, y, log_x=True, log_y=False)
    fig.update_layout(
        xaxis_title="Number of unique fragments",
        yaxis_title="TSS enrichment score",
    )

    return render_plot(fig, width, height, **kwargs)

def frag_size_distr(
    adata: AnnData | np.ndarray,
    use_rep: str = "frag_size_distr",
    max_recorded_size: int = 1000,
    **kwargs,
) -> 'plotly.graph_objects.Figure' | None:
    """ Plot the fragment size distribution.
    """
    import plotly.graph_objects as go

    if is_anndata(adata):
        if use_rep not in adata.uns or len(adata.uns[use_rep]) <= max_recorded_size:
            logging.info("Computing fragment size distribution...")
            snapatac2.metrics.frag_size_distr(adata, add_key=use_rep, max_recorded_size=max_recorded_size)
        data = adata.uns[use_rep]
    else:
        data = adata
    data = data[:max_recorded_size+1]

    x, y = zip(*enumerate(data))
    # Make a line plot
    fig = go.Figure()
    fig.add_trace(go.Scatter(x=x[1:], y=y[1:], mode='lines'))
    fig.update_layout(
        xaxis_title="Fragment size",
        yaxis_title="Count",
    )
    return render_plot(fig, **kwargs)

def spectral_eigenvalues(
    adata: AnnData,
    width: int = 600,
    height: int = 400,
    show: bool = True,
    interactive: bool = True,
    out_file: str | None = None,
) -> 'plotly.graph_objects.Figure' | None:
    """Plot the eigenvalues of spectral embedding.

    Parameters
    ----------
    adata
        Annotated data matrix.
    width
        The width of the plot
    height
        The height of the plot
    show
        Show the figure.
    interactive
        Whether to make interactive plot
    out_file
        Path of the output file for saving the output image, end with
        '.svg' or '.pdf' or '.png' or '.html'.

    Returns
    -------
    'plotly.graph_objects.Figure' | None
        If `show=False` and `out_file=None`, an `plotly.graph_objects.Figure` will be 
        returned, which can then be further customized using the plotly API.
    """
 
    import plotly.express as px
    import pandas as pd

    data = adata.uns["spectral_eigenvalue"]

    df = pd.DataFrame({"Component": map(str, range(1, data.shape[0] + 1)), "Eigenvalue": data})
    fig = px.scatter(df, x="Component", y="Eigenvalue", template="plotly_white")
    n = find_elbow(data)
    adata.uns["num_eigen"] = n
    fig.add_vline(x=n)

    return render_plot(fig, width, height, interactive, show, out_file)

def regions(
    adata: AnnData | AnnDataSet,
    groupby: str | list[str],
    peaks: dict[str, list[str]],
    width: float = 600,
    height: float = 400,
    show: bool = True,
    interactive: bool = True,
    out_file: str | None = None,
) -> 'plotly.graph_objects.Figure' | None:
    """
    Parameters
    ----------
    adata
        Annotated data matrix.
    groupby
        Group the cells into different groups. If a `str`, groups are obtained from
        `.obs[groupby]`.
    peaks
        Peaks of each group.
    width
        The width of the plot
    height
        The height of the plot
    show
        Show the figure
    interactive
        Whether to make interactive plot
    out_file
        Path of the output file for saving the output image, end with
        '.svg' or '.pdf' or '.png' or '.html'.

    Returns
    -------
    'plotly.graph_objects.Figure' | None
        If `show=False` and `out_file=None`, an `plotly.graph_objects.Figure` will be 
        returned, which can then be further customized using the plotly API.
    """
    import polars as pl
    import plotly.graph_objects as go

    peaks = np.concatenate([[x for x in p] for p in peaks.values()])
    n = len(peaks)
    if n > 50000:
        logging.warning(f"Input contains {n} peaks, only 50000 peaks will be plotted.")
        np.random.seed(0)
        indices = np.random.choice(n, 50000, replace=False)
        peaks = peaks[sorted(indices)]

    count = aggregate_X(adata, groupby=groupby, normalize="RPKM")
    names = count.obs_names
    count = pl.DataFrame(count.X.T)
    count.columns = list(names)
    idx_map = {x: i for i, x in enumerate(adata.var_names)}
    idx = [idx_map[x] for x in peaks]
    mat = np.log2(1 + count.to_numpy()[idx, :])

    trace = go.Heatmap(
        x=count.columns,
        y=peaks,
        z=mat,
        type='heatmap',
        colorscale='Viridis',
        colorbar={ "title": "log2(1 + RPKM)" },
    )
    data = [trace]
    layout = {
        "yaxis": { "visible": False, "autorange": "reversed" },
        "xaxis": { "title": groupby },
    }
    fig = go.Figure(data=data, layout=layout)
    return render_plot(fig, width, height, interactive, show, out_file)

def umap(
    adata: AnnData | np.ndarray,
    color: str | np.ndarray | None = None,
    use_rep: str = "X_umap",
    marker_size: float = None,
    marker_opacity: float = 1,
    sample_size: int | None = None,
    **kwargs,
) -> 'plotly.graph_objects.Figure' | None:
    """Plot the UMAP embedding.

    Parameters
    ----------
    adata
        Annotated data matrix.
    color
        If the input is a string, it will be used the key to retrieve values from
        `obs`.
    use_rep
        Use the indicated representation in `.obsm`.
    marker_size
        Size of the dots.
    marker_opacity
        Opacity of the dots.
    sample_size
        If the number of cells is larger than `sample_size`, a random sample of
        `sample_size` cells will be used for plotting.
    kwargs        
        Additional arguments passed to :func:`~snapatac2.pl.render_plot` to
        control the final plot output. Please see :func:`~snapatac2.pl.render_plot`
        for details.

    Returns
    -------
    'plotly.graph_objects.Figure' | None
        If `show=False` and `out_file=None`, an `plotly.graph_objects.Figure` will be 
        returned, which can then be further customized using the plotly API.
    """
    from natsort import index_natsorted

    embedding = adata.obsm[use_rep] if is_anndata(adata) else adata
    if isinstance(color, str):
        groups = adata.obs[color].to_numpy()
    else:
        groups = color
        color = "color"
    
    if sample_size is not None and embedding.shape[0] > sample_size:
        idx = np.random.choice(embedding.shape[0], sample_size, replace=False)
        embedding = embedding[idx, :]
        if groups is not None: groups = groups[idx]

    if groups is not None:
        idx = index_natsorted(groups)
        embedding = embedding[idx, :]
        groups = [groups[i] for i in idx]

    if marker_size is None:
        num_points = embedding.shape[0]
        marker_size = (1000 / num_points)**(1/3) * 3

    if embedding.shape[1] >= 3:
        return scatter3d(embedding[:, 0], embedding[:, 1], embedding[:, 2], color=groups,
            x_label="UMAP-1", y_label="UMAP-2", z_label="UMAP-3", color_label=color,
            marker_size=marker_size, marker_opacity=marker_opacity, **kwargs)
    else:
        return scatter(embedding[:, 0], embedding[:, 1], color=groups,
            x_label="UMAP-1", y_label="UMAP-2", color_label=color,
            marker_size=marker_size, marker_opacity=marker_opacity, **kwargs)

def motif_enrichment(
    enrichment: list(str, 'pl.DataFrame'),
    min_log_fc: float = 1,
    max_fdr: float = 0.01,
    **kwargs,
) -> 'plotly.graph_objects.Figure' | None:
    """Plot the motif enrichment result.

    Parameters
    ----------
    enrichment
        Motif enrichment result.
    min_log_fc
        Retain motifs that satisfy: log2-fold-change >= `min_log_fc`.
    max_fdr
        Retain motifs that satisfy: FDR <= `max_fdr`.
    kwargs        
        Additional arguments passed to :func:`~snapatac2.pl.render_plot` to
        control the final plot output. Please see :func:`~snapatac2.pl.render_plot`
        for details.

    Returns
    -------
    'plotly.graph_objects.Figure' | None
        If `show=False` and `out_file=None`, an `plotly.graph_objects.Figure` will be 
        returned, which can then be further customized using the plotly API.
    """
 
    import pandas as pd
    
    fc = np.vstack([df['log2(fold change)'] for df in enrichment.values()])
    filter1 = np.apply_along_axis(lambda x: np.any(np.abs(x) >= min_log_fc), 0, fc)
    
    fdr = np.vstack([df['adjusted p-value'] for df in enrichment.values()])
    filter2 = np.apply_along_axis(lambda x: np.any(x <= max_fdr), 0, fdr)

    passed = np.logical_and(filter1, filter2)
    
    sign = np.sign(fc[:, passed])
    pvals = np.vstack([df['p-value'].to_numpy()[passed] for df in enrichment.values()])
    minval = np.min(pvals[np.nonzero(pvals)])
    pvals = np.clip(pvals, minval, None)
    pvals = sign * np.log(-np.log10(pvals))

    df = pd.DataFrame(
        pvals.T,
        columns=list(enrichment.keys()),
        index=next(iter(enrichment.values()))['id'].to_numpy()[passed],
    )
      
    return heatmap(
        df.to_numpy(),
        row_names=df.index,
        column_names=df.columns,
        colorscale='RdBu_r',
        **kwargs,
    )

def coverage(
    adata: AnnData,
    region: str,
    groupby: str | list[str],
    out_file: str | None = None,
):
    """
    Plot the coverage tracks for different groups of cells. This function requires
    `matplotlib` to be installed.

    This is a simple implementation for quick visualization. For more advanced
    visualization, please consider exporting the data using 
    :func:`~snapatac2.ex.export_coverage` and visualize it using a genome browser.

    Parameters
    ----------
    adata
        Annotated data matrix.
    region
        Genomic region, e.g. 'chr1:100000-200000'.
    groupby
        Group the cells into different groups. If a `str`, groups are obtained from
        `.obs[groupby]`.
    out_file
        Path of the output file for saving the output image.

    See Also
    --------
    export_coverage
    :func:`~snapatac2.ex.export_coverage` 
    """

    from matplotlib import pyplot as plt

    groupby = adata.obs[groupby] if isinstance(groupby, str) else groupby
    groupby = [x for x in groupby]
    signal_values = []
    track_names = []
    for k, v in sorted(list(internal.get_coverage(adata, region, groupby).items())):
        track_names.append(k)
        signal_values.append(v)
    signal_values = np.array(signal_values)

    start, end = region.split(":")[1].split("-")
    start = int(start)
    end = int(end)
    height_per_track = 1.2
    width = 6

    n_tracks, n_points = signal_values.shape
    fig, axes = plt.subplots(
        n_tracks, 1,
        figsize=[width, n_tracks * height_per_track],
        sharex=True,
        constrained_layout=True,
    )

    if n_tracks == 1:
        axes = [axes]

    # Compute global max for y-axis scaling
    global_max = signal_values.max()

    cmap = plt.get_cmap("tab10")
    for i, (ax, signal) in enumerate(zip(axes, signal_values)):
        color = cmap(i % 10)   # cycle through colors if >10 tracks
        ax.fill_between(range(n_points), 0, signal, color=color)
        ax.set_title(track_names[i], fontsize=7)
        ax.spines[['top', 'right']].set_visible(False)
        ax.set_ylim(0, global_max)

    axes[-1].set_xticks([0, n_points - 1])
    axes[-1].set_xticklabels([str(start), str(end)])
    axes[-1].set_xlabel(region)

    fig.supylabel("RPM")

    if out_file is None:
        plt.show()
    else:
        plt.savefig(out_file, dpi=300, bbox_inches='tight')