"""Factor plotting functions for scDEF.

This module provides factor-related plotting functions for scDEF models.
"""

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from matplotlib.lines import Line2D
import scanpy as sc
import copy
from scipy.cluster.hierarchy import ward, leaves_list
from scipy.spatial.distance import pdist
from scipy.stats import norm
from typing import (
    Optional,
    Sequence,
    Literal,
    Mapping,
    Dict,
    List,
    Tuple,
    Any,
    Union,
    TYPE_CHECKING,
)
from matplotlib.figure import Figure
from matplotlib.axes import Axes
import pandas as pd
from ..utils import data_utils

if TYPE_CHECKING:
    from scdef.models._scdef import scDEF


def obs_factor_dotplot(
    model: "scDEF",
    obs_key: str,
    layer_idx: int,
    cluster_rows: bool = True,
    cluster_cols: bool = True,
    figsize: Tuple[float, float] = (8, 2),
    s_min: int = 100,
    s_max: int = 500,
    titlesize: int = 12,
    labelsize: int = 12,
    legend_fontsize: int = 12,
    legend_titlesize: int = 12,
    cmap: str = "viridis",
    logged: bool = False,
    width_ratios: List[float] = [5, 1, 1],
    show_ylabel: bool = True,
    show: bool = True,
) -> Optional[Figure]:
    """Plot dotplot showing factor assignments for observations.

    Args:
        model: scDEF model instance
        obs_key: key in model.adata.obs to use for grouping
        layer_idx: layer index to plot
        cluster_rows: whether to cluster rows
        cluster_cols: whether to cluster columns
        figsize: figure size
        s_min: minimum circle size
        s_max: maximum circle size
        titlesize: title font size
        labelsize: label font size
        legend_fontsize: legend font size
        legend_titlesize: legend title font size
        cmap: colormap name
        logged: whether to log transform colors
        width_ratios: width ratios for subplots
        show_ylabel: whether to show y-axis label
        show: whether to show the plot

    Returns:
        Figure object if show is False, None otherwise
    """
    # For each obs, compute the average cell score on each factor among the cells that attach to that obs, use as color
    # And compute the fraction of cells in the obs that attach to each factor, use as circle size
    layer_name = model.layer_names[layer_idx]

    obs = model.adata.obs[obs_key].unique()
    n_obs = len(obs)
    n_factors = len(model.factor_lists[layer_idx])

    df_rows = []  # noqa: F841
    c = np.zeros((n_obs, n_factors))
    s = np.zeros((n_obs, n_factors))
    for i, obs_val in enumerate(obs):
        cells_from_obs = model.adata.obs.index[
            np.where(model.adata.obs[obs_key] == obs_val)[0]
        ]
        n_cells_obs = len(cells_from_obs)
        for factor in range(n_factors):
            factor_name = model.factor_names[layer_idx][factor]
            cells_attached = model.adata.obs.index[
                np.where(
                    model.adata.obs.loc[cells_from_obs][f"{layer_name}"] == factor_name
                )[0]
            ]
            if len(cells_attached) == 0:
                average_weight = 0  # np.nan
                fraction_attached = 0  # np.nan
            else:
                average_weight = np.mean(
                    model.adata.obs.loc[cells_from_obs][f"{factor_name}_score"]
                )
                fraction_attached = len(cells_attached) / n_cells_obs
            c[i, factor] = average_weight
            s[i, factor] = fraction_attached

    ylabels = obs
    xlabels = model.factor_names[layer_idx]

    if cluster_rows:
        Z = ward(pdist(s))
        hclust_index = leaves_list(Z)
        s = s[hclust_index]
        c = c[hclust_index]
        ylabels = ylabels[hclust_index]

    if cluster_cols:
        Z = ward(pdist(s.T))
        hclust_index = leaves_list(Z)
        s = s[:, hclust_index]
        c = c[:, hclust_index]
        xlabels = xlabels[hclust_index]

    x, y = np.meshgrid(np.arange(len(xlabels)), np.arange(len(ylabels)))

    fig, axes = plt.subplots(1, 3, figsize=figsize, width_ratios=width_ratios)
    plt.sca(axes[0])
    ax = plt.gca()
    s = s / np.max(s)
    s *= s_max
    s += s_max / s_min
    if logged:
        c = np.log(c)
    plt.scatter(x, y, c=c, s=s, cmap=cmap)

    ax.set(
        xticks=np.arange(n_factors),
        yticks=np.arange(n_obs),
        xticklabels=xlabels,
        yticklabels=ylabels,
    )
    ax.tick_params(axis="both", which="major", labelsize=labelsize)
    ax.set_xticks(np.arange(n_factors + 1) - 0.5, minor=True)
    ax.set_yticks(np.arange(n_obs + 1) - 0.5, minor=True)
    ax.grid(which="minor")

    if show_ylabel:
        ax.set_ylabel(obs_key, rotation=270, labelpad=20.0, fontsize=labelsize)
    plt.xlabel("Factor", fontsize=labelsize)
    plt.title(f"Layer {layer_idx}\n", fontsize=titlesize)

    # Make legend
    map_f_to_s = {"0": 5, "25": 7, "50": 9, "75": 11, "100": 13}
    plt.sca(axes[1])
    ax = plt.gca()
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.spines["bottom"].set_visible(False)
    ax.spines["left"].set_visible(False)

    ax.get_xaxis().set_ticks([])
    ax.get_yaxis().set_ticks([])
    circles = [
        Line2D(
            [],
            [],
            color="white",
            marker="o",
            markersize=map_f_to_s[f],
            markerfacecolor="gray",
        )
        for f in map_f_to_s.keys()
    ]
    lg = plt.legend(  # noqa: F841
        circles[::-1],
        list(map_f_to_s.keys())[::-1],
        numpoints=1,
        loc=2,
        frameon=False,
        fontsize=legend_fontsize,
    )
    plt.title("Fraction of \ncells in group (%)", fontsize=legend_titlesize)

    plt.sca(axes[2])
    ax = plt.gca()
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.spines["bottom"].set_visible(False)
    ax.spines["left"].set_visible(False)

    ax.get_xaxis().set_ticks([])
    ax.get_yaxis().set_ticks([])
    cb = plt.colorbar(ax=axes[0], cmap=cmap)
    cb.ax.set_title("Average\ncell score", fontsize=legend_titlesize)
    cb.ax.tick_params(labelsize=legend_fontsize)
    plt.grid("off")
    if show:
        plt.show()


from .trajectory import multilevel_paga  # noqa: E402,F401


def layers_continuous_obs(
    model: "scDEF",
    obs_keys: Union[str, List[str]],
    obs_mats: Dict[str, Dict[int, np.ndarray]],
    sort_layer_factors: bool = True,
    orders: Optional[List[np.ndarray]] = None,
    layers: Optional[List[int]] = None,
    vmax: Optional[float] = None,
    vmin: Optional[float] = None,
    cb_title: str = "",
    cb_title_fontsize: int = 10,
    fontsize: int = 12,
    title_fontsize: int = 12,
    pad: float = 0.1,
    shrink: float = 0.7,
    figsize: Tuple[float, float] = (10, 4),
    xticks_rotation: float = 90.0,
    cmap: Optional[str] = None,
    show: bool = True,
    rasterized: bool = False,
    **kwargs: Any,
) -> None:
    """Plot continuous observation matrices across layers.

    Args:
        model: scDEF model instance
        obs_keys: observation keys to plot
        obs_mats: observation matrices dictionary
        sort_layer_factors: whether to sort factors by layer
        orders: custom factor orders
        layers: which layers to plot
        vmax: maximum value for colormap
        vmin: minimum value for colormap
        cb_title: colorbar title
        cb_title_fontsize: colorbar title font size
        fontsize: font size for labels
        title_fontsize: title font size
        pad: padding for colorbar
        shrink: shrink factor for colorbar
        figsize: figure size
        xticks_rotation: rotation angle for x-axis ticks
        cmap: colormap name
        show: whether to show the plot
        rasterized: whether to rasterize the plot
        **kwargs: additional plotting keyword arguments
    """
    if not isinstance(obs_keys, list):
        obs_keys = [obs_keys]

    if layers is None:
        layers = [i for i in range(0, model.n_layers) if len(model.factor_lists[i]) > 1]

    n_layers = len(layers)

    if sort_layer_factors:
        layer_factor_orders = model.get_layer_factor_orders()
    else:
        if orders is not None:
            layer_factor_orders = orders
        else:
            layer_factor_orders = [
                np.arange(len(model.factor_lists[i])) for i in range(model.n_layers)
            ]

    n_factors = [len(model.factor_lists[idx]) for idx in layers]
    n_obs = len(obs_keys)
    fig, axs = plt.subplots(
        n_obs,
        n_layers,
        figsize=figsize,
        gridspec_kw={"width_ratios": n_factors},
    )
    # Ensure axs is always 2D for later indexing
    if n_obs == 1 and n_layers == 1:
        axs = np.array([[axs]])
    elif n_obs == 1 or n_layers == 1:
        axs = np.atleast_2d(axs)
    else:
        axs = axs.reshape((n_obs, n_layers))
    for i in layers:
        axs[0][i].set_title(f"Layer {i}", fontsize=title_fontsize)
        for j, obs_key in enumerate(obs_keys):
            ax = axs[j][i]
            mat = obs_mats[obs_key][i]
            mat = mat[:, layer_factor_orders[i]]
            axplt = ax.pcolormesh(
                mat, vmax=vmax, vmin=vmin, cmap=cmap, rasterized=rasterized
            )

            if j == len(obs_keys) - 1:
                xlabels = model.factor_names[i]
                xlabels = np.array(xlabels)[layer_factor_orders[i]]
                ax.set_xticks(
                    np.arange(len(xlabels)) + 0.5,
                    xlabels,
                    rotation=xticks_rotation,
                    fontsize=fontsize,
                )
            else:
                ax.set(xticks=[])

            if i == 0:
                ax.set_yticks(
                    [0.5],
                    [obs_key],
                )
            else:
                ax.set(yticks=[])

    plt.subplots_adjust(wspace=0.05)
    plt.subplots_adjust(hspace=0.05)

    cb = fig.colorbar(axplt, ax=axs.ravel().tolist(), pad=pad, shrink=shrink)
    cb.ax.set_title(cb_title, fontsize=cb_title_fontsize)
    if show:
        plt.show()


def layers_obs(
    model: "scDEF",
    obs_keys: Union[str, List[str]],
    obs_mats: Dict[str, Dict[int, np.ndarray]],
    obs_clusters: Dict[str, np.ndarray],
    obs_vals_dict: Dict[str, List[str]],
    sort_layer_factors: bool = True,
    orders: Optional[List[np.ndarray]] = None,
    layers: Optional[List[int]] = None,
    vmax: Optional[float] = None,
    vmin: Optional[float] = None,
    cb_title: str = "",
    cb_title_fontsize: int = 10,
    fontsize: int = 12,
    title_fontsize: int = 12,
    pad: float = 0.1,
    shrink: float = 0.7,
    figsize: Tuple[float, float] = (10, 4),
    xticks_rotation: float = 90.0,
    cmap: Optional[str] = None,
    show: bool = True,
    rasterized: bool = False,
    **kwargs: Any,
) -> None:
    """Plot observation matrices across layers.

    Args:
        model: scDEF model instance
        obs_keys: observation keys to plot
        obs_mats: observation matrices dictionary
        sort_layer_factors: whether to sort factors by layer
        orders: custom factor orders
        layers: which layers to plot
        vmax: maximum value for colormap
        vmin: minimum value for colormap
        cb_title: colorbar title
        cb_title_fontsize: colorbar title font size
        fontsize: font size for labels
        title_fontsize: title font size
        pad: padding for colorbar
        shrink: shrink factor for colorbar
        figsize: figure size
        xticks_rotation: rotation angle for x-axis ticks
        cmap: colormap name
        show: whether to show the plot
        rasterized: whether to rasterize the plot
        **kwargs: additional plotting keyword arguments
    """
    if not isinstance(obs_keys, list):
        obs_keys = [obs_keys]

    if layers is None:
        layers = [i for i in range(0, model.n_layers) if len(model.factor_lists[i]) > 1]

    n_layers = len(layers)

    if sort_layer_factors:
        layer_factor_orders = model.get_layer_factor_orders()
    else:
        if orders is not None:
            layer_factor_orders = orders
        else:
            layer_factor_orders = [
                np.arange(len(model.factor_lists[i])) for i in range(model.n_layers)
            ]

    n_factors = [len(model.factor_lists[idx]) for idx in layers]
    n_obs = [len(obs_clusters[obs_key]) for obs_key in obs_keys]
    fig, axs = plt.subplots(
        len(obs_keys),
        n_layers,
        figsize=figsize,
        gridspec_kw={"width_ratios": n_factors, "height_ratios": n_obs},
    )
    axs = axs.reshape((len(obs_keys), n_layers))
    for i in layers:
        axs[0][i].set_title(f"Layer {i}", fontsize=title_fontsize)
        for j, obs_key in enumerate(obs_keys):
            ax = axs[j][i]
            mat = obs_mats[obs_key][i]
            mat = mat[obs_clusters[obs_key]][:, layer_factor_orders[i]]
            axplt = ax.pcolormesh(
                mat, vmax=vmax, vmin=vmin, cmap=cmap, rasterized=rasterized
            )

            if j == len(obs_keys) - 1:
                xlabels = model.factor_names[i]
                xlabels = np.array(xlabels)[layer_factor_orders[i]]
                ax.set_xticks(
                    np.arange(len(xlabels)) + 0.5,
                    xlabels,
                    rotation=xticks_rotation,
                    fontsize=fontsize,
                )
            else:
                ax.set(xticks=[])

            if i == 0:
                ylabels = np.array(obs_vals_dict[obs_keys[j]])[
                    obs_clusters[obs_keys[j]]
                ]
                ax.set_yticks(
                    np.arange(len(ylabels)) + 0.5,
                    ylabels,
                )
            else:
                ax.set(yticks=[])

            if i == n_layers - 1:
                ax.yaxis.set_label_position("right")
                ax.set_ylabel(obs_key, rotation=270, labelpad=20.0, fontsize=fontsize)

    plt.subplots_adjust(wspace=0.05)
    plt.subplots_adjust(hspace=0.05)

    cb = fig.colorbar(axplt, ax=axs.ravel().tolist(), pad=pad, shrink=shrink)
    cb.ax.set_title(cb_title, fontsize=cb_title_fontsize)
    if show:
        plt.show()


def pathway_scores(
    model: "scDEF",
    pathways: pd.DataFrame,
    top_genes: Optional[int] = 20,
    **kwargs: Any,
) -> None:
    """Plot the association between a set of cell annotations and a set of gene signatures.

    Args:
        model: scDEF model instance
        pathways: a pandas DataFrame containing PROGENy pathways
        top_genes: number of top genes to consider
        **kwargs: plotting keyword arguments
    """
    (
        obs_mats,
        obs_clusters,
        obs_vals_dict,
        joined_mats,
    ) = data_utils.prepare_pathway_factor_scores(
        model,
        pathways,
        top_genes=top_genes,
        **kwargs,
    )

    if "vmax" not in kwargs:
        kwargs["vmax"] = joined_mats.max()
    if "vmin" not in kwargs:
        kwargs["vmin"] = joined_mats.min()

    layers_obs(
        model,
        ["Pathway"],
        obs_mats,
        obs_clusters,
        obs_vals_dict,
        **kwargs,
    )


def signatures_scores(
    model: "scDEF",
    obs_keys: Sequence[str],
    markers: Mapping[str, Sequence[str]],
    top_genes: Optional[int] = 10,
    hierarchy: Optional[Dict[str, Sequence[str]]] = None,
    **kwargs: Any,
) -> None:
    """Plot the association between a set of cell annotations and a set of gene signatures.

    Args:
        model: scDEF model instance
        obs_keys: the keys in model.adata.obs to use
        markers: a dictionary with keys corresponding to model.adata.obs[obs_keys] and values to gene lists
        top_genes: number of genes to consider in the score computations
        hierarchy: the polytree to restrict the associations to
        **kwargs: plotting keyword arguments
    """
    obs_mats, obs_clusters, obs_vals_dict = data_utils.prepare_obs_factor_scores(
        model,
        obs_keys,
        data_utils.get_signature_scores,
        markers=markers,
        top_genes=top_genes,
        hierarchy=hierarchy,
    )
    layers_obs(model, obs_keys, obs_mats, obs_clusters, obs_vals_dict, **kwargs)


def obs_scores(
    model: "scDEF",
    obs_keys: Sequence[str],
    hierarchy: Optional[Dict[str, Sequence[str]]] = None,
    mode: Literal["f1", "fracs", "weights"] = "fracs",
    vmax: Optional[float] = None,
    vmin: Optional[float] = None,
    **kwargs: Any,
) -> None:
    """Plot the association between a set of cell annotations and factors.

    Args:
        model: scDEF model instance
        obs_keys: the keys in model.adata.obs to use
        hierarchy: the polytree to restrict the associations to
        mode: whether to compute scores based on assignments or weights
        **kwargs: plotting keyword arguments
    """
    if mode == "f1":
        f = data_utils.get_assignment_scores
    elif mode == "fracs":
        f = data_utils.get_assignment_fracs
    elif mode == "weights":
        f = data_utils.get_weight_scores
    else:
        raise ValueError("`mode` must be one of ['f1', 'fracs', 'weights']")

    obs_mats, obs_clusters, obs_vals_dict = data_utils.prepare_obs_factor_scores(
        model,
        obs_keys,
        f,
        hierarchy=hierarchy,
    )

    if mode == "f1" or mode == "fracs":
        vmax = 1.0
        vmin = 0.0

    layers_obs(
        model,
        obs_keys,
        obs_mats,
        obs_clusters,
        obs_vals_dict,
        vmax=vmax,
        vmin=vmin,
        **kwargs,
    )


def continuous_obs_scores(
    model: "scDEF",
    obs_keys: Sequence[str],
    mode: Literal["correlations"] = "correlations",
    vmax: Optional[float] = None,
    vmin: Optional[float] = None,
    **kwargs: Any,
) -> None:
    """Plot the correlations between a set of cell annotations and factors.

    Args:
        model: scDEF model instance
        obs_keys: the keys in model.adata.obs to use
        mode: how to compute scores
        **kwargs: plotting keyword arguments
    """
    if mode == "correlations":
        f = data_utils.get_correlations
    else:
        raise ValueError("`mode` must be one of ['correlations']")

    obs_mats = data_utils.prepare_continuous_factor_scores(
        model,
        obs_keys,
        f,
    )

    layers_continuous_obs(
        model,
        obs_keys,
        obs_mats,
        vmax=vmax,
        vmin=vmin,
        **kwargs,
    )


def umap(
    model: "scDEF",
    color: Union[str, List[str]] = [],
    layers: Optional[List[int]] = None,
    figsize: Tuple[float, float] = (16, 4),
    fontsize: int = 12,
    legend_fontsize: int = 10,
    rasterized: bool = True,
    n_legend_cols: int = 1,
    factor_subset: Optional[List[str]] = None,
    show: bool = True,
) -> None:
    """Plot pre-computed UMAPs for different layers.

    UMAP embeddings must have been computed first via ``scdef.tl.umap``.

    Args:
        model: scDEF model instance
        color: color key(s) to use for coloring
        layers: which layers to plot
        figsize: figure size
        fontsize: font size for labels
        legend_fontsize: legend font size
        rasterized: whether to rasterize the plot
        n_legend_cols: number of columns in legend
        factor_subset: subset of factors to plot
        show: whether to show the plot
    """
    if layers is None:
        layers = [
            i
            for i in range(model.n_layers - 1, -1, -1)
            if len(model.factor_lists[i]) > 1
        ]

    n_layers = len(layers)

    if not isinstance(color, list):
        color = [color]

    n_rows = len(color)
    if n_rows == 0:
        n_rows = 1

    fig, axes = plt.subplots(n_rows, n_layers, figsize=figsize)
    for layer in layers:
        layer_name = model.layer_names[layer]
        umap_key = f"X_umap_{layer_name}"
        if umap_key not in model.adata.obsm:
            raise KeyError(
                f"UMAP embedding '{umap_key}' not found in model.adata.obsm. "
                f"Run scdef.tl.umap(model) first."
            )
        # Temporarily set X_umap so sc.pl.umap can find it
        model.adata.obsm["X_umap"] = model.adata.obsm[umap_key]

        for row in range(len(color)):
            if n_rows > 1:
                ax = axes[row, layer]
            else:
                ax = axes[layer]
            legend_loc = None
            if layer == n_layers - 1:
                legend_loc = "right margin"
            ax = sc.pl.umap(
                model.adata,
                color=[color[row]],
                frameon=False,
                show=False,
                ax=ax,
                legend_loc=legend_loc,
            )
            if row == 0:
                ax.set_title(f"Layer {layer}", fontsize=fontsize)
            else:
                ax.set_title("")

            if layer == n_layers - 1:
                leg = ax.legend(
                    loc="center left",
                    bbox_to_anchor=(1, 0.5),
                    frameon=False,
                    title_fontsize=legend_fontsize,
                    fontsize=legend_fontsize,
                    title=color[row],
                    ncols=n_legend_cols,
                )
                leg._legend_box.align = "left"

    if show:
        plt.show()


def factors_bars(
    model: "scDEF",
    obs_keys: Union[str, List[str]],
    sort_layer_factors: bool = True,
    orders: Optional[List[np.ndarray]] = None,
    sharey: bool = True,
    layers: Optional[List[int]] = None,
    vmax: Optional[float] = None,
    vmin: Optional[float] = None,
    fontsize: int = 12,
    title_fontsize: int = 12,
    legend_fontsize: int = 8,
    figsize: Tuple[float, float] = (10, 4),
    total: bool = False,
    show: bool = True,
) -> None:
    """Plot factor scores as bar charts.

    Args:
        model: scDEF model instance
        obs_keys: observation keys to plot
        sort_layer_factors: whether to sort factors by layer
        orders: custom factor orders
        sharey: whether to share y-axis across subplots
        layers: which layers to plot
        vmax: maximum value for y-axis
        vmin: minimum value for y-axis
        fontsize: font size for labels
        title_fontsize: title font size
        legend_fontsize: legend font size
        figsize: figure size
        total: whether to plot total scores
        show: whether to show the plot
    """
    if not isinstance(obs_keys, list):
        obs_keys = [obs_keys]

    if layers is None:
        layers = [i for i in range(0, model.n_layers) if len(model.factor_lists[i]) > 1]

    n_layers = len(layers)

    if sort_layer_factors:
        layer_factor_orders = model.get_layer_factor_orders()
    else:
        if orders is not None:
            layer_factor_orders = orders
        else:
            layer_factor_orders = [
                np.arange(len(model.factor_lists[i])) for i in range(model.n_layers)
            ]

    n_factors = [len(model.factor_lists[idx]) for idx in layers]
    fig, axs = plt.subplots(
        len(obs_keys),
        n_layers,
        figsize=figsize,
        gridspec_kw={"width_ratios": n_factors},
        sharey=sharey,
    )
    # Ensure axs is always 2D for later indexing
    if len(obs_keys) == 1 and n_layers == 1:
        axs = np.array([[axs]])
    elif len(obs_keys) == 1 or n_layers == 1:
        axs = np.atleast_2d(axs)
    else:
        axs = axs.reshape((len(obs_keys), n_layers))
    # Get data using the new utility functions
    obs_mats, obs_clusters, obs_vals_dict = data_utils.prepare_obs_factor_scores(
        model,
        obs_keys,
        data_utils.get_assignment_fracs,
        total=total,
    )

    for i in layers:
        axs[0][i].set_title(f"Layer {i}", fontsize=title_fontsize)
        for j, obs_key in enumerate(obs_keys):
            ax = axs[j][i]
            mat = obs_mats[obs_key][i]
            mat = mat[:, layer_factor_orders[i]]

            # Plot bars for each observation
            for idx, obs in enumerate(obs_vals_dict[obs_key]):
                obs_idx = np.where(model.adata.obs[obs_key].cat.categories == obs)[0][0]
                color = model.adata.uns[f"{obs_key}_colors"][obs_idx]
                y = mat[idx]
                if idx == 0:
                    ax.bar(np.arange(len(y)), y, color=color, label=obs)
                else:
                    ax.bar(
                        np.arange(len(y)),
                        y,
                        bottom=mat[:idx].sum(axis=0),
                        color=color,
                        label=obs,
                    )

            ax.set_xticks(np.arange(len(model.factor_names[i])))
            ax.set_xticklabels(
                np.array(model.factor_names[i])[layer_factor_orders[i]],
                rotation=90,
                fontsize=fontsize,
            )
            if i == 0:
                ax.set_ylabel(obs_key, fontsize=fontsize)
            if vmax is not None:
                ax.set_ylim(0, vmax)
            if vmin is not None:
                ax.set_ylim(vmin, ax.get_ylim()[1])
            # Remove top and right axis frame
            ax.spines["top"].set_visible(False)
            ax.spines["right"].set_visible(False)

    plt.tight_layout()
    if show:
        plt.show()


def cell_entropies(
    model: "scDEF",
    thres: float = 0.9,
    show: bool = True,
) -> Optional[Figure]:
    """Plot cell entropies and factor numbers across layers.

    Args:
        model: scDEF model instance
        thres: Threshold for cumulative sum calculation
        show: Whether to show the plot
    """
    entropies = []
    factor_ns = []
    for layer_idx in range(4):
        layer_name = model.layer_names[layer_idx]
        nf = len(model.factor_lists[layer_idx])
        a = (
            model.adata.obsm[f"X_{layer_name}factors"]
            / np.sum(model.adata.obsm[f"X_{layer_name}factors"], axis=1)[:, None]
        )
        a_sorted = np.vstack(
            [a[i, np.argsort(a, axis=1)[:, ::-1][i]] for i in range(a.shape[0])]
        )
        a_cumsums = np.cumsum(a_sorted, axis=1)
        n_factors = np.array(
            [(np.where(a_cumsums[i] > thres)[0][0] + 1) for i in range(a.shape[0])]
        )
        factor_ns.append(n_factors)
        entropy = np.array(
            [np.sum(-np.log(a[i]) * a[i]) / np.log(nf) for i in range(a.shape[0])]
        )
        entropies.append(entropy)
    fig, axes = plt.subplots(1, 2)
    plt.sca(axes[0])
    plt.boxplot(entropies)
    plt.sca(axes[1])
    plt.boxplot(factor_ns)
    if show:
        plt.show()
    else:
        return fig


def factor_genes(
    model: "scDEF",
    thres: float = 0.9,
    show: bool = True,
) -> Optional[Figure]:
    """Plot number of genes in factors across layers.

    Args:
        model: scDEF model instance
        thres: threshold for cumulative sum calculation
        show: whether to show the plot

    Returns:
        Figure object if show is False, None otherwise
    """
    layer_kept_n_genes = []
    layer_removed_n_genes = []
    factor_n_kept_genes = dict()
    for layer_idx in range(len(model.layer_sizes)):
        layer_name = f"{model.layer_names[layer_idx]}"

        term_scores = model.pmeans[f"{model.layer_names[0]}W"]

        if layer_idx > 0:
            term_scores = model.pmeans[f"{model.layer_names[layer_idx]}W"]
            for layer in range(layer_idx - 1, 0, -1):
                lower_mat = model.pmeans[f"{model.layer_names[layer]}W"]
                term_scores = term_scores.dot(lower_mat)
            term_scores = term_scores.dot(model.pmeans[f"{model.layer_names[0]}W"])

        kept_factors_n_genes = []
        removed_factors_n_genes = []
        f_idx = 0
        for factor in range(model.layer_sizes[layer_idx]):
            vals = term_scores[factor] / np.sum(term_scores[factor])
            # sort
            vals_sorted = vals[np.argsort(vals)[::-1]]
            if factor in model.factor_lists[layer_idx]:
                kept_factors_n_genes.append(
                    np.where(np.cumsum(vals_sorted) > thres)[0][0]
                )
                factor_n_kept_genes[f"{layer_name}{f_idx}"] = kept_factors_n_genes[-1]
                f_idx += 1
            else:
                removed_factors_n_genes.append(
                    np.where(np.cumsum(vals_sorted) > thres)[0][0]
                )

        layer_kept_n_genes.append(kept_factors_n_genes)
        layer_removed_n_genes.append(removed_factors_n_genes)

    m = np.max(list(factor_n_kept_genes.values()))
    for f in factor_n_kept_genes:
        factor_n_kept_genes[f] = factor_n_kept_genes[f] / m

    fig, axes = plt.subplots(1, 4, figsize=(6, 3), sharey=True)
    for i in range(4):
        plt.sca(axes[i])
        plt.boxplot(
            [layer_kept_n_genes[i], layer_removed_n_genes[i]],
            tick_labels=["Kept", "Removed"],
        )
    plt.suptitle("Number of genes in factors")
    if show:
        plt.show()
    else:
        return fig


def factor_gini(
    model: "scDEF",
    idx: int,
    thres: float = 0.9,
    show: bool = True,
) -> Optional[Axes]:
    """Plot Gini coefficient for a specific factor.

    Args:
        model: scDEF model instance
        idx: Factor index to plot
        thres: Threshold for cumulative sum calculation
        show: Whether to show the plot
    """
    a = model.pmeans[f"{model.layer_names[0]}W"][idx]
    norm_a = a / np.sum(a)
    vals_sorted = norm_a[np.argsort(norm_a)[::-1]]
    plt.plot(np.cumsum(vals_sorted))
    plt.axvline(np.where(np.cumsum(vals_sorted) > thres)[0][0], color="gray")
    plt.title(model.pmeans["brd"][idx])
    if show:
        plt.show()
    else:
        return plt.gca()


def factor_gene_uncertainty_boxplot(
    model: "scDEF",
    factor: Union[int, str],
    layer_idx: int = 0,
    max_genes: Optional[int] = 50,
    mc_samples_upper: int = 100,
    random_seed: int = 0,
    whisker_quantiles: Tuple[float, float] = (0.05, 0.95),
    sort_by: Literal["mean", "confidence"] = "mean",
    color_by_confidence: bool = False,
    confidence_tau_quantile: float = 0.99,
    confidence_cmap: str = "viridis",
    confidence_vmin: float = 0.0,
    confidence_vmax: float = 1.0,
    add_confidence_colorbar: bool = True,
    show_confidence_cutoff_line: bool = False,
    confidence_include_threshold: float = 0.9,
    confidence_cutoff_line_kwargs: Optional[Dict[str, Any]] = None,
    show_tau_quantile_line: bool = False,
    tau_quantile_line_kwargs: Optional[Dict[str, Any]] = None,
    xtick_rotation: float = 90.0,
    figsize: Tuple[float, float] = (12, 4),
    ax: Optional[Axes] = None,
    show: bool = True,
) -> Optional[Axes]:
    """Plot per-gene uncertainty boxes for a factor using posterior mean/variance.

    Genes are sorted by posterior mean loading or confidence for the selected factor.
    For each gene, a box is drawn from an approximated posterior distribution of
    ``W`` using the stored posterior mean and variance:
      - box: 25th to 75th percentile
      - median: 50th percentile
      - whiskers: configurable quantiles (default 5th/95th)

    Args:
        model: scDEF model instance
        factor: factor index (within kept factors of the layer) or factor name
        layer_idx: layer index to visualize
        max_genes: maximum number of genes to plot; if None, plot all genes
        mc_samples_upper: number of posterior samples for ``layer_idx > 0``
        random_seed: random seed for upper-layer posterior sampling
        whisker_quantiles: lower/upper quantiles for whiskers
        sort_by: sorting criterion for genes; ``"mean"`` (default) sorts by
            posterior mean loading, ``"confidence"`` sorts by posterior
            confidence ``P(score > tau)``
        color_by_confidence: whether to color each box by posterior confidence
            ``P(W > tau)``
        confidence_tau_quantile: factor-wise quantile used to define ``tau`` for
            confidence coloring
        confidence_cmap: matplotlib colormap name for confidence coloring
        confidence_vmin: lower bound of confidence colormap normalization
        confidence_vmax: upper bound of confidence colormap normalization
        add_confidence_colorbar: whether to add a confidence colorbar when
            ``color_by_confidence`` is True
        show_confidence_cutoff_line: whether to draw a vertical line marking the
            last plotted gene with confidence >= ``confidence_include_threshold``
        confidence_include_threshold: confidence threshold used for the cutoff
            line (default 0.9)
        confidence_cutoff_line_kwargs: optional kwargs passed to ``ax.axvline``
            for styling the cutoff line
        show_tau_quantile_line: whether to draw a horizontal line at ``tau``,
            where ``tau`` is the quantile threshold set by
            ``confidence_tau_quantile``
        tau_quantile_line_kwargs: optional kwargs passed to ``ax.axhline`` for
            styling the tau quantile line
        xtick_rotation: x tick label rotation in degrees
        figsize: figure size if ax is None
        ax: matplotlib axis to draw on
        show: whether to show the plot

    Returns:
        Axes object if show is False, None otherwise.
    """
    if layer_idx < 0 or layer_idx >= model.n_layers:
        raise ValueError(f"layer_idx must be in [0, {model.n_layers - 1}].")
    if mc_samples_upper <= 0:
        raise ValueError("mc_samples_upper must be > 0.")
    q_lo, q_hi = whisker_quantiles
    if not (0.0 < q_lo < 0.5 and 0.5 < q_hi < 1.0):
        raise ValueError("whisker_quantiles must satisfy 0 < q_lo < 0.5 < q_hi < 1.")
    if not (0.0 < confidence_tau_quantile < 1.0):
        raise ValueError("confidence_tau_quantile must be in (0, 1).")
    if sort_by not in {"mean", "confidence"}:
        raise ValueError("sort_by must be one of {'mean', 'confidence'}.")
    if not (0.0 < confidence_include_threshold < 1.0):
        raise ValueError("confidence_include_threshold must be in (0, 1).")
    if show_confidence_cutoff_line and sort_by != "confidence":
        raise ValueError("show_confidence_cutoff_line requires sort_by='confidence'.")

    if isinstance(factor, str):
        if factor not in model.factor_names[layer_idx]:
            raise ValueError(
                f"Factor '{factor}' not found in layer {layer_idx} factor names."
            )
        factor_idx = model.factor_names[layer_idx].index(factor)
    else:
        factor_idx = int(factor)
    if factor_idx < 0 or factor_idx >= len(model.factor_lists[layer_idx]):
        raise IndexError("factor index out of bounds for kept factors in this layer.")

    abs_factor_idx = int(model.factor_lists[layer_idx][factor_idx])
    layer_name = model.layer_names[layer_idx]
    genes = np.asarray(model.adata.var_names)
    eps = 1e-12

    if layer_idx == 0:
        all_means = np.asarray(model.pmeans[f"{layer_name}W"], dtype=float)[
            abs_factor_idx
        ]
        all_variances = np.asarray(model.pvars[f"{layer_name}W"], dtype=float)[
            abs_factor_idx
        ]
        all_variances = np.maximum(all_variances, 0.0)
        tau = float(np.quantile(all_means, confidence_tau_quantile))
        all_confidences = norm.cdf((all_means - tau) / np.sqrt(all_variances + eps))
        if sort_by == "confidence":
            order = np.lexsort((-all_means, -all_confidences))
        else:
            order = np.argsort(all_means)[::-1]
        if max_genes is not None:
            order = order[: int(max_genes)]
        means = all_means[order]
        variances = all_variances[order]
        genes = genes[order]
        confidences = all_confidences[order]

        mean_safe = np.maximum(means, eps)
        sigma2 = np.log1p(variances / np.maximum(mean_safe**2, eps))
        sigma = np.sqrt(np.maximum(sigma2, 0.0))
        mu_log = np.log(mean_safe) - 0.5 * sigma2

        def _q(p):
            return np.exp(mu_log + sigma * norm.ppf(p))

        q1 = _q(0.25)
        med = _q(0.50)
        q3 = _q(0.75)
        whislo = _q(q_lo)
        whishi = _q(q_hi)
    else:
        from jax import random

        if hasattr(model, "logger"):
            model.logger.info(
                "Computing upper-layer uncertainty boxplot with Monte Carlo "
                "(layer=%s, samples=%s).",
                layer_idx,
                mc_samples_upper,
            )

        _, mean_scores = model.get_rankings(
            layer_idx=layer_idx,
            top_genes=len(genes),
            genes=True,
            return_scores=True,
            sorted_scores=False,
        )
        all_means = np.asarray(mean_scores[factor_idx], dtype=float)

        base_rng = random.PRNGKey(int(random_seed))
        sample_scores_all = np.empty((int(mc_samples_upper), len(genes)), dtype=float)
        for s_idx in range(int(mc_samples_upper)):
            rng = random.fold_in(base_rng, factor_idx * int(mc_samples_upper) + s_idx)
            _, sampled = model.get_signature_sample(
                rng,
                factor_idx=factor_idx,
                layer_idx=layer_idx,
                top_genes=len(model.adata.var_names),
                return_scores=True,
            )
            sample_scores_all[s_idx, :] = np.asarray(sampled, dtype=float)

        tau = float(np.quantile(all_means, confidence_tau_quantile))
        all_confidences = np.mean(sample_scores_all > tau, axis=0)
        if sort_by == "confidence":
            order = np.lexsort((-all_means, -all_confidences))
        else:
            order = np.argsort(all_means)[::-1]
        if max_genes is not None:
            order = order[: int(max_genes)]

        means = all_means[order]
        genes = genes[order]
        confidences = all_confidences[order]
        sample_scores = sample_scores_all[:, order]

        q1 = np.quantile(sample_scores, 0.25, axis=0)
        med = np.quantile(sample_scores, 0.50, axis=0)
        q3 = np.quantile(sample_scores, 0.75, axis=0)
        whislo = np.quantile(sample_scores, q_lo, axis=0)
        whishi = np.quantile(sample_scores, q_hi, axis=0)

    box_stats = [
        {
            "label": str(gene),
            "med": float(med[i]),
            "q1": float(q1[i]),
            "q3": float(q3[i]),
            "whislo": float(whislo[i]),
            "whishi": float(whishi[i]),
            "fliers": [],
        }
        for i, gene in enumerate(genes)
    ]

    if ax is None:
        fig_w = max(figsize[0], 0.2 * len(box_stats))
        _, ax = plt.subplots(1, 1, figsize=(fig_w, figsize[1]))
    artists = ax.bxp(box_stats, showfliers=False, patch_artist=color_by_confidence)

    if color_by_confidence:
        cmap = plt.get_cmap(confidence_cmap)
        conf_norm = mcolors.Normalize(
            vmin=confidence_vmin, vmax=confidence_vmax, clip=True
        )
        for i, box in enumerate(artists["boxes"]):
            box.set_facecolor(cmap(conf_norm(confidences[i])))
            box.set_alpha(0.85)
        if add_confidence_colorbar:
            sm = plt.cm.ScalarMappable(norm=conf_norm, cmap=cmap)
            sm.set_array([])
            cbar = ax.figure.colorbar(sm, ax=ax, pad=0.01)
            cbar.set_label(f"P(W > tau), tau=q{confidence_tau_quantile:.2f}")

    if show_confidence_cutoff_line:
        included = np.where(confidences >= confidence_include_threshold)[0]
        if len(included) > 0:
            line_kwargs: Dict[str, Any] = {
                "color": "black",
                "linestyle": "--",
                "linewidth": 1.0,
                "alpha": 0.9,
            }
            if confidence_cutoff_line_kwargs is not None:
                line_kwargs.update(confidence_cutoff_line_kwargs)
            # Draw after the last included gene position.
            ax.axvline(float(included[-1]) + 0.5, **line_kwargs)

    if show_tau_quantile_line:
        q_label = np.format_float_positional(
            confidence_tau_quantile, precision=6, fractional=True, trim="-"
        )
        tau_line_kwargs_final: Dict[str, Any] = {
            "color": "gray",
            "linestyle": ":",
            "linewidth": 1.0,
            "alpha": 0.9,
            "label": f"tau (q={q_label})",
        }
        if tau_quantile_line_kwargs is not None:
            tau_line_kwargs_final.update(tau_quantile_line_kwargs)
        ax.axhline(float(tau), **tau_line_kwargs_final)
        ax.legend(frameon=False, fontsize=8, loc="upper right")

    factor_name = model.factor_names[layer_idx][factor_idx]
    ax.set_title(f"{factor_name}: posterior loading uncertainty")
    ax.set_ylabel("W loading")
    ax.tick_params(axis="x", labelrotation=xtick_rotation)
    ax.grid(axis="y", alpha=0.2)
    if show:
        plt.show()
        return None
    return ax
