"""Trajectory/PAGA plotting utilities for scDEF."""

from typing import Optional, List, Tuple, Any, TYPE_CHECKING, Sequence, Union, Dict
import copy
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib as mpl
import scanpy as sc
from scipy.ndimage import uniform_filter1d
from sklearn.preprocessing import minmax_scale

from ..tools.trajectory import multilevel_paga as compute_multilevel_paga
from ..tools.factor import get_stored_confident_signatures

if TYPE_CHECKING:
    from scdef.models._scdef import scDEF


def multilevel_paga(
    model: "scDEF",
    neighbors_rep: Optional[str] = "X_L0",
    layers: Optional[List[int]] = None,
    figsize: Optional[Tuple[float, float]] = (16, 4),
    reuse_pos: Optional[bool] = True,
    recompute: bool = False,
    fontsize: Optional[int] = 12,
    show: Optional[bool] = True,
    **paga_kwargs: Any,
) -> None:
    """Plot cached multilevel PAGA graphs across scDEF layers."""
    if layers is None:
        layers = [
            i
            for i in range(model.n_layers - 1, -1, -1)
            if len(model.factor_lists[i]) > 1
        ]

    if len(layers) == 0:
        print("Cannot run PAGA on 0 layers.")
        return

    cache = model.adata.uns.get("multilevel_paga", None)
    cache_layers = [] if cache is None else cache.get("layers", [])
    need_recompute = (
        recompute
        or cache is None
        or cache.get("neighbors_rep") != neighbors_rep
        or cache.get("reuse_pos") != bool(reuse_pos)
        or any(int(layer) not in cache_layers for layer in layers)
    )
    if need_recompute:
        compute_multilevel_paga(
            model,
            neighbors_rep=neighbors_rep,
            layers=layers,
            reuse_pos=reuse_pos,
            **paga_kwargs,
        )
        cache = model.adata.uns.get("multilevel_paga", None)

    if cache is None or "results" not in cache:
        raise RuntimeError(
            "Multilevel PAGA cache not found. Run scdef.tl.multilevel_paga(model) first."
        )

    n_layers = len(layers)
    fig, axes = plt.subplots(1, n_layers, figsize=figsize)
    if n_layers == 1:
        axes = [axes]
    old_paga = copy.deepcopy(model.adata.uns.get("paga", None))
    layout = cache.get("layout", "fa")
    for i, layer_idx in enumerate(layers):
        ax = axes[i]
        layer_name = model.layer_names[layer_idx]
        if layer_name not in cache["results"]:
            raise KeyError(
                f"Cached PAGA for layer '{layer_name}' not found. Recompute with scdef.tl.multilevel_paga."
            )
        entry = cache["results"][layer_name]
        model.adata.uns["paga"] = copy.deepcopy(entry["paga"])
        sc.pl.paga(
            model.adata,
            init_pos=np.array(entry["pos"]),
            layout=layout,
            ax=ax,
            show=False,
            **paga_kwargs,
        )
        ax.set_title(f"Layer {layer_idx} PAGA", fontsize=fontsize)
    if old_paga is None:
        model.adata.uns.pop("paga", None)
    else:
        model.adata.uns["paga"] = old_paga
    if show:
        plt.show()


def plot_trajectory_heatmap(
    model: "scDEF",
    factor_path: Sequence[Union[int, str]],
    layer_idx: int = 0,
    l0_path: Optional[Sequence[Union[int, str]]] = None,
    genes_per_factor: Optional[int] = 3,
    smoothing: int = 50,
    figwidth: float = 8,
    gene_height: float = 0.28,
    block_spacing: int = 1,
    annotation_obs_key: Optional[Union[str, Sequence[str]]] = None,
    subset_obs_key: Optional[str] = None,
    subset_obs: Optional[Union[str, Sequence[str]]] = None,
    heatmap_cmap: str = "RdYlBu_r",
    factor_heatmap_cmap: str = "viridis",
    colorbar_gap: float = 0.16,
    xlabel: str = "Cells",
    save: Optional[str] = None,
    show: bool = True,
):
    """Plot stacked trajectory heatmap from precomputed confident signatures.

    Requires ``scd.tl.set_confident_signatures(model)`` to be run beforehand.

    Cell sorting:
        Cells are restricted to those assigned to ``factor_path`` (and optional
        ``subset_obs`` filter), then ordered by a soft trajectory progress score.
        For each selected cell, the function takes its factor probabilities
        (from ``X_<layer>_probs``), keeps only the columns for the path factors,
        and computes a probability-weighted average of path positions
        ``0..len(path)-1``. This gives a continuous progress coordinate used to
        sort cells from early to late along the path.

        If ``layer_idx > 0`` and ``l0_path`` is provided, the sorting weights are
        computed on layer 0 probabilities instead. If that yields near-zero total
        path weight for all selected cells, sorting falls back to the
        ``layer_idx`` path probabilities.
    """
    layer_name = model.layer_names[layer_idx]
    kept_names = list(model.factor_names[layer_idx])

    def _resolve_path_names(
        path: Sequence[Union[int, str]], layer_names: List[str], layer_idx_resolve: int
    ) -> List[str]:
        resolved: List[str] = []
        for f in path:
            if isinstance(f, str):
                if f not in layer_names:
                    raise ValueError(
                        f"Factor '{f}' not found in layer {layer_idx_resolve}."
                    )
                resolved.append(f)
            else:
                f_idx = int(f)
                if f_idx < 0 or f_idx >= len(layer_names):
                    raise IndexError(
                        f"Factor index {f_idx} out of bounds for layer {layer_idx_resolve}."
                    )
                resolved.append(layer_names[f_idx])
        return resolved

    path_names = _resolve_path_names(factor_path, kept_names, layer_idx)
    if len(path_names) == 0:
        raise ValueError("factor_path must contain at least one factor.")
    if block_spacing < 0:
        raise ValueError("block_spacing must be >= 0.")
    if colorbar_gap < 0.0:
        raise ValueError("colorbar_gap must be >= 0.")

    subset_mask = np.ones(model.adata.n_obs, dtype=bool)
    if subset_obs is not None and subset_obs_key is None:
        raise ValueError("subset_obs_key must be provided when subset_obs is set.")
    if subset_obs_key is not None:
        if subset_obs_key not in model.adata.obs.columns:
            raise KeyError(f"{subset_obs_key} not found in adata.obs.")
        if subset_obs is None:
            raise ValueError("subset_obs must be provided when subset_obs_key is set.")
        if isinstance(subset_obs, str):
            subset_vals = [subset_obs]
        else:
            subset_vals = [str(v) for v in subset_obs]
        if len(subset_vals) == 0:
            raise ValueError("subset_obs must contain at least one value.")
        subset_mask = (
            model.adata.obs[subset_obs_key].astype(str).isin(subset_vals).values
        )

    path_mask = model.adata.obs[layer_name].isin(path_names).values & subset_mask
    if np.count_nonzero(path_mask) == 0:
        raise ValueError(
            "No cells found for the provided factor_path after applying subset filters."
        )

    sort_layer_name = layer_name
    sort_names = path_names
    sort_kept_names = kept_names
    use_l0_sort = layer_idx > 0 and l0_path is not None
    if use_l0_sort:
        l0_kept_names = list(model.factor_names[0])
        sort_layer_name = model.layer_names[0]
        sort_names = _resolve_path_names(l0_path, l0_kept_names, 0)
        if len(sort_names) == 0:
            raise ValueError("l0_path must contain at least one factor.")
        sort_kept_names = l0_kept_names

    X_probs = np.asarray(model.adata.obsm[f"X_{sort_layer_name}_probs"], dtype=float)
    factor_pos = {name: idx for idx, name in enumerate(sort_kept_names)}
    path_cols = np.array([factor_pos[name] for name in sort_names], dtype=int)
    path_weights = X_probs[path_mask][:, path_cols]
    if np.all(np.sum(path_weights, axis=1) <= 1e-12):
        if use_l0_sort and hasattr(model, "logger"):
            model.logger.warning(
                "l0_path provided but selected cells have near-zero weight on this "
                "path. Falling back to layer %s factor_path ordering.",
                layer_idx,
            )
        X_probs_fb = np.asarray(model.adata.obsm[f"X_{layer_name}_probs"], dtype=float)
        factor_pos_fb = {name: idx for idx, name in enumerate(kept_names)}
        path_cols_fb = np.array([factor_pos_fb[name] for name in path_names], dtype=int)
        path_weights = X_probs_fb[path_mask][:, path_cols_fb]
    denom = np.maximum(np.sum(path_weights, axis=1), 1e-12)
    ranks = np.arange(path_weights.shape[1], dtype=float)
    progress = np.sum(path_weights * ranks[None, :], axis=1) / denom
    selected_cells = np.where(path_mask)[0]
    sorted_cells = selected_cells[np.argsort(progress)]
    t_path = model.adata[sorted_cells].copy()

    confident_sigs = get_stored_confident_signatures(
        model, layer_idx=layer_idx, max_genes=genes_per_factor
    )

    score_cols = [f"{name}_score" for name in path_names]
    missing_scores = [col for col in score_cols if col not in t_path.obs.columns]
    if len(missing_scores) > 0:
        raise KeyError(f"Score columns not found in adata.obs: {missing_scores}")

    block_matrices: List[Tuple[np.ma.MaskedArray, np.ma.MaskedArray]] = []
    block_labels: List[List[str]] = []
    block_has_genes: List[bool] = []
    for factor_name in path_names:
        block_rows: List[np.ndarray] = []
        labels: List[str] = []
        kinds: List[str] = []
        score_col = f"{factor_name}_score"
        score_vals = uniform_filter1d(
            np.asarray(t_path.obs[score_col].values, dtype=float), size=smoothing
        )
        block_rows.append(minmax_scale(score_vals))
        labels.append(factor_name)
        kinds.append("factor")

        genes = [
            g for g in confident_sigs.get(factor_name, []) if g in t_path.var_names
        ]
        for gene in genes:
            expr = t_path[:, [gene]].X
            if hasattr(expr, "toarray"):
                expr = expr.toarray()
            expr_vals = uniform_filter1d(
                np.asarray(expr, dtype=float).ravel(), size=smoothing
            )
            block_rows.append(minmax_scale(expr_vals))
            labels.append(f"  {gene}")
            kinds.append("gene")

        block_arr = np.vstack(block_rows)
        factor_rows = np.asarray([k == "factor" for k in kinds], dtype=bool)
        gene_rows = np.asarray([k == "gene" for k in kinds], dtype=bool)

        genes_masked = np.ma.array(
            block_arr, mask=np.broadcast_to(~gene_rows[:, None], block_arr.shape)
        )
        factors_masked = np.ma.array(
            block_arr, mask=np.broadcast_to(~factor_rows[:, None], block_arr.shape)
        )

        block_matrices.append((genes_masked, factors_masked))
        block_labels.append(labels)
        block_has_genes.append(bool(np.any(gene_rows)))

    if len(block_matrices) == 0:
        raise ValueError("No rows to plot. Check factor_path and confident signatures.")

    obs_keys: List[str] = []
    if annotation_obs_key is not None:
        if isinstance(annotation_obs_key, str):
            obs_keys = [annotation_obs_key]
        else:
            obs_keys = [str(k) for k in annotation_obs_key]
        if len(obs_keys) == 0:
            raise ValueError(
                "annotation_obs_key must be a non-empty string or sequence of strings."
            )

    obs_tracks: List[Dict[str, object]] = []
    for key in obs_keys:
        if key not in t_path.obs.columns:
            raise KeyError(f"{key} not found in adata.obs.")
        cat = pd.Categorical(t_path.obs[key])
        categories = cat.categories.tolist()
        uns_key = f"{key}_colors"
        if uns_key in t_path.uns:
            cat_to_color = dict(
                zip(
                    t_path.uns.get(f"{key}_categories", categories),
                    t_path.uns[uns_key],
                )
            )
        else:
            cmap_fb = plt.get_cmap("tab10", max(len(categories), 1))
            cat_to_color = {c: cmap_fb(i) for i, c in enumerate(categories)}
        rgb = np.array([mpl.colors.to_rgb(cat_to_color[c]) for c in t_path.obs[key]])
        obs_tracks.append(
            {
                "key": key,
                "categories": categories,
                "cat_to_color": cat_to_color,
                "rgb": rgb,
            }
        )

    nrows = len(obs_tracks) + len(block_matrices)
    height_ratios = [0.3] * len(obs_tracks)
    for labels in block_labels:
        height_ratios.append(len(labels) * gene_height)

    total_height = sum(height_ratios) + 1.5
    fig = plt.figure(figsize=(figwidth, total_height))
    gs = fig.add_gridspec(
        nrows=nrows,
        ncols=4,
        height_ratios=height_ratios,
        width_ratios=[figwidth - 1.6, 0.22, 0.24, 0.22],
        hspace=0.05 + 0.08 * block_spacing,
        wspace=colorbar_gap,
    )

    row = 0
    for obs in obs_tracks:
        ax_obs = fig.add_subplot(gs[row, 0])
        ax_obs_cb1 = fig.add_subplot(gs[row, 1])
        ax_obs_cb_gap = fig.add_subplot(gs[row, 2])
        ax_obs_cb2 = fig.add_subplot(gs[row, 3])
        ax_obs_cb1.axis("off")
        ax_obs_cb_gap.axis("off")
        ax_obs_cb2.axis("off")
        ax_obs.imshow(
            obs["rgb"][np.newaxis, :, :], aspect="auto", interpolation="nearest"
        )
        ax_obs.set_yticks([0])
        ax_obs.set_yticklabels([obs["key"]], fontsize=9)
        ax_obs.set_xticks([])
        legend_patches = [
            mpl.patches.Patch(color=obs["cat_to_color"][c], label=c)
            for c in obs["categories"]
        ]
        ax_obs.legend(
            handles=legend_patches,
            fontsize=8,
            frameon=False,
            loc="lower left",
            bbox_to_anchor=(0.0, 1.05),
            ncol=max(1, min(len(obs["categories"]), 8)),
            borderaxespad=0,
        )
        row += 1

    hm_row_start = row
    im_genes_ref = None
    im_factors_ref = None
    any_gene_rows = False
    for i, labels in enumerate(block_labels):
        ax_hm = fig.add_subplot(gs[hm_row_start + i, 0])
        genes_masked, factors_masked = block_matrices[i]

        im_genes = ax_hm.imshow(
            genes_masked, aspect="auto", cmap=heatmap_cmap, interpolation="nearest"
        )
        im_factors = ax_hm.imshow(
            factors_masked,
            aspect="auto",
            cmap=factor_heatmap_cmap,
            interpolation="nearest",
        )
        ax_hm.set_yticks(range(len(labels)))
        ax_hm.set_yticklabels(labels, fontsize=8)
        if i == len(block_labels) - 1:
            ax_hm.set_xlabel(xlabel, fontsize=10)
        ax_hm.set_xticks([])

        im_factors_ref = im_factors
        if block_has_genes[i]:
            im_genes_ref = im_genes
            any_gene_rows = True

        if i < len(block_labels) - 1:
            ax_hm.axhline(len(labels) - 0.5, color="white", linewidth=1.0, alpha=0.8)

    ax_hm_cb_gene = fig.add_subplot(gs[hm_row_start:, 1])
    ax_hm_cb_factor = fig.add_subplot(gs[hm_row_start:, 3])
    if any_gene_rows and im_genes_ref is not None:
        plt.colorbar(
            im_genes_ref, cax=ax_hm_cb_gene, label="Smoothed gene expression in path"
        )
        ax_hm_cb_gene.yaxis.labelpad = 10
        ax_hm_cb_gene.yaxis.label.set_size(8)
        ax_hm_cb_gene.tick_params(labelsize=7)
    else:
        ax_hm_cb_gene.axis("off")
    if im_factors_ref is not None:
        plt.colorbar(
            im_factors_ref, cax=ax_hm_cb_factor, label="Smoothed factor score in path"
        )
        ax_hm_cb_factor.yaxis.labelpad = 10
        ax_hm_cb_factor.yaxis.label.set_size(8)
        ax_hm_cb_factor.tick_params(labelsize=7)
    else:
        ax_hm_cb_factor.axis("off")

    if save:
        plt.savefig(save, bbox_inches="tight", dpi=150)
    if show:
        plt.show()
        return None
    return fig
