from __future__ import annotations
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from anndata import AnnData
    from typing import List, Literal
    from matplotlib.axes import Axes

import matplotlib.lines as mlines
import matplotlib.pyplot as plt
import numpy as np
import pysankey
import scanpy
import seaborn as sns

from matplotlib import gridspec
from matplotlib.markers import MarkerStyle
from packaging.version import Version


def remodeling_score(
    remodeling_score: np.ndarray,
    show: bool | None = None,
    save: bool | str | None = None,
) -> List[plt.Axes] | None:
    """Visualise the distribution of protein *remodeling scores*.

    The function creates a vertical layout containing a histogram (top) and a
    compact box plot (bottom).  It mirrors the visualisation published in
    Hein *et al.* 2024 and serves as a quick overview of score
    distribution and potential outliers.

    Parameters
    ----------
    remodeling_score
        One-dimensional array of remodeling scores, typically generated by
        :func:`~grassp.tl.remodeling_score`.
    show
        If ``True`` (default) the plot is shown and the function returns ``None``.
    save
        Path or boolean forwarded to :func:`scanpy.pl._utils.savefig_or_show`.
        If ``True``, a PNG file named ``remodeling_score.png`` is written.

    Returns
    -------
    When ``show`` is ``False``, the list ``[ax_hist, ax_box]`` is returned to
    enable further customisation.  Otherwise ``None``.
    """

    # Create grid layout
    gs = gridspec.GridSpec(2, 1, height_ratios=[5, 1])
    # Histogram
    ax0 = plt.subplot(gs[0])
    sns.histplot(remodeling_score, ax=ax0, kde=False)
    ax0.set(xlabel="")
    # turn off x tick labels
    ax0.set_xticklabels([])

    # Boxplot
    ax1 = plt.subplot(gs[1])
    sns.boxplot(
        x=remodeling_score,
        ax=ax1,
        flierprops=dict(
            marker="o", markeredgecolor="orange", markerfacecolor="none", markersize=6
        ),
    )
    ax1.set(xlabel="Remodeling score")
    axs = [ax0, ax1]

    show = scanpy.settings.autoshow if show is None else show
    scanpy.pl._utils.savefig_or_show("remodeling_score", show=show, save=save)
    if show:
        return None
    return axs


remodeling_legend = [
    mlines.Line2D(
        [], [], color="black", marker="*", linestyle="None", label="remodeled_proteins"
    ),
    mlines.Line2D(
        [], [], color="grey", linestyle="-", linewidth=2, label="remodeling trajectory"
    ),
]


def _get_cluster_colors(data: AnnData, color_key: str = "leiden") -> np.ndarray[str, Any]:
    """Get cluster colors for a given color key.

    Parameters
    ----------
    data
        Annotated data matrix
    color_key
        Key in data.obs to use for coloring

    Returns
    -------
    np.ndarray
        Array of colors for the given color key
    """
    if f"{color_key}_colors" not in data.uns.keys():

        if Version(scanpy.__version__) >= Version("1.12.0"):
            scanpy.pl._utils.set_default_colors_for_categorical_obs(  # pylint: disable=no-member
                data, color_key  # pylint: disable=no-member
            )  # pylint: disable=no-member
        else:
            scanpy.pl._utils._set_default_colors_for_categorical_obs(  # pylint: disable=no-member
                data, color_key  # pylint: disable=no-member
            )  # pylint: disable=no-member
    return np.array(
        [data.uns[f"{color_key}_colors"][x] for x in data.obs[color_key].cat.codes]
    )


def aligned_umap(
    data: AnnData,
    data2: AnnData,
    highlight_hits: List[str] | np.ndarray[bool, Any] | None = None,
    highlight_annotation_col: str | None = None,
    aligned_umap_key: str = "X_aligned_umap",
    data1_label: str = "data1",
    data2_label: str = "data2",
    color_by: Literal["perturbation", "cluster"] = "perturbation",
    data1_color: str = "#C7E8F9",
    data2_color: str = "#FFCCC2",
    figsize: tuple[float, float] = (8.25, 6),
    size: int = 80,
    alpha: float = 0.4,
    ax: plt.Axes | None = None,
    show: bool | None = None,
    save: bool | str | None = None,
) -> plt.Axes | None:
    """Side-by-side visualisation of aligned UMAP embeddings.

    Given two datasets embedded into a common aligned UMAP space (see
    :func:`~grassp.tl.aligned_umap`), the function plots both embeddings in a
    single scatter plot and optionally draws **remodeling trajectories** for
    highlighted proteins.

    Parameters
    ----------
    data, data2
        Two :class:`~anndata.AnnData` objects with the same set/order of
        observations and a shared key in ``.obsm`` (*aligned_umap_key*).
    highlight_hits
        Proteins to emphasise with a star marker and a line connecting their
        positions between the two datasets.  Can be a list of observation
        names or a boolean mask.
    highlight_annotation_col
        Optional column in ``data.obs`` used for text labels next to the
        highlighted points.
    aligned_umap_key
        Key in ``.obsm`` containing the aligned coordinates (default:
        ``"X_aligned_umap"``).
    data1_label, data2_label
        Legend labels for the two datasets.
    color_by
        If ``"perturbation"`` (default) both datasets are assigned uniform
        colors (``data1_color``, ``data2_color``).  If ``"cluster"`` cluster-
        specific colors defined in ``data.obs[color_key]`` are used.
    size, alpha, figsize, ax
        Standard matplotlib styling parameters.
    show, save
        Forwarded to :func:`scanpy.pl._utils.savefig_or_show`.

    Returns
    -------
    Returns the Axes object if ``show`` is ``False``.
    """

    if ax is None:
        fig, ax = plt.subplots(figsize=figsize)

    if color_by == "cluster":
        data1_colors = _get_cluster_colors(data, data1_color)
        data2_colors = _get_cluster_colors(data2, data2_color)

        # Create cluster legend handles
        cluster_handles = []
        unique_clusters = data.obs[data1_color].cat.categories
        cluster_colors = data.uns[f"{data1_color}_colors"]
        for cluster, color in zip(unique_clusters, cluster_colors):
            cluster_handles.append(
                mlines.Line2D(
                    [],
                    [],
                    color=color,
                    marker="o",
                    linestyle="None",
                    markersize=8,
                    label=cluster,
                )
            )
    else:
        data1_colors = data1_color
        data2_colors = data2_color

    embedding1 = data.obsm[aligned_umap_key]
    embedding2 = data2.obsm[aligned_umap_key]

    # Plot the two embeddings as scatter plots
    ax.scatter(
        np.asarray(embedding1)[:, 0],
        np.asarray(embedding1)[:, 1],
        c=data1_colors,
        s=size,
        alpha=alpha,
        label=data1_label,
        marker=MarkerStyle("."),
        linewidths=0,
        edgecolor=None,
    )
    ax.scatter(
        np.asarray(embedding2)[:, 0],
        np.asarray(embedding2)[:, 1],
        c=data2_colors,
        s=size,
        alpha=alpha,
        label=data2_label,
        marker=MarkerStyle("+"),
        linewidths=1,
        edgecolor=None,
    )

    if highlight_hits is not None:
        embedding1_hits = np.asarray(embedding1)[highlight_hits]
        embedding2_hits = np.asarray(embedding2)[highlight_hits]
        # Plot trajectory lines
        for i, (start, end) in enumerate(zip(embedding1_hits, embedding2_hits)):
            # Draw line
            ax.plot(
                [start[0], end[0]],
                [start[1], end[1]],
                color="grey",
                linewidth=0.7,
                alpha=0.5,
            )
            # Draw marker at the end point
            ax.scatter(
                start[0],
                start[1],
                color="black",
                s=30,
                marker=MarkerStyle("*"),
                edgecolor=None,
            )
            if highlight_annotation_col is not None:
                # Add annotation
                ax.annotate(
                    str(data.obs.loc[highlight_hits, highlight_annotation_col].iloc[i]),
                    (start[0], start[1]),
                    color="black",
                    fontsize=5,
                    bbox=dict(facecolor="white", edgecolor="none", alpha=0.7, pad=0.5),
                    ha="right",
                    va="bottom",
                    xytext=(5, 5),
                    textcoords="offset points",
                )

    # Combine scatter plot legend with remodeling legend
    handles, labels = ax.get_legend_handles_labels()
    combined_handles = handles + remodeling_legend

    # Add legends to the right of the plot
    if color_by == "cluster":
        # First legend for dataset and remodeling markers
        ax.legend(handles=combined_handles, bbox_to_anchor=(1.15, 1), loc="upper left")
        # Second legend for clusters
        ax.legend(
            handles=cluster_handles,
            bbox_to_anchor=(1.15, 0.5),
            loc="center left",
            title="Clusters",
        )
    else:
        # Single legend if not coloring by cluster
        ax.legend(handles=combined_handles, bbox_to_anchor=(1.15, 1), loc="upper left")

    # Adjust layout to prevent legend cutoff
    plt.tight_layout()

    ax.set_xlabel(f"{aligned_umap_key.replace('X_', '')}1")
    ax.set_ylabel(f"{aligned_umap_key.replace('X_', '')}2")
    ax.set_xticks([])
    ax.set_yticks([])

    show = scanpy.settings.autoshow if show is None else show
    scanpy.pl._utils.savefig_or_show("aligned_umap", show=show, save=save)
    if show:
        return None
    return ax


def remodeling_sankey(
    data: AnnData,
    data2: AnnData,
    cluster_key: str = "leiden",
    ax: plt.Axes | None = None,
    aspect: int = 20,
    fontsize: int = 12,
    figsize: tuple[float, float] = (10, 11),
    show: bool | None = None,
    save: bool | str | None = None,
) -> plt.Axes | None:
    """Sankey diagram of cluster re-assignment between two conditions.

    Parameters
    ----------
    data, data2
        Two :class:`~anndata.AnnData` objects with identical observation
        indices and a categorical *cluster_key* in ``.obs``.
    cluster_key
        Column name that encodes cluster memberships.
    aspect, fontsize, figsize, ax
        Styling options passed to :func:`pysankey.sankey` or matplotlib.
    show, save
        Behaviour identical to other plotting functions.

    Returns
    -------
    Axes containing the Sankey plot if ``show`` is ``False``.
    """
    if ax is None:
        fig, ax = plt.subplots(figsize=figsize)

    # Check that the anndata objects are aligned
    assert (data.obs_names == data2.obs_names).all()

    pysankey.sankey(
        left=data.obs[cluster_key],
        right=data2.obs[cluster_key],
        aspect=aspect,
        # colorDict=colorDict,
        fontsize=fontsize,
        color_gradient=False,
        # leftLabels=[
        #     "nucleus",
        #     "cytosol",
        #     "mitochondrion",
        #     "ER",
        #     "plasma memb. & actin",
        #     "endo-lysosome & trans-Golgi",
        #     "ERGIC/Golgi",
        #     "translation/RNA granules",
        #     "peroxisome",
        # ],
        # rightLabels=[
        #     "nucleus",
        #     "cytosol",
        #     "mitochondrion",
        #     "ER",
        #     "plasma memb. & actin",
        #     "endo-lysosome & trans-Golgi",
        #     "translation/RNA granules",
        #     "peroxisome",
        #     "COPI vesicle",
        # ],
        ax=ax,
    )

    show = scanpy.settings.autoshow if show is None else show
    scanpy.pl._utils.savefig_or_show("remodeling_sankey", show=show, save=save)
    if show:
        return None
    return ax


def mr_plot(
    data: AnnData,
    mr_key: str = "mr_scores",
    ax: plt.Axes = None,
    m_cutoffs: list[float] = [2, 3, 4],
    r_cutoffs: list[float] = [0.68, 0.81, 0.93],
    highlight_hits: bool = True,
    highlight_proteins: list[str] = [],
    show: bool | None = None,
    save: bool | str | None = None,
    **kwargs,
) -> Axes | None:
    """MR-plot for simultaneous visualisation of *M* and *R* scores.

    *M* corresponds to ``-log10(q-value)`` and reflects statistical
    significance; *R* is a measure of consistencey of the re-localisation between replicates.

    Parameters
    ----------
    data
        AnnData object after :func:`grassp.tl.mr_score`.
    mr_key
        Prefix of the keys written by ``mr_score`` (defaults to
        ``"mr_scores"`` resulting in ``mr_scores_M`` and ``mr_scores_R`` in
        ``.obs``).
    m_cutoffs, r_cutoffs
        Horizontal/vertical guideline positions (lenient → stringent).
    highlight_hits
        If ``True`` mark proteins passing the *lenient* thresholds in red.
    highlight_proteins
        List of proteins to highlight. Not compatible with ``highlight_hits``.
    ax, show, save, **kwargs
        Standard matplotlib/scanpy plotting options.

    Returns
    -------
    Returns the Axes object if ``show`` is ``False``.
    """
    if ax is None:
        fig, ax = plt.subplots(figsize=(10, 6))

    try:
        m_scores = data.obs[f"{mr_key}_M"]
        r_scores = data.obs[f"{mr_key}_R"]
    except KeyError:
        raise ValueError(
            f"MR scores not found in data.obs['{mr_key}_M/R'], run gr.tl.mr_score first"
        )

    # Plot data
    ax.scatter(m_scores, r_scores, alpha=0.5, s=10, color="black", marker=".", **kwargs)

    if highlight_proteins != []:
        # assert not highlight_hits, (
        #     "highlight_proteins and highlight_hits cannot be used together"
        # )
        hits = data.obs_names.isin(highlight_proteins)
    if highlight_hits:
        hits = (m_scores >= m_cutoffs[0]) & (r_scores >= r_cutoffs[0])
        ax.scatter(m_scores[hits], r_scores[hits], color="red", s=20, marker=".")

    # Add cutoff lines
    colors = ["gray", "darkgray", "lightgray"]
    for m_cut, color in zip(m_cutoffs, colors):
        ax.axvline(m_cut, color=color, linestyle="--", alpha=0.5)
    for r_cut, color in zip(r_cutoffs, colors):
        ax.axhline(r_cut, color=color, linestyle="--", alpha=0.5)

    # Set x-axis limits
    ax.set_xlim(0, np.max(m_scores) + 5)

    # Set labels
    ax.set_xlabel("M score (-log10 Q-value)")
    ax.set_ylabel("R score (minimum correlation)")
    ax.set_title("MR Plot")

    show = scanpy.settings.autoshow if show is None else show
    scanpy.pl._utils.savefig_or_show("remodeling_score", show=show, save=save)
    if show:
        return None
    return ax
