import collections
import enum
import math
import os
import pathlib
from typing import Any, Dict, List, Literal, Optional, Sequence, Tuple, Union

import numpy as np
import pandas as pd
from scipy.ndimage.filters import convolve

import matplotlib as mpl
import matplotlib.pyplot as plt
import seaborn as sns
from matplotlib import cm, colors
from matplotlib.ticker import FormatStrFormatter
from mpl_toolkits.axes_grid1.inset_locator import inset_axes

from anndata import AnnData

from cellrank import logging as logg
from cellrank._utils import Lineage
from cellrank._utils._docs import d, inject_docs
from cellrank._utils._enum import DEFAULT_BACKEND, Backend_t, ModeEnum
from cellrank._utils._parallelize import _get_n_cores
from cellrank._utils._utils import (
    _check_collection,
    _genesymbols,
    _min_max_scale,
    _unique_order_preserving,
    save_fig,
    valuedispatch,
)
from cellrank.pl._utils import (
    _callback_type,
    _create_callbacks,
    _create_models,
    _fit_bulk,
    _get_backend,
    _get_categorical_colors,
    _input_model_type,
    _return_model_type,
    _time_range_type,
)

__all__ = ["heatmap"]

_N_XTICKS = 10


class HeatmapMode(ModeEnum):
    GENES = enum.auto()
    LINEAGES = enum.auto()


@d.dedent
@inject_docs(m=HeatmapMode)
@_genesymbols
def heatmap(
    adata: AnnData,
    model: _input_model_type,
    genes: Sequence[str],
    time_key: str,
    lineages: Optional[Union[str, Sequence[str]]] = None,
    backward: bool = False,
    mode: Literal["genes", "lineages"] = HeatmapMode.LINEAGES,
    time_range: Optional[Union[_time_range_type, List[_time_range_type]]] = None,
    callback: _callback_type = None,
    cluster_key: Optional[Union[str, Sequence[str]]] = None,
    show_fate_probabilities: bool = False,
    cluster_genes: bool = False,
    keep_gene_order: bool = False,
    scale: bool = True,
    n_convolve: Optional[int] = 5,
    show_all_genes: bool = False,
    cbar: bool = True,
    lineage_height: float = 0.33,
    fontsize: Optional[float] = None,
    xlabel: Optional[str] = None,
    title: Optional[str] = None,
    cmap: colors.ListedColormap = cm.viridis,
    dendrogram: bool = True,
    return_genes: bool = False,
    return_models: bool = False,
    return_figure: bool = False,
    n_jobs: Optional[int] = 1,
    backend: Backend_t = DEFAULT_BACKEND,
    show_progress_bar: bool = True,
    figsize: Optional[Tuple[float, float]] = None,
    dpi: Optional[int] = None,
    save: Optional[Union[str, pathlib.Path]] = None,
    gene_order: Optional[Sequence[str]] = None,
    **kwargs: Any,
) -> Optional[Union[Dict[str, pd.DataFrame], Tuple[_return_model_type, Dict[str, pd.DataFrame]]]]:
    """Plot a heatmap of smoothed gene expression along specified lineages.

    .. seealso::
        - See :doc:`../../../notebooks/tutorials/estimators/800_gene_trends` on how to
          visualize the gene trends.

    This requires a pseudotemporal ordering of cellular dynamics, computed using a method like
    DPT :cite:`haghverdi:16` or Palantir :cite:`setty:19`. The function combines the pseudotemporal
    ordering with CellRank's fate probabilities to visualize each gene's expression along specific
    trajectories. In the heatmap, genes are ordered according to their peak in pseudotime, which
    emphasizes expression cascades of sequential activation.

    Parameters
    ----------
    %(adata)s
    %(model)s
    %(genes)s
    time_key
        Key in :attr:`~anndata.AnnData.obs` where the pseudotime is stored.
    lineages
        Names of the lineages for which to plot. If :obj:`None`, plot all lineages.
    %(backward)s
    mode
        Valid options are:

        - ``{m.LINEAGES!r}`` - group by ``genes`` for each lineage in ``lineages``.
        - ``{m.GENES!r}`` - group by ``lineages`` for each gene in ``genes``.
    %(time_range)s
        This can also be specified on per-lineage basis.
    %(gene_symbols)s
    %(model_callback)s
    cluster_key
        Key in :attr:`~anndata.AnnData.obs` containing categorical observations to be plotted on the top of heatmap.
        Only available when ``mode = {m.LINEAGES!r}``.
    show_fate_probabilities
        Whether to also plot fate probabilities alongside the smoothed expression.
        Only available when ``mode = {m.LINEAGES!r}``.
    cluster_genes
        Whether to cluster genes using :func:`~seaborn.clustermap` when ``mode = 'lineages'``.
    keep_gene_order
        Whether to keep the gene order for later lineages after the first was sorted.
        Only available when ``cluster_genes = False`` and ``mode = {m.LINEAGES!r}``.
    scale
        Whether to normalize the gene expression to :math:`[0, 1]`.
    n_convolve
        Size of the convolution window when smoothing fate probabilities.
    show_all_genes
        Whether to show all genes on y-axis.
    cbar
        Whether to show the colorbar.
    lineage_height
        Height of a bar when ``mode = {m.GENES!r}``.
    fontsize
        Size of the title's font.
    xlabel
        Label on the x-axis. If :obj:`None`, it is determined based on ``time_key``.
    title
        Title of the figure.
    cmap
        Colormap to use when visualizing the smoothed expression.
    dendrogram
        Whether to show dendrogram when ``cluster_genes = True``.
    return_genes
        Whether to return the sorted or clustered genes. Only available when ``mode = {m.LINEAGES!r}``.
    %(return_models)s
    return_figure
        Whether to return the figure object. Sets ``return_genes = True``
    %(parallel)s
    %(plotting)s
    kwargs
        Keyword arguments for :meth:`~cellrank.models.BaseModel.prepare`.

    Returns
    -------
    %(plots_or_returns_models)s

    - If ``return_genes = True`` and ``mode = {m.LINEAGES!r}``, returns a :class:`~pandas.DataFrame`
    containing the clustered or sorted genes.
    - If ``return_figure = True``, returns a tuple containing the figure and genes.
    """

    def find_indices(series: pd.Series, values) -> List[int]:
        def find_nearest(array: np.ndarray, value: float) -> int:
            ix = np.searchsorted(array, value, side="left")
            if ix > 0 and (ix == len(array) or math.fabs(value - array[ix - 1]) < math.fabs(value - array[ix])):
                return int(ix - 1)
            return int(ix)

        series = series.sort_values(ascending=True)
        return list(series.iloc[[find_nearest(series.values, v) for v in values]].index)

    def subset_lineage(lname: str, rng: np.ndarray) -> np.ndarray:
        time_series = adata.obs[time_key]
        ixs = find_indices(time_series, rng)
        lin = Lineage.from_adata(adata[ixs], backward=backward)
        lin = lin[lname].X.squeeze(1).copy()
        if n_convolve is None:
            return lin.copy()
        return convolve(lin, np.ones(n_convolve) / n_convolve, mode="nearest")

    def create_col_colors(lname: str, rng: np.ndarray) -> Tuple[np.ndarray, colors.Colormap, colors.Normalize]:
        color = probs[lname].colors[0]
        lin = subset_lineage(lname, rng)

        h, _, v = colors.rgb_to_hsv(colors.to_rgb(color))
        end_color = colors.hsv_to_rgb([h, 1, v])

        lineage_cmap = colors.LinearSegmentedColormap.from_list("lineage_cmap", ["#ffffff", end_color], N=len(rng))
        norm = colors.Normalize(vmin=np.min(lin), vmax=np.max(lin))
        scalar_map = cm.ScalarMappable(cmap=lineage_cmap, norm=norm)

        return (
            np.array([colors.to_hex(c) for c in scalar_map.to_rgba(lin)]),
            lineage_cmap,
            norm,
        )

    def create_col_categorical_color(cluster_key: str, rng: np.ndarray) -> np.ndarray:
        cols, mapper = _get_categorical_colors(adata, cluster_key)
        ixs = find_indices(adata.obs[time_key], rng)

        return np.array([colors.to_hex(mapper[v]) for v in adata[ixs, :].obs[cluster_key].values])

    def create_cbar(
        ax,
        x_delta: float,
        cmap: colors.Colormap,
        norm: colors.Normalize,
        label: Optional[str] = None,
    ) -> plt.Axes:
        cax = inset_axes(
            ax,
            width="1%",
            height="100%",
            loc="lower right",
            bbox_to_anchor=(x_delta, 0, 1, 1),
            bbox_transform=ax.transAxes,
        )

        _ = mpl.colorbar.ColorbarBase(
            cax,
            cmap=cmap,
            norm=norm,
            label=label,
            ticks=np.linspace(norm.vmin, norm.vmax, 5),
        )

        return cax

    @valuedispatch
    def _plot_heatmap(mode: HeatmapMode) -> plt.Figure:
        raise NotImplementedError(mode.value)

    @_plot_heatmap.register(HeatmapMode.GENES)
    def _() -> Tuple[plt.Figure, None]:
        def color_fill_rec(ax, xs, y1, y2, colors=None, cmap=cmap, **kwargs) -> None:
            colors = colors if cmap is None else cmap(colors)

            x = 0
            for i, (color, x, y1, y2) in enumerate(zip(colors, xs, y1, y2)):  # noqa: B020
                dx = (xs[i + 1] - xs[i]) if i < len(x) else (xs[-1] - xs[-2])
                ax.add_patch(plt.Rectangle((x, y1), dx, y2 - y1, color=color, ec=color, **kwargs))

            ax.plot(x, y2, lw=0)

        fig, axes = plt.subplots(
            nrows=len(genes) + show_fate_probabilities,
            figsize=(12, len(genes) + len(lineages) * lineage_height) if figsize is None else figsize,
            dpi=dpi,
            constrained_layout=True,
        )
        axes = np.ravel([axes])

        if show_fate_probabilities:
            data["fate probability"] = data[next(iter(data.keys()))]

        for ax, (gene, models) in zip(axes, data.items()):
            if scale:
                vmin, vmax = 0, 1
            else:
                c = np.array([m.y_test for m in models.values()])
                vmin, vmax = np.nanmin(c), np.nanmax(c)

            norm = colors.Normalize(vmin=vmin, vmax=vmax)

            ix = 0
            ys = [ix]

            if gene == "fate probability":
                norm = colors.Normalize(vmin=0, vmax=1)
                for ln, x in ((ln, m.x_test) for ln, m in models.items()):
                    y = np.ones_like(x)
                    c = subset_lineage(ln, x.squeeze())

                    color_fill_rec(ax, x, y * ix, y * (ix + lineage_height), colors=norm(c))

                    ix += lineage_height
                    ys.append(ix)
            else:
                for x, c in ((m.x_test, m.y_test) for m in models.values()):
                    y = np.ones_like(x)
                    c = _min_max_scale(c) if scale else c

                    color_fill_rec(ax, x, y * ix, y * (ix + lineage_height), colors=norm(c))

                    ix += lineage_height
                    ys.append(ix)

            xs = np.array([m.x_test for m in models.values()])
            x_min, x_max = np.min(xs), np.max(xs)
            ax.set_xticks(np.linspace(x_min, x_max, _N_XTICKS))

            ax.set_yticks(np.array(ys[:-1]) + lineage_height / 2)
            ax.spines["left"].set_position(("data", 0))  # move the left spine to the rectangles to get nicer yticks
            ax.set_yticklabels(models.keys(), ha="right")

            ax.set_title(gene, fontdict={"fontsize": fontsize})
            ax.set_ylabel("lineage")

            for pos in ["top", "bottom", "left", "right"]:
                ax.spines[pos].set_visible(False)

            if cbar:
                cax, _ = mpl.colorbar.make_axes(ax)
                _ = mpl.colorbar.ColorbarBase(
                    cax,
                    ticks=np.linspace(vmin, vmax, 5),
                    norm=norm,
                    cmap=cmap,
                    label="value" if gene == "fate probability" else "scaled expression" if scale else "expression",
                )

            ax.tick_params(
                top=False,
                bottom=False,
                left=True,
                right=False,
                labelleft=True,
                labelbottom=False,
            )

        ax.xaxis.set_major_formatter(FormatStrFormatter("%.3f"))
        ax.tick_params(
            top=False,
            bottom=True,
            left=True,
            right=False,
            labelleft=True,
            labelbottom=True,
        )
        ax.set_xlabel(xlabel)
        if title is not None:
            ax.set_title(title)

        return fig, None

    @_plot_heatmap.register(HeatmapMode.LINEAGES)
    def _(gene_order: Optional[Sequence[str]] = None) -> Tuple[List[plt.Figure], pd.DataFrame]:
        data_t = collections.defaultdict(dict)  # transpose
        for gene, lns in data.items():
            for ln, y in lns.items():
                data_t[ln][gene] = y

        figs = []
        sorted_genes = pd.DataFrame() if return_genes else None

        for lname, models in data_t.items():
            xs = np.array([m.x_test for m in models.values()])
            x_min, x_max = np.nanmin(xs), np.nanmax(xs)

            df = pd.DataFrame([m.y_test for m in models.values()], index=models.keys())
            df.index.name = "genes"

            if not cluster_genes:
                if gene_order is not None:
                    gene_order = [gene for gene in gene_order if gene in df.index]
                    df = df.loc[gene_order]
                else:
                    max_sort = np.argsort(np.argmax(df.apply(_min_max_scale, axis=1).values, axis=1))
                    df = df.iloc[max_sort, :]
                    if keep_gene_order:
                        gene_order = df.index

            cat_colors = None
            if cluster_key is not None:
                cat_colors = np.stack(
                    [create_col_categorical_color(c, np.linspace(x_min, x_max, df.shape[1])) for c in cluster_key],
                    axis=0,
                )

            if show_fate_probabilities:
                col_colors, col_cmap, col_norm = create_col_colors(lname, np.linspace(x_min, x_max, df.shape[1]))
                if cat_colors is not None:
                    col_colors = np.vstack([cat_colors, col_colors[None, :]])
            else:
                col_colors, col_cmap, col_norm = cat_colors, None, None

            row_cluster = cluster_genes and df.shape[0] > 1
            show_clust = row_cluster and dendrogram

            g = sns.clustermap(
                data=df,
                cmap=cmap,
                figsize=(10, min(len(genes) / 8 + 1, 10)) if figsize is None else figsize,
                xticklabels=False,
                row_cluster=row_cluster,
                col_colors=col_colors,
                colors_ratio=0,
                col_cluster=False,
                cbar_pos=None,
                yticklabels=show_all_genes or "auto",
                standard_scale=0 if scale else None,
            )

            if cbar:
                cax = create_cbar(
                    g.ax_heatmap,
                    0.1,
                    cmap=cmap,
                    norm=colors.Normalize(
                        vmin=0 if scale else np.min(df.values),
                        vmax=1 if scale else np.max(df.values),
                    ),
                    label="scaled expression" if scale else "expression",
                )
                g.fig.add_axes(cax)

                if col_cmap is not None and col_norm is not None:
                    cax = create_cbar(
                        g.ax_heatmap,
                        0.25,
                        cmap=col_cmap,
                        norm=col_norm,
                        label="fate probability",
                    )
                    g.fig.add_axes(cax)

            if g.ax_col_colors:
                main_bbox = _get_ax_bbox(g.fig, g.ax_heatmap)
                n_bars = show_fate_probabilities + (len(cluster_key) if cluster_key is not None else 0)
                _set_ax_height_to_cm(
                    g.fig,
                    g.ax_col_colors,
                    height=min(5, max(n_bars * main_bbox.height / len(df), 0.25 * n_bars)),
                )
                g.ax_col_colors.set_title(lname, fontdict={"fontsize": fontsize})
            else:
                g.ax_heatmap.set_title(lname, fontdict={"fontsize": fontsize})

            g.ax_col_dendrogram.set_visible(False)  # gets rid of top free space
            g.ax_heatmap.yaxis.tick_left()
            g.ax_heatmap.yaxis.set_label_position("right")
            # fmt: off
            g.ax_heatmap.set_xlabel(xlabel)
            g.ax_heatmap.set_xticks(np.linspace(0, len(df.columns), _N_XTICKS))
            g.ax_heatmap.set_xticklabels([round(n, 3) for n in np.linspace(x_min, x_max, _N_XTICKS)])
            if title is not None:
                g.ax_heatmap.set_title(title)
            # fmt: on
            if show_clust:
                # robustly show dendrogram, because gene names can be long
                g.ax_row_dendrogram.set_visible(True)
                dendro_box = g.ax_row_dendrogram.get_position()
                pad = 0.005
                bb = g.ax_heatmap.yaxis.get_tightbbox(g.fig.canvas.get_renderer()).transformed(
                    g.fig.transFigure.inverted()
                )
                dendro_box.x0 = bb.x0 - dendro_box.width - pad
                dendro_box.x1 = bb.x0 - pad
                g.ax_row_dendrogram.set_position(dendro_box)
            else:
                g.ax_row_dendrogram.set_visible(False)

            if return_genes:
                sorted_genes[lname] = (
                    df.index[g.dendrogram_row.reordered_ind]
                    if hasattr(g, "dendrogram_row") and g.dendrogram_row is not None
                    else df.index
                )

            figs.append(g)

        return figs, sorted_genes

    mode = HeatmapMode(mode)

    probs = Lineage.from_adata(adata, backward=backward)
    if lineages is None:
        lineages = probs.names
    elif isinstance(lineages, str):
        lineages = [lineages]
    lineages = _unique_order_preserving(lineages)
    probs = probs[lineages]

    if cluster_key is not None:
        if isinstance(cluster_key, str):
            cluster_key = [cluster_key]
        cluster_key = _unique_order_preserving(cluster_key)

    if isinstance(genes, str):
        genes = [genes]
    genes = _unique_order_preserving(genes)
    _check_collection(adata, genes, "var_names", use_raw=kwargs.get("use_raw", False))

    kwargs["backward"] = backward
    kwargs["time_key"] = time_key
    models = _create_models(model, genes, lineages)
    all_models, data, genes, lineages = _fit_bulk(
        models,
        _create_callbacks(adata, callback, genes, lineages, **kwargs),
        genes,
        lineages,
        time_range,
        return_models=True,  # always return (better error messages)
        filter_all_failed=True,
        parallel_kwargs={
            "show_progress_bar": show_progress_bar,
            "n_jobs": _get_n_cores(n_jobs, len(genes)),
            "backend": _get_backend(models, backend),
        },
        **kwargs,
    )

    xlabel = time_key if xlabel is None else xlabel

    logg.debug(f"Plotting `{mode!r}` heatmap")
    if gene_order is not None and mode != HeatmapMode.LINEAGES:
        gene_order = [gene for gene in gene_order if gene in genes]

    if mode == HeatmapMode.LINEAGES:
        fig, genes = _plot_heatmap(mode, gene_order)
    else:
        fig, genes = _plot_heatmap(mode)

    if return_figure:
        return fig, genes

    if save is not None and fig is not None:
        if isinstance(fig, plt.Figure):
            save_fig(fig, save)
        elif len(fig) == 1:
            save_fig(fig[0], save)
        else:
            for ln, f in zip(lineages, fig):
                save_fig(f, os.path.join(save, f"lineage_{ln}"))

    if return_genes and mode == HeatmapMode.LINEAGES:
        return (all_models, genes) if return_models else genes
    if return_models:
        return all_models


def _get_ax_bbox(fig: plt.Figure, ax: plt.Axes):
    return ax.get_window_extent().transformed(fig.dpi_scale_trans.inverted())


def _set_ax_height_to_cm(fig: plt.Figure, ax: plt.Axes, height: float) -> None:
    from mpl_toolkits.axes_grid1 import Divider, Size

    height /= 2.54  # cm to inches

    bbox = _get_ax_bbox(fig, ax)

    hori = [Size.Fixed(bbox.x0), Size.Fixed(bbox.width), Size.Fixed(bbox.x1)]
    vert = [Size.Fixed(bbox.y0), Size.Fixed(height), Size.Fixed(bbox.y1)]

    divider = Divider(fig, (0.0, 0.0, 1.0, 1.0), hori, vert, aspect=False)

    ax.set_axes_locator(divider.new_locator(nx=1, ny=1))
