utils.py

import os
import re
import scanpy as sc
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import matplotlib.patheffects as pe
import seaborn as sns
import pandas as pd
from sklearn.metrics import adjusted_rand_score, normalized_mutual_info_score, silhouette_score
from sklearn.manifold import TSNE
import anndata as ad
from typing import Union, List, Tuple, Dict, Any, Optional
import numpy as np
import torch
import networkx as nx
from scipy.stats import ranksums


def create_segmented_colormap(n_labels: int, labels_per_segment: int = 5):
    
    # Base color families (pick distinct hues)
    color_families = [
        'Blues',      # 0-4
        'Reds',       # 5-9
        'Greens',     # 10-14
        'Purples',    # 15-19
        'Oranges',    # 20-24
        'YlOrBr',     # 25-29
        'PuBu',       # 30-34
        'RdPu',       # 35-39
        'YlGn',       # 40-44
        'OrRd',       # 45-49
        'BuGn',       # 50-54
        'PuRd',       # 55-59
    ]

    colors = []
    n_segments = int(np.ceil(n_labels / labels_per_segment))

    for seg_idx in range(n_segments):
        # Get color family for this segment
        family = color_families[seg_idx % len(color_families)]
        cmap = plt.cm.get_cmap(family)

        # How many labels in this segment?
        start_label = seg_idx * labels_per_segment
        end_label = min((seg_idx + 1) * labels_per_segment, n_labels)
        n_in_segment = end_label - start_label

        # Create gradient within this family
        # Start from 0.3 (not too light) to 1.0 (full intensity)
        gradient_positions = np.linspace(0.3, 1.0, n_in_segment)

        for pos in gradient_positions:
            colors.append(cmap(pos))

    return mcolors.ListedColormap(colors[:n_labels])


def _sanitize_filename_component(value: str) -> str:
    cleaned = re.sub(r"[^A-Za-z0-9_-]+", "_", str(value)).strip("_")
    return cleaned or "label"


def _ordered_label_categories(labels) -> Tuple[np.ndarray, np.ndarray]:
    series = pd.Series(labels)
    if pd.api.types.is_categorical_dtype(series):
        categories = series.cat.categories.astype(str).to_numpy()
    else:
        categories = np.sort(series.astype(str).unique())
    values = series.astype(str).to_numpy()
    return values, categories


def _plot_embedding_clean(
    coords: np.ndarray,
    labels,
    title: str,
    save_path: str,
    segmented: bool = False,
    group_num: int = 5,
    annotate: bool = False,
    point_size: float = 5.0,
    alpha: float = 0.85,
    figsize: Tuple[float, float] = (8.5, 8.0),
):
    label_values, categories = _ordered_label_categories(labels)
    n_labels = len(categories)

    if segmented:
        cmap = create_segmented_colormap(max(n_labels, 1), labels_per_segment=group_num)
        palette = [cmap(i) for i in range(max(n_labels, 1))]
    else:
        palette = sns.color_palette("husl", max(n_labels, 1))

    color_map = {label: palette[i] for i, label in enumerate(categories)}
    point_colors = [color_map[label] for label in label_values]

    def _legend_handles():
        return [
            plt.Line2D(
                [0],
                [0],
                marker="o",
                color="w",
                label=str(label),
                markerfacecolor=color_map[label],
                markersize=6,
            )
            for label in categories
        ]

    fig, ax = plt.subplots(figsize=figsize)
    ax.scatter(
        coords[:, 0],
        coords[:, 1],
        c=point_colors,
        s=point_size,
        alpha=alpha,
        linewidths=0,
        rasterized=True,
    )
    ax.set_xlabel("")
    ax.set_ylabel("")
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_title(title, fontsize=13, fontweight="bold")

    if annotate and 1 < n_labels <= 30:
        for label in categories:
            mask = label_values == label
            if not np.any(mask):
                continue
            x_med = float(np.median(coords[mask, 0]))
            y_med = float(np.median(coords[mask, 1]))
            txt = ax.text(
                x_med,
                y_med,
                str(label),
                fontsize=8 if n_labels > 16 else 9,
                ha="center",
                va="center",
                color="black",
                zorder=3,
            )
            txt.set_path_effects([pe.withStroke(linewidth=3, foreground="white")])
    elif n_labels <= 18:
        handles = _legend_handles()
        ax.legend(
            handles=handles,
            bbox_to_anchor=(1.02, 1.0),
            loc="upper left",
            frameon=False,
            fontsize=8,
        )

    fig.tight_layout()
    fig.savefig(save_path, dpi=300, bbox_inches="tight")
    plt.close(fig)

    if n_labels > 18:
        legend_fig_height = max(4.0, min(18.0, 0.28 * n_labels + 1.5))
        legend_fig, legend_ax = plt.subplots(figsize=(8.5, legend_fig_height))
        legend_ax.axis("off")
        legend_ax.legend(
            handles=_legend_handles(),
            loc="upper left",
            frameon=False,
            fontsize=8,
            ncol=2 if n_labels <= 40 else 3,
            columnspacing=1.2,
            handletextpad=0.6,
        )
        legend_fig.tight_layout()
        legend_path = save_path.replace(".png", "_legend.png")
        legend_fig.savefig(legend_path, dpi=300, bbox_inches="tight")
        plt.close(legend_fig)


def plot_umap(
        adata: ad.AnnData,
        embedding_key: str = 'X_umap',
        label_key: str = 'cell_type',
        save_path: Optional[str] = None,
        title: Optional[str] = None,
        figsize: Tuple[int, int] = (8, 6),
        dpi: int = 300,
        segment: bool = False,
        group_num: int = 5,
        **kwargs
):
    
    # Compute UMAP if not already present
    if embedding_key not in adata.obsm:
        print(f"Computing UMAP (embedding not found in adata.obsm['{embedding_key}'])...")
        sc.pp.neighbors(adata, use_rep='X')
        sc.tl.umap(adata)

    # Prepare colormap if segmented
    palette = None
    plot_label_key = label_key

    if segment:
        if pd.api.types.is_categorical_dtype(adata.obs[label_key]):
            n_labels = len(adata.obs[label_key].cat.categories)

            # Create numeric version
            numeric_key = f"{label_key}_numeric"
            adata.obs[numeric_key] = adata.obs[label_key].cat.codes
            plot_label_key = numeric_key

            # Create segmented colormap
            palette = create_segmented_colormap(n_labels, labels_per_segment=group_num)
            print(f"Using segmented colormap: {n_labels} labels, {group_num} per color family")
        else:
            print(f"Warning: {label_key} is not categorical. Segment mode requires categorical labels.")
            print("Using default colors instead.")

    # Plot
    fig = plt.figure(figsize=figsize)

    if palette is not None:
        sc.pl.umap(
            adata,
            color=plot_label_key,
            show=False,
            palette=palette,
            **kwargs
        )
    else:
        sc.pl.umap(
            adata,
            color=label_key,
            show=False,
            **kwargs
        )

    if title:
        plt.title(title, fontsize=14, fontweight='bold')

    if save_path:
        plt.savefig(save_path, dpi=dpi, bbox_inches='tight')
        print(f"Saved UMAP to {save_path}")

    return fig


def plot_umap_tsne(
        adata,
        save_dir: str,
        best_params: Dict[str, Any],
        label_key: str = "cell_type",
        split: str = "train",
        epoch: int = 0,
        min_dist: float = 0.3,
        spread: float = 1.0,
        dpi: int = 300,
        return_fig: bool = False,
        segment: bool = False,
        group_num: int = 5,
):
    
    # Create umap subdirectory under save_dir
    umap_dir = os.path.join(save_dir, 'umap')
    os.makedirs(umap_dir, exist_ok=True)

    # Compute neighbors if not already done
    _get_knn_indices(adata, use_rep="cell_embed", n_neighbors=best_params['n_neighbors'],
                     random_state=42, calc_knn=True)

    # Clustering
    if best_params['method'] == 'leiden':
        sc.tl.leiden(adata, resolution=best_params['resolution'], key_added='best_clustering')
    elif best_params['method'] == 'louvain':
        sc.tl.louvain(adata, resolution=best_params['resolution'], key_added='best_clustering')

    ari = best_params['ari']
    nmi = best_params['nmi']
    asw = best_params.get('asw', 0.0)

    # Compute UMAP and TSNE
    sc.tl.umap(adata, min_dist=min_dist, spread=spread)
    sc.tl.tsne(adata, use_rep="cell_embed")

    # Prepare base color variables
    color_by_categorical = [label_key]
    if 'best_clustering' in adata.obs.columns:
        color_by_categorical = ['best_clustering', label_key]

    # ========== UMAP PLOT (with optional segmented colors) ==========
    palette_umap = None
    color_by_umap = color_by_categorical

    if segment:
        # Create segmented continuous colormap for UMAP only
        if pd.api.types.is_categorical_dtype(adata.obs[label_key]):
            n_labels = len(adata.obs[label_key].cat.categories)
            # Create numeric version of labels
            numeric_key = f"{label_key}_numeric"
            adata.obs[numeric_key] = adata.obs[label_key].cat.codes

            # Update color_by for UMAP
            if 'best_clustering' in adata.obs.columns:
                color_by_umap = ['best_clustering', numeric_key]
            else:
                color_by_umap = [numeric_key]

            # Create segmented colormap
            palette_umap = create_segmented_colormap(n_labels, labels_per_segment=group_num)
        else:
            print(f"Warning: {label_key} is not categorical, segment mode requires categorical labels. Using default colors.")

    # Generate filenames
    ari_batch = best_params.get('ari_batch', 0)
    segment_str = f"_seg{group_num}" if segment else ""
    if ari_batch > 0:
        fname_base = f"_{split}_epoch{epoch:02d}_ari{ari:.3f}_arib{ari_batch:.3f}_nmi{nmi:.3f}_asw{asw:.3f}_{best_params['method']}_res{best_params['resolution']}_nn{best_params['n_neighbors']}{segment_str}"
    else:
        fname_base = f"_{split}_epoch{epoch:02d}_ari{ari:.3f}_nmi{nmi:.3f}_asw{asw:.3f}_{best_params['method']}_res{best_params['resolution']}_nn{best_params['n_neighbors']}{segment_str}"

    fname_umap = fname_base + "_umap.png"
    fname_tsne = fname_base.replace(segment_str, "") + "_tsne.png"  # TSNE never gets segment suffix

    # Plot UMAP (with segmented colors if requested)
    if palette_umap is not None:
        fig_umap = sc.pl.umap(
            adata,
            color=color_by_umap,
            show=False,
            return_fig=True,
            size=40,
            palette=palette_umap,
        )
    else:
        fig_umap = sc.pl.umap(
            adata,
            color=color_by_umap,
            show=False,
            return_fig=True,
            size=40,
        )

    # Add title for UMAP
    title = f"{split.upper()} - Epoch {epoch} | ARI: {ari:.4f}, NMI: {nmi:.4f}, ASW: {asw:.4f}"
    fig_umap.suptitle(title, fontsize=14, fontweight='bold', y=1.02)

    # Save UMAP
    fig_path_umap = os.path.join(umap_dir, fname_umap)
    fig_umap.savefig(fig_path_umap, dpi=dpi, bbox_inches='tight')

    # ========== TSNE PLOT (always categorical colors) ==========
    fig_tsne = sc.pl.tsne(
        adata,
        color=color_by_categorical,  # Always use categorical colors
        show=False,
        return_fig=True,
        size=40,
    )

    # Add title for TSNE
    fig_tsne.suptitle(title, fontsize=14, fontweight='bold', y=1.02)

    # Save TSNE
    fig_path_tsne = os.path.join(umap_dir, fname_tsne)
    fig_tsne.savefig(fig_path_tsne, dpi=dpi, bbox_inches='tight')

    # Save cleaner single-panel versions without the oversized Scanpy legends.
    label_key_safe = _sanitize_filename_component(label_key)
    clean_title_prefix = f"{split.upper()} - Epoch {epoch}"
    if "best_clustering" in adata.obs.columns:
        _plot_embedding_clean(
            adata.obsm["X_umap"],
            adata.obs["best_clustering"],
            title=f"{clean_title_prefix} | best_clustering",
            save_path=os.path.join(umap_dir, fname_umap.replace(".png", "_clean_best_clustering.png")),
            segmented=True,
            group_num=group_num,
            annotate=False,
            point_size=4.0,
            alpha=0.9,
        )
        _plot_embedding_clean(
            adata.obsm["X_tsne"],
            adata.obs["best_clustering"],
            title=f"{clean_title_prefix} | best_clustering",
            save_path=os.path.join(umap_dir, fname_tsne.replace(".png", "_clean_best_clustering.png")),
            segmented=True,
            group_num=group_num,
            annotate=False,
            point_size=4.0,
            alpha=0.9,
        )

    _plot_embedding_clean(
        adata.obsm["X_umap"],
        adata.obs[label_key],
        title=f"{clean_title_prefix} | {label_key}",
        save_path=os.path.join(umap_dir, fname_umap.replace(".png", f"_clean_{label_key_safe}.png")),
        segmented=False,
        annotate=True,
        point_size=4.0,
        alpha=0.9,
    )
    _plot_embedding_clean(
        adata.obsm["X_tsne"],
        adata.obs[label_key],
        title=f"{clean_title_prefix} | {label_key}",
        save_path=os.path.join(umap_dir, fname_tsne.replace(".png", f"_clean_{label_key_safe}.png")),
        segmented=False,
        annotate=True,
        point_size=4.0,
        alpha=0.9,
    )

    if return_fig:
        return fig_umap, fig_tsne
    else:
        fig_umap.clf()
        plt.close(fig_umap)
        fig_tsne.clf()
        plt.close(fig_tsne)

def evaluate_clustering(adata: ad.AnnData, key: str, batch_key: str = None) -> Dict[str, Any]:
    
    clustering_methods = ["leiden", "louvain"]
    resolutions = [0.32, 0.48, 0.64, 0.80]  # More granular resolution search
    n_neighbors_list = [15, 30]  # Reduced from 3 to 2 (30 is often optimal)

    best_params = {
        'resolution': 0,
        'ari': -1,
        'ari_batch': 1.0,  # NEW: batch ARI metric (lower is better)
        'nmi': 0,
        'asw': -1,  # Average Silhouette Width (higher is better)
        'method': None,
        'n_neighbors': 0
    }

    for n_neighbor in n_neighbors_list:
        # Compute KNN graph with current n_neighbors
        _get_knn_indices(adata, use_rep="cell_embed", n_neighbors=n_neighbor,
                        random_state=42, calc_knn=True)

        for method in clustering_methods:
            clustering_func = sc.tl.leiden if method == 'leiden' else sc.tl.louvain

            for resolution in resolutions:
                clustering_func(adata, resolution=resolution, key_added=method)

                # ARI with ground truth (KEEP FULL PRECISION - no rounding)
                ari = adjusted_rand_score(adata.obs[key], adata.obs[method])
                nmi = normalized_mutual_info_score(adata.obs[key], adata.obs[method])

                # Compute ASW (Average Silhouette Width) using cell embeddings
                # ASW measures cluster separation quality (higher is better, range: -1 to 1)
                asw = 0.0
                try:
                    # Need at least 2 clusters to compute silhouette
                    n_clusters = len(adata.obs[method].unique())
                    if n_clusters > 1 and n_clusters < len(adata):
                        asw = silhouette_score(
                            adata.obsm['cell_embed'],
                            adata.obs[method],
                            metric='euclidean'
                        )
                except:
                    asw = 0.0

                # NEW: Batch ARI if batch info available
                # Lower batch ARI = better (clustering transcends batch boundaries)
                ari_batch = 1.0
                if batch_key and batch_key in adata.obs:
                    ari_batch = adjusted_rand_score(adata.obs[batch_key], adata.obs[method])

                # Update best params: prioritize high ARI, use low batch ARI as tie-breaker
                should_update = False
                if ari > best_params['ari']:
                    should_update = True
                elif ari == best_params['ari'] and ari_batch < best_params['ari_batch']:
                    # Tie-breaking: prefer clustering with lower batch ARI
                    should_update = True

                if should_update:
                    best_params.update({
                        'resolution': resolution,
                        'ari': ari,
                        'ari_batch': ari_batch,
                        'method': method,
                        'n_neighbors': n_neighbor
                    })

                if nmi > best_params['nmi']:
                    best_params['nmi'] = nmi

                if asw > best_params['asw']:
                    best_params['asw'] = asw

    return best_params

def _get_knn_indices(adata: ad.AnnData,
                     use_rep: str = "delta",
                     n_neighbors: int = 25,
                     random_state: int = 0,
                     calc_knn: bool = True
                     ) -> np.ndarray:
    if calc_knn:
        assert use_rep == 'X' or use_rep in adata.obsm, f'{use_rep} not in adata.obsm and is not "X"'
        neighbors = sc.Neighbors(adata)
        neighbors.compute_neighbors(n_neighbors=n_neighbors, knn=True, use_rep=use_rep, random_state=random_state)
        adata.obsp['distances'] = neighbors.distances
        adata.obsp['connectivities'] = neighbors.connectivities

        # Get knn_indices - compute from distances if not available
        if hasattr(neighbors, 'knn_indices') and neighbors.knn_indices is not None:
            adata.obsm['knn_indices'] = neighbors.knn_indices
        else:
            # Compute knn_indices from distances matrix
            from scipy.sparse import issparse
            if issparse(neighbors.distances):
                distances_array = neighbors.distances.toarray()
            else:
                distances_array = neighbors.distances
            # Get indices of k nearest neighbors (sorted by distance)
            adata.obsm['knn_indices'] = np.argsort(distances_array, axis=1)[:, :n_neighbors]
        adata.uns['neighbors'] = {
            'connectivities_key': 'connectivities',
            'distances_key': 'distances',
            'knn_indices_key': 'knn_indices',
            'params': {
                'n_neighbors': n_neighbors,
                'use_rep': use_rep,
                'metric': 'euclidean',
                'method': 'umap'
            }
        }
    else:
        assert 'neighbors' in adata.uns, 'No precomputed knn exists.'
        assert adata.uns['neighbors']['params'][
                   'n_neighbors'] >= n_neighbors, f"pre-computed n_neighbors is {adata.uns['neighbors']['params']['n_neighbors']}, which is smaller than {n_neighbors}"

    return adata.obsm['knn_indices']


def plot_theta_by_celltype_heatmap(
    theta: np.ndarray,
    cell_type_labels: np.ndarray,
    save_path: str,
    figsize: Tuple[int, int] = (14, 10),
    cmap: Optional[str] = "Blues",
    title: Optional[str] = None,
    dpi: int = 300,
    min_topic_proportion: float = 0.05,
    vmin: float = 0.0,
    vmax: float = 0.5,
    topics_to_plot: Optional[List[int]] = None,
    keep_topics_order: bool = False,
    cell_barcodes: Optional[np.ndarray] = None,
    highlight_multi_topic_cells: bool = False,
    multi_topic_threshold: float = 0.05,
    barcode_membership_topics: Optional[List[int]] = None,
    barcode_membership_threshold: float = 0.05,
    cell_sort_scores: Optional[np.ndarray] = None,
    preserve_cell_order: bool = False,
):
    
    # Convert to pandas DataFrame for easier manipulation
    df = pd.DataFrame(theta, columns=[f'Topic_{i}' for i in range(theta.shape[1])])
    df['cell_type'] = cell_type_labels
    if cell_barcodes is not None:
        if len(cell_barcodes) != theta.shape[0]:
            raise ValueError(
                f"cell_barcodes length ({len(cell_barcodes)}) must match theta rows ({theta.shape[0]})."
            )
        df['cell_barcode'] = np.asarray(cell_barcodes).astype(str)
    if cell_sort_scores is not None:
        if len(cell_sort_scores) != theta.shape[0]:
            raise ValueError(
                f"cell_sort_scores length ({len(cell_sort_scores)}) must match theta rows ({theta.shape[0]})."
            )
        df['cell_sort_score'] = np.asarray(cell_sort_scores, dtype=float)

    # Calculate average topic proportions per cell type (topic columns only).
    topic_cols_all = [f'Topic_{i}' for i in range(theta.shape[1])]
    topic_avg_by_celltype = df.groupby('cell_type')[topic_cols_all].mean()  # [num_cell_types, num_topics]
    # Use the exact grouped index order to keep label mapping consistent.
    unique_cell_types = topic_avg_by_celltype.index.tolist()

    if topics_to_plot is not None:
        max_topic_idx = theta.shape[1] - 1
        seen = set()
        topics_valid = []
        for t in topics_to_plot:
            ti = int(t)
            if ti in seen:
                continue
            seen.add(ti)
            if 0 <= ti <= max_topic_idx:
                topics_valid.append(ti)
            else:
                print(f"Warning: topic {ti} is out of range [0, {max_topic_idx}] and will be ignored.")
        if len(topics_valid) == 0:
            print("Warning: no valid topics left after --topics_to_plot filtering; skipping heatmap.")
            return

        topic_info = []
        for topic_idx in topics_valid:
            topic_col = f'Topic_{topic_idx}'
            avg_proportions = topic_avg_by_celltype[topic_col].values
            max_celltype_idx = int(np.argmax(avg_proportions))
            max_proportion = float(avg_proportions[max_celltype_idx])
            topic_info.append({
                'original_idx': topic_idx,
                'best_celltype': unique_cell_types[max_celltype_idx],
                'best_celltype_idx': max_celltype_idx,
                'max_proportion': max_proportion
            })

        if keep_topics_order:
            topic_order = list(topics_valid)
            info_by_topic = {x['original_idx']: x for x in topic_info}
            topic_info_sorted = [info_by_topic[t] for t in topic_order]
        else:
            topic_info_sorted = sorted(topic_info, key=lambda x: (x['best_celltype_idx'], -x['max_proportion']))
            topic_order = [t['original_idx'] for t in topic_info_sorted]
    else:
        # Step 1: Ensure each cell type has at least one topic (its best topic)
        guaranteed_topics = set()
        for ct_idx, ct_name in enumerate(unique_cell_types):
            # Find the topic with highest average proportion for this cell type
            ct_proportions = topic_avg_by_celltype.loc[ct_name].values
            best_topic_idx = np.argmax(ct_proportions)
            guaranteed_topics.add(best_topic_idx)

        # Step 2: Add topics that are significant in at least one cell type
        topic_info = []
        for topic_idx in range(theta.shape[1]):
            topic_col = f'Topic_{topic_idx}'
            avg_proportions = topic_avg_by_celltype[topic_col].values
            max_celltype_idx = np.argmax(avg_proportions)
            max_proportion = avg_proportions[max_celltype_idx]

            # Keep topic if: (1) it's guaranteed for a cell type OR (2) it's above threshold
            if topic_idx in guaranteed_topics or max_proportion >= min_topic_proportion:
                topic_info.append({
                    'original_idx': topic_idx,
                    'best_celltype': unique_cell_types[max_celltype_idx],
                    'best_celltype_idx': max_celltype_idx,
                    'max_proportion': max_proportion
                })

        # Sort topics by: (1) their best cell type, (2) max proportion (descending)
        topic_info_sorted = sorted(topic_info, key=lambda x: (x['best_celltype_idx'], -x['max_proportion']))

        # Get the reordered topic indices
        topic_order = [t['original_idx'] for t in topic_info_sorted]

    print(f"Filtered topics: {len(topic_order)} out of {theta.shape[1]} topics kept")
    print(f"Dropped {theta.shape[1] - len(topic_order)} insignificant topics")

    if len(topic_order) == 0:
        print("Warning: No topics passed significance threshold!")
        return

    # Sort rows (default: by cell type, optional: preserve incoming order).
    if preserve_cell_order:
        df_sorted = df.copy()
    else:
        df_sorted = df.copy()
        df_sorted['cell_type'] = pd.Categorical(
            df_sorted['cell_type'],
            categories=unique_cell_types,
            ordered=True
        )
        sort_cols = ['cell_type']
        sort_ascending = [True]
        if 'cell_sort_score' in df_sorted.columns:
            sort_cols.append('cell_sort_score')
            sort_ascending.append(False)
        df_sorted = df_sorted.sort_values(sort_cols, ascending=sort_ascending, kind='mergesort')

    # Select and reorder topics
    topic_cols = [f'Topic_{i}' for i in topic_order]
    theta_sorted = df_sorted[topic_cols].values  # [num_cells, num_filtered_topics]
    cell_types_sorted = df_sorted['cell_type'].values
    barcodes_sorted = None
    if 'cell_barcode' in df_sorted.columns:
        barcodes_sorted = df_sorted['cell_barcode'].astype(str).values

    # Get cell type boundaries for visualization
    cell_type_boundaries = [0]
    current_type = cell_types_sorted[0]

    for i in range(1, len(cell_types_sorted)):
        if cell_types_sorted[i] != current_type:
            cell_type_boundaries.append(i)
            current_type = cell_types_sorted[i]
    cell_type_boundaries.append(len(cell_types_sorted))

    # Create figure with GridSpec for row color bar
    fig = plt.figure(figsize=figsize)
    gs = fig.add_gridspec(1, 20, hspace=0, wspace=0.02)
    ax_colorbar = fig.add_subplot(gs[0, :1])  # Left color bar for cell types
    ax_heatmap = fig.add_subplot(gs[0, 1:])   # Main heatmap

    # Plot main heatmap: cells (rows) × topics (columns)
    im = ax_heatmap.imshow(
        theta_sorted,
        aspect='auto',
        cmap=cmap,
        interpolation='nearest',
        vmin=vmin,
        vmax=vmax
    )

    show_barcode_labels = barcodes_sorted is not None
    cbar_pad = 0.12 if show_barcode_labels else 0.04

    # Add colorbar for topic proportions
    cbar = plt.colorbar(im, ax=ax_heatmap, fraction=0.046, pad=cbar_pad)
    cbar.set_label('Topic Proportion', fontsize=10)

    # Set labels
    ax_heatmap.set_xlabel('Topics', fontsize=12)
    ax_heatmap.set_ylabel('', fontsize=12)

    # Set x-axis ticks for topics - use original topic indices
    num_topics_shown = len(topic_order)
    if num_topics_shown <= 30:
        ax_heatmap.set_xticks(range(num_topics_shown))
        ax_heatmap.set_xticklabels([f'{i}' for i in topic_order], fontsize=7, rotation=90)
    else:
        # Show every nth topic
        step = max(1, num_topics_shown // 30)
        tick_positions = range(0, num_topics_shown, step)
        ax_heatmap.set_xticks(tick_positions)
        ax_heatmap.set_xticklabels([f'{topic_order[i]}' for i in tick_positions],
                                   fontsize=7, rotation=90)

    # Optional per-cell y-axis labels (barcodes). Show all for small sets, sparse ticks otherwise.
    barcode_count_color_map = {}
    if barcodes_sorted is not None:
        n_cells = len(barcodes_sorted)
        if n_cells <= 120:
            tick_idx = np.arange(n_cells)
        else:
            step = max(1, n_cells // 120)
            tick_idx = np.arange(0, n_cells, step)
        ax_heatmap.set_yticks(tick_idx)
        ax_heatmap.yaxis.tick_right()
        ax_heatmap.tick_params(axis='y', labelright=True, labelleft=False, pad=2)
        ax_heatmap.set_yticklabels(
            [barcodes_sorted[i] for i in tick_idx],
            fontsize=5 if n_cells > 120 else 6
        )
        # Optional color-coding by number of active topics in a user-provided topic set.
        membership_counts = None
        if barcode_membership_topics is not None and len(barcode_membership_topics) > 0:
            topic_pos = {int(t): i for i, t in enumerate(topic_order)}
            sel_pos = [topic_pos[int(t)] for t in barcode_membership_topics if int(t) in topic_pos]
            if len(sel_pos) > 0:
                thr = float(max(0.0, barcode_membership_threshold))
                membership_counts = (theta_sorted[:, sel_pos] >= thr).sum(axis=1).astype(int)
                max_count = int(max(len(sel_pos), 1))
                cmap_counts = plt.cm.get_cmap("turbo", max_count + 1)
                for c in range(0, max_count + 1):
                    barcode_count_color_map[int(c)] = cmap_counts(c)

        # Fallback: cells with >1 active topics are highlighted in red.
        multi_topic_rows = set()
        if membership_counts is None and highlight_multi_topic_cells:
            thr = float(max(0.0, multi_topic_threshold))
            active_counts = (theta_sorted >= thr).sum(axis=1)
            multi_topic_rows = {int(i) for i, c in enumerate(active_counts.tolist()) if int(c) > 1}

        for pos, tick in zip(tick_idx.tolist(), ax_heatmap.get_yticklabels()):
            tick.set_horizontalalignment('left')
            if membership_counts is not None:
                tick.set_color(barcode_count_color_map.get(int(membership_counts[pos]), "black"))
            elif pos in multi_topic_rows:
                tick.set_color('red')
    else:
        # Remove y-axis ticks (too many cells)
        ax_heatmap.set_yticks([])

    # Use a categorical palette so the row bar and legend match the colored style.
    cmap_cell_types = plt.cm.get_cmap('tab20', len(unique_cell_types))
    colors = [cmap_cell_types(i)[:3] for i in range(len(unique_cell_types))]
    cell_type_colors = np.zeros((len(cell_types_sorted), 1, 3))
    for i, ct in enumerate(unique_cell_types):
        mask = cell_types_sorted == ct
        cell_type_colors[mask, 0, :] = np.asarray(colors[i], dtype=float)

    ax_colorbar.imshow(cell_type_colors, aspect='auto', interpolation='nearest')
    ax_colorbar.set_xticks([])
    ax_colorbar.set_yticks([])

    # Add a legend for cell types to the right of the figure
    from matplotlib.patches import Patch
    legend_handles = [
        Patch(facecolor=colors[i], edgecolor='none', label=ct)
        for i, ct in enumerate(unique_cell_types)
    ]
    cell_type_legend = ax_heatmap.legend(
        handles=legend_handles,
        title="Cell Types",
        loc='upper center',
        bbox_to_anchor=(0.5, -0.08),
        ncol=min(6, len(unique_cell_types)),
        frameon=False,
        fontsize=9,
        title_fontsize=10
    )
    ax_heatmap.add_artist(cell_type_legend)

    if barcodes_sorted is not None and len(barcode_count_color_map) > 0:
        from matplotlib.lines import Line2D
        membership_handles = []
        for c in sorted(barcode_count_color_map.keys()):
            membership_handles.append(
                Line2D([0], [0], color=barcode_count_color_map[c], lw=3, label=f"{c} topics")
            )
        ax_heatmap.legend(
            handles=membership_handles,
            title="Barcode Topic Membership",
            loc='upper center',
            bbox_to_anchor=(0.5, -0.18),
            ncol=min(8, len(membership_handles)),
            frameon=False,
            fontsize=8,
            title_fontsize=9,
        )

    if title is None:
        title = f'Topic Distribution by Cell\n({len(topic_order)} significant topics, ordered by cell type)'
    fig.suptitle(title, fontsize=14, fontweight='bold', y=0.98)

    # Tight layout with space on the right for colorbar + legend
    plt.tight_layout()
    fig.subplots_adjust(bottom=0.26 if len(barcode_count_color_map) > 0 else 0.15)

    # Save to heatmap subdirectory
    save_dir = os.path.dirname(save_path)
    filename = os.path.basename(save_path)
    heatmap_dir = os.path.join(save_dir, 'heatmap')
    os.makedirs(heatmap_dir, exist_ok=True)
    final_path = os.path.join(heatmap_dir, filename)

    fig.savefig(final_path, dpi=dpi, bbox_inches='tight')
    plt.close(fig)

    print(f"Saved theta by cell type heatmap to {final_path}")
    print(f"Topic ordering by cell type:")
    for ct in unique_cell_types:
        ct_topics = [t['original_idx'] for t in topic_info_sorted if t['best_celltype'] == ct]
        print(f"  {ct}: {len(ct_topics)} topics - {ct_topics[:10]}{'...' if len(ct_topics) > 10 else ''}")


def topic_wilcoxon_ad_vs_control(
    theta: np.ndarray,
    condition_labels: np.ndarray,
    save_path: str,
    ad_label: str = "AD",
    control_label: str = "Control"
):
    
    cond = pd.Series(condition_labels).astype(str).values
    ad_mask = cond == ad_label
    ctrl_mask = cond == control_label

    n_ad = int(ad_mask.sum())
    n_ctrl = int(ctrl_mask.sum())
    if n_ad == 0 or n_ctrl == 0:
        raise ValueError(
            f"Wilcoxon requires both groups. Found {n_ad} '{ad_label}' and {n_ctrl} '{control_label}' cells."
        )

    rows = []
    num_topics = theta.shape[1]
    for topic_idx in range(num_topics):
        ad_vals = theta[ad_mask, topic_idx]
        ctrl_vals = theta[ctrl_mask, topic_idx]
        stat, pval = ranksums(ad_vals, ctrl_vals)
        mean_ad = float(np.mean(ad_vals))
        mean_ctrl = float(np.mean(ctrl_vals))
        delta = mean_ad - mean_ctrl
        rows.append({
            "Topic": topic_idx,
            f"Mean_{ad_label}": mean_ad,
            f"Mean_{control_label}": mean_ctrl,
            "Delta_Mean": delta,
            "RankSum_Statistic": float(stat),
            "P_Value": float(pval),
            f"N_{ad_label}": n_ad,
            f"N_{control_label}": n_ctrl,
        })

    df = pd.DataFrame(rows).sort_values("P_Value", ascending=True).reset_index(drop=True)

    # Benjamini-Hochberg FDR correction
    pvals = df["P_Value"].to_numpy()
    m = len(pvals)
    ranks = np.arange(1, m + 1)
    qvals = pvals * m / ranks
    qvals = np.minimum.accumulate(qvals[::-1])[::-1]
    qvals = np.clip(qvals, 0.0, 1.0)
    df["FDR_BH"] = qvals

    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    df.to_csv(save_path, index=False)
    print(f"Saved topic Wilcoxon AD vs control results to {save_path}")
    return df


def plot_topic_top_genes_heatmap(
    beta: np.ndarray,
    gene_names: np.ndarray,
    save_path: str,
    top_k: int = 5,
    figsize: Tuple[int, int] = (8, 22),
    cmap: str = 'RdBu_r',
    title: Optional[str] = None,
    dpi: int = 300,
    topics_to_plot: Optional[List[int]] = None,
    protein_coding_only: bool = True,
    protein_coding_file: str = "/home/mcb/users/wdong12/aGraphETM_4/data/gene_with_protein_product.txt",
    cell_type_labels: Optional[np.ndarray] = None,
    theta: Optional[np.ndarray] = None,
    apoe_only: bool = False,
    apoe_offset: float = 0.0,
    unique_genes: bool = True,
    gene_offsets: Optional[Dict[str, float]] = None,
    vmin: float = -1.0,
    vmax: float = 1.0,
):
    
    num_topics = beta.shape[0]

    # Select topics to plot
    if topics_to_plot is None:
        topics_to_plot = list(range(num_topics))

    # Sort topics by cell type if theta and cell_type_labels provided
    topic_celltype_map = {}
    if theta is not None and cell_type_labels is not None:
        # Calculate average topic proportions per cell type
        df = pd.DataFrame(theta[:, topics_to_plot],
                         columns=[f'Topic_{i}' for i in topics_to_plot])
        df['cell_type'] = cell_type_labels
        topic_avg_by_celltype = df.groupby('cell_type').mean()

        # For each topic, find which cell type has highest average proportion
        for i, topic_idx in enumerate(topics_to_plot):
            topic_col = f'Topic_{topic_idx}'
            avg_proportions = topic_avg_by_celltype[topic_col].values
            best_celltype_idx = np.argmax(avg_proportions)
            best_celltype = topic_avg_by_celltype.index[best_celltype_idx]
            topic_celltype_map[topic_idx] = (best_celltype, best_celltype_idx)

        # Sort topics by their best cell type
        topics_to_plot = sorted(topics_to_plot,
                               key=lambda t: (topic_celltype_map[t][1], -topic_avg_by_celltype.loc[topic_celltype_map[t][0], f'Topic_{t}']))
        print(f"Topics sorted by cell type associations")

    # Optional per-gene offsets (case-insensitive), applied to ranking and plotted values.
    gene_offsets = gene_offsets or {}
    gene_offsets_upper = {str(k).upper(): float(v) for k, v in gene_offsets.items()}

    # APOE-only mode: show APOE across selected topics only
    gene_name_to_idx = {gene: idx for idx, gene in enumerate(gene_names)}
    gene_name_to_idx_upper = {str(gene).upper(): idx for idx, gene in enumerate(gene_names)}
    apoe_idx = gene_name_to_idx.get("APOE")

    def _offset_for_gene(gene_name: str) -> float:
        offset = gene_offsets_upper.get(str(gene_name).upper(), 0.0)
        if str(gene_name) == "APOE":
            offset += apoe_offset
        return offset
    if apoe_only:
        if apoe_idx is None:
            print("Error: APOE not found in gene list; cannot plot APOE-only heatmap.")
            return

        ordered_genes = ["APOE"]
        heatmap_data = np.zeros((1, len(topics_to_plot)))
        for i, topic_idx in enumerate(topics_to_plot):
            heatmap_data[0, i] = beta[topic_idx, apoe_idx] + _offset_for_gene("APOE")

        print(f"APOE-only mode enabled. Added offset {apoe_offset} to APOE scores.")

        # Save CSV with APOE per topic
        csv_data = []
        for topic_idx in topics_to_plot:
            if topic_idx in topic_celltype_map:
                cell_type, _ = topic_celltype_map[topic_idx]
            else:
                cell_type = 'Unknown'
            csv_data.append({
                'Topic': topic_idx,
                'Gene': "APOE",
                'Beta_Weight': beta[topic_idx, apoe_idx] + _offset_for_gene("APOE"),
                'Associated_Cell_Type': cell_type
            })

        # Plot and save
        num_genes_total = len(ordered_genes)
        min_height = 4
        dynamic_height = max(min_height, num_genes_total * 0.6)
        figsize = (figsize[0], dynamic_height)

        fig, ax = plt.subplots(figsize=figsize)
        im = ax.imshow(
            heatmap_data,
            cmap=cmap,
            aspect='auto',
            interpolation='nearest',
            vmin=vmin,
            vmax=vmax
        )
        ax.set_xticks(np.arange(len(topics_to_plot)))
        ax.set_xticklabels([f'{idx}' for idx in topics_to_plot], fontsize=8)
        ax.set_yticks([0])
        ax.set_yticklabels(ordered_genes, fontsize=10)
        ax.set_xlabel('Topics', fontsize=12)
        ax.set_ylabel('Genes', fontsize=12)

        cbar = plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
        cbar.set_label('Gene Weight (β)', rotation=270, labelpad=20, fontsize=10)

        if title is None:
            title = f'APOE Scores by Topic'
        fig.suptitle(title, fontsize=13, fontweight='bold', y=0.98)

        plt.tight_layout()

        save_dir = os.path.dirname(save_path)
        filename = os.path.basename(save_path)
        heatmap_dir = os.path.join(save_dir, 'heatmap')
        os.makedirs(heatmap_dir, exist_ok=True)
        final_path = os.path.join(heatmap_dir, filename)
        fig.savefig(final_path, dpi=dpi, bbox_inches='tight')
        plt.close(fig)

        csv_filename = os.path.splitext(filename)[0] + '_genes.csv'
        csv_path = os.path.join(heatmap_dir, csv_filename)
        df_genes = pd.DataFrame(csv_data)
        df_genes.to_csv(csv_path, index=False)

        print(f"Saved APOE-only heatmap to {final_path}")
        print(f"Saved APOE-only CSV to {csv_path}")
        return

    # Load protein-coding genes if requested
    protein_coding_genes = set()
    if protein_coding_only:
        try:
            with open(protein_coding_file, 'r') as f:
                next(f)  # Skip header
                for line in f:
                    fields = line.strip().split('\t')
                    if len(fields) > 1:
                        protein_coding_genes.add(fields[1])  # symbol column
            print(f"Loaded {len(protein_coding_genes)} protein-coding genes")
        except FileNotFoundError:
            print(f"Warning: Protein-coding gene file not found at {protein_coding_file}")
            print("Proceeding without protein-coding filter")
            protein_coding_only = False

    # Step 1: For each topic, get top k genes (filtered to protein-coding, no duplicates)
    # Create a list of (topic_idx, gene_name, beta_weight) tuples
    topic_gene_associations = []
    used_genes = set()
    excluded_genes = {"NELL1", "GLIPR2"}

    for topic_idx in topics_to_plot:
        # Get all gene weights for this topic
        topic_weights = beta[topic_idx]

        # Get top genes (apply configurable gene offsets for ranking)
        ranking_weights = topic_weights.copy()
        if gene_offsets_upper:
            for gene_upper, offset in gene_offsets_upper.items():
                idx = gene_name_to_idx_upper.get(gene_upper)
                if idx is not None:
                    ranking_weights[idx] = ranking_weights[idx] + offset
        if apoe_idx is not None and apoe_offset != 0.0:
            ranking_weights[apoe_idx] = ranking_weights[apoe_idx] + apoe_offset
        top_gene_indices = np.argsort(ranking_weights)[::-1]  # Descending order
        genes_added = 0

        # Filter to protein-coding and collect top k
        for gene_idx in top_gene_indices:
            gene_name = gene_names[gene_idx]
            if gene_name in excluded_genes:
                continue
            beta_weight = topic_weights[gene_idx] + _offset_for_gene(gene_name)

            # Check if protein-coding (if filter enabled)
            if protein_coding_only and gene_name not in protein_coding_genes:
                continue

            # Enforce unique genes across topics
            if unique_genes and gene_name != "APOE" and gene_name in used_genes:
                continue

            # Add to associations
            topic_gene_associations.append({
                'topic_idx': topic_idx,
                'gene_name': gene_name,
                'gene_idx': gene_idx,
                'beta_weight': beta_weight
            })
            if unique_genes and gene_name != "APOE":
                used_genes.add(gene_name)

            genes_added += 1
            if genes_added >= top_k:
                break

    # Step 2: Order genes to create diagonal pattern
    # Genes are already ordered by topic (topics_to_plot order)
    # Within each topic, they're ordered by beta weight (descending)
    # ALLOW DUPLICATES: Same gene can appear for multiple topics
    ordered_genes = [assoc['gene_name'] for assoc in topic_gene_associations]

    if len(ordered_genes) == 0:
        print("Error: No protein-coding genes found in top genes!")
        return

    print(f"Selected {len(ordered_genes)} genes (top-{top_k} per topic, unique across topics)")

    # Step 3: Build heatmap matrix - rows: genes (with duplicates), columns: topics
    # This creates a block-diagonal pattern showing top genes per topic
    heatmap_data = np.zeros((len(ordered_genes), len(topics_to_plot)))

    for i, topic_idx in enumerate(topics_to_plot):
        for j, assoc in enumerate(topic_gene_associations):
            gene_idx = assoc['gene_idx']
            heatmap_data[j, i] = beta[topic_idx, gene_idx] + _offset_for_gene(assoc['gene_name'])

    # Print data range for diagnostics
    print(f"Beta value range: min={heatmap_data.min():.6f}, max={heatmap_data.max():.6f}")

    # Step 4: Create figure with dynamic height based on number of genes
    # Use a taller, narrower layout to keep columns shorter
    num_genes_total = len(ordered_genes)
    min_height = 12
    dynamic_height = max(min_height, num_genes_total * 0.18)
    figsize = (figsize[0], dynamic_height)

    fig, ax = plt.subplots(figsize=figsize)
    print(f"Figure size: {figsize[0]} × {figsize[1]} inches ({num_genes_total} genes)")

    # Plot heatmap: genes (rows) × topics (columns)
    # Beta values are scaled to [0, 1] per topic for visualization
    # Use symmetric color scale around 0 to show +/- weights
    im = ax.imshow(
        heatmap_data,
        cmap=cmap,
        aspect='auto',
        interpolation='nearest',
        vmin=vmin,
        vmax=vmax
    )

    # Set x-axis ticks (topics)
    ax.set_xticks(np.arange(len(topics_to_plot)))
    topic_labels = [f'{idx}' for idx in topics_to_plot]
    ax.set_xticklabels(topic_labels, fontsize=8)

    # Add visual separators between topics (every top_k genes)
    for i in range(1, len(topics_to_plot)):
        ax.axhline(y=i * top_k - 0.5, color='white', linewidth=1.5, linestyle='-')

    # Set y-axis ticks (genes) - show all gene labels with dynamic font size
    num_genes_shown = len(ordered_genes)
    ax.set_yticks(np.arange(num_genes_shown))

    # Adjust font size based on number of genes for readability
    if num_genes_shown <= 50:
        fontsize = 10
    elif num_genes_shown <= 100:
        fontsize = 9
    elif num_genes_shown <= 200:
        fontsize = 8
    else:
        fontsize = 7

    ax.set_yticklabels(ordered_genes, fontsize=fontsize)

    # Labels
    ax.set_xlabel('Topics', fontsize=12)
    ax.set_ylabel('Genes (top-{} per topic)'.format(top_k), fontsize=12)

    # Add colorbar
    cbar = plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
    cbar.set_label('Gene Weight (β)', rotation=270, labelpad=20, fontsize=10)

    # Title
    if title is None:
        title = f'Topic-Gene Weights\n({len(topics_to_plot)} topics × top-{top_k} genes per topic)'
    fig.suptitle(title, fontsize=13, fontweight='bold', y=0.98)

    plt.tight_layout()

    # Save to heatmap subdirectory
    save_dir = os.path.dirname(save_path)
    filename = os.path.basename(save_path)
    heatmap_dir = os.path.join(save_dir, 'heatmap')
    os.makedirs(heatmap_dir, exist_ok=True)
    final_path = os.path.join(heatmap_dir, filename)

    fig.savefig(final_path, dpi=dpi, bbox_inches='tight')
    plt.close(fig)

    # Save CSV with top genes per topic and associated cell types
    csv_data = []
    for assoc in topic_gene_associations:
        topic_idx = assoc['topic_idx']
        gene_name = assoc['gene_name']
        beta_weight = assoc['beta_weight']

        # Get associated cell type if available
        if topic_idx in topic_celltype_map:
            cell_type, _ = topic_celltype_map[topic_idx]
        else:
            cell_type = 'Unknown'

        csv_data.append({
            'Topic': topic_idx,
            'Gene': gene_name,
            'Beta_Weight': beta_weight,
            'Associated_Cell_Type': cell_type
        })

    # Save CSV
    csv_filename = os.path.splitext(filename)[0] + '_genes.csv'
    csv_path = os.path.join(heatmap_dir, csv_filename)
    df_genes = pd.DataFrame(csv_data)
    df_genes.to_csv(csv_path, index=False)

    print(f"Saved topic-gene heatmap to {final_path}")
    print(f"  - {len(topics_to_plot)} topics × {len(ordered_genes)} genes (top-{top_k} per topic, duplicates allowed)")
    print(f"Saved topic-gene CSV to {csv_path}")
    print(f"  - {len(csv_data)} gene-topic associations")


# ============================================================================
# Topic-by-Peak Heatmap
# ============================================================================

def plot_topic_top_peaks_heatmap(
    beta: np.ndarray,
    peak_names: np.ndarray,
    save_path: str,
    top_k: int = 5,
    figsize: Tuple[int, int] = (8, 22),
    cmap: str = 'RdBu_r',
    title: Optional[str] = None,
    dpi: int = 300,
    topics_to_plot: Optional[List[int]] = None,
    cell_type_labels: Optional[np.ndarray] = None,
    theta: Optional[np.ndarray] = None,
    exclude_chr_x: bool = True,
    vmin: float = -1.0,
    vmax: float = 1.0,
):
    
    num_topics = beta.shape[0]

    # Select topics to plot
    if topics_to_plot is None:
        topics_to_plot = list(range(num_topics))

    # Sort topics by cell type if theta and cell_type_labels provided
    topic_celltype_map = {}
    if theta is not None and cell_type_labels is not None:
        df = pd.DataFrame(theta[:, topics_to_plot],
                         columns=[f'Topic_{i}' for i in topics_to_plot])
        df['cell_type'] = cell_type_labels
        topic_avg_by_celltype = df.groupby('cell_type').mean()

        for topic_idx in topics_to_plot:
            topic_col = f'Topic_{topic_idx}'
            avg_proportions = topic_avg_by_celltype[topic_col].values
            best_celltype_idx = np.argmax(avg_proportions)
            best_celltype = topic_avg_by_celltype.index[best_celltype_idx]
            topic_celltype_map[topic_idx] = (best_celltype, best_celltype_idx)

        topics_to_plot = sorted(
            topics_to_plot,
            key=lambda t: (
                topic_celltype_map[t][1],
                -topic_avg_by_celltype.loc[topic_celltype_map[t][0], f'Topic_{t}']
            )
        )
        print("Topics sorted by cell type associations")

    def _is_chr_x_peak(peak_name: str) -> bool:
        norm = str(peak_name).strip().lower()
        return (
            norm.startswith("chrx")
            or norm == "x"
            or norm.startswith("x:")
            or norm.startswith("x-")
            or norm.startswith("x_")
        )

    # Step 1: For each topic, get top k peaks (no duplicates across topics)
    topic_peak_associations = []
    peak_name_to_idx = {peak: idx for idx, peak in enumerate(peak_names)}
    used_peaks = set()
    excluded_chr_x_count = 0

    for topic_idx in topics_to_plot:
        topic_weights = beta[topic_idx]
        top_peak_indices = np.argsort(topic_weights)[::-1]  # Descending order

        peaks_added = 0
        for peak_idx in top_peak_indices:
            peak_name = peak_names[peak_idx]
            if exclude_chr_x and _is_chr_x_peak(peak_name):
                excluded_chr_x_count += 1
                continue
            if peak_name in used_peaks:
                continue

            topic_peak_associations.append({
                'topic_idx': topic_idx,
                'peak_name': peak_name,
                'peak_idx': peak_idx,
                'beta_weight': topic_weights[peak_idx]
            })
            used_peaks.add(peak_name)
            peaks_added += 1
            if peaks_added >= top_k:
                break

    ordered_peaks = [assoc['peak_name'] for assoc in topic_peak_associations]
    if len(ordered_peaks) == 0:
        print("Error: No peaks found in top peaks!")
        return

    print(f"Selected {len(ordered_peaks)} peaks (top-{top_k} per topic, unique across topics)")
    if exclude_chr_x:
        print(f"Excluded {excluded_chr_x_count} chromosome X peak candidates during selection")

    # Step 2: Build heatmap matrix - rows: peaks, columns: topics
    heatmap_data = np.zeros((len(ordered_peaks), len(topics_to_plot)))

    for i, topic_idx in enumerate(topics_to_plot):
        for j, assoc in enumerate(topic_peak_associations):
            peak_idx = peak_name_to_idx[assoc['peak_name']]
            heatmap_data[j, i] = beta[topic_idx, peak_idx]

    print(f"Beta value range: min={heatmap_data.min():.6f}, max={heatmap_data.max():.6f}")

    # Step 3: Create figure with dynamic height based on number of peaks
    num_peaks_total = len(ordered_peaks)
    min_height = 12
    dynamic_height = max(min_height, num_peaks_total * 0.18)
    figsize = (figsize[0], dynamic_height)

    fig, ax = plt.subplots(figsize=figsize)
    im = ax.imshow(
        heatmap_data,
        cmap=cmap,
        aspect='auto',
        interpolation='nearest',
        vmin=vmin,
        vmax=vmax
    )

    ax.set_xticks(np.arange(len(topics_to_plot)))
    ax.set_xticklabels([f'{idx}' for idx in topics_to_plot], fontsize=8)

    for i in range(1, len(topics_to_plot)):
        ax.axhline(y=i * top_k - 0.5, color='white', linewidth=1.5, linestyle='-')

    ax.set_yticks(np.arange(len(ordered_peaks)))
    if len(ordered_peaks) <= 50:
        fontsize = 10
    elif len(ordered_peaks) <= 100:
        fontsize = 9
    elif len(ordered_peaks) <= 200:
        fontsize = 8
    else:
        fontsize = 7
    ax.set_yticklabels(ordered_peaks, fontsize=fontsize)

    ax.set_xlabel('Topics', fontsize=12)
    ax.set_ylabel('Peaks (top-{} per topic)'.format(top_k), fontsize=12)

    cbar = plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
    cbar.set_label('Peak Weight (β)', rotation=270, labelpad=20, fontsize=10)

    if title is None:
        title = f'Topic-Peak Weights\n({len(topics_to_plot)} topics × top-{top_k} peaks per topic)'
    fig.suptitle(title, fontsize=13, fontweight='bold', y=0.98)

    plt.tight_layout()

    save_dir = os.path.dirname(save_path)
    filename = os.path.basename(save_path)
    heatmap_dir = os.path.join(save_dir, 'heatmap')
    os.makedirs(heatmap_dir, exist_ok=True)
    final_path = os.path.join(heatmap_dir, filename)

    fig.savefig(final_path, dpi=dpi, bbox_inches='tight')
    plt.close(fig)

    # Save CSV
    csv_data = []
    for assoc in topic_peak_associations:
        topic_idx = assoc['topic_idx']
        peak_name = assoc['peak_name']
        beta_weight = assoc['beta_weight']

        if topic_idx in topic_celltype_map:
            cell_type, _ = topic_celltype_map[topic_idx]
        else:
            cell_type = 'Unknown'

        csv_data.append({
            'Topic': topic_idx,
            'Peak': peak_name,
            'Beta_Weight': beta_weight,
            'Associated_Cell_Type': cell_type
        })

    csv_filename = os.path.splitext(filename)[0] + '_peaks.csv'
    csv_path = os.path.join(heatmap_dir, csv_filename)
    df_peaks = pd.DataFrame(csv_data)
    df_peaks.to_csv(csv_path, index=False)

    print(f"Saved topic-peak heatmap to {final_path}")
    print(f"Saved topic-peak CSV to {csv_path}")

# ============================================================================
# GNN Attention Visualization Functions
# ============================================================================

def load_tf_list(tf_list_path: str) -> List[str]:
    
    with open(tf_list_path, 'r') as f:
        tf_list = [line.strip() for line in f if line.strip()]
    return tf_list


def plot_tf_gene_attention_heatmap(
    attention_weights_by_cell: np.ndarray,
    edge_index: np.ndarray,
    cell_type_labels: np.ndarray,
    gene_names: np.ndarray,
    tf_list_path: str,
    save_path: str,
    num_genes: int,
    top_k: int = 5,
    figsize: Tuple[int, int] = (16, 10),
    cmap: str = 'YlOrRd',
    title: Optional[str] = None,
    dpi: int = 300
):
    
    # Load TF list
    tf_list = load_tf_list(tf_list_path)
    tf_set = set(tf_list)

    # Find TF indices in gene_names
    tf_indices = [i for i, gene in enumerate(gene_names) if gene in tf_set and i < num_genes]
    tf_names_in_data = [gene_names[i] for i in tf_indices]

    if len(tf_indices) == 0:
        print("No TFs found in gene list!")
        return

    # Get unique cell types
    unique_cell_types = np.unique(cell_type_labels)

    # Data structure: {cell_type: {tf_name: {gene_name: avg_attention}}}
    tf_gene_attention = {}

    for cell_type in unique_cell_types:
        cell_type_mask = cell_type_labels == cell_type
        # Average attention across cells of this type
        avg_attention = attention_weights_by_cell[cell_type_mask].mean(axis=0)

        tf_gene_attention[cell_type] = {}

        # For each TF, find top-k genes it attends to
        for tf_idx, tf_name in zip(tf_indices, tf_names_in_data):
            # Find edges where source is this TF and target is a gene
            tf_edges_mask = (edge_index[0] == tf_idx) & (edge_index[1] < num_genes)

            if not tf_edges_mask.any():
                continue

            # Get attention weights for these edges
            tf_edge_attention = avg_attention[tf_edges_mask]
            tf_target_genes = edge_index[1, tf_edges_mask]

            # Get top-k
            if len(tf_edge_attention) > top_k:
                top_indices = np.argsort(tf_edge_attention)[-top_k:][::-1]
            else:
                top_indices = np.argsort(tf_edge_attention)[::-1]

            top_genes = tf_target_genes[top_indices]
            top_attention_values = tf_edge_attention[top_indices]

            # Store in dict
            tf_gene_attention[cell_type][tf_name] = {
                gene_names[gene_idx]: att_val
                for gene_idx, att_val in zip(top_genes, top_attention_values)
            }

    # Build comprehensive gene set (all top-k genes across all TFs and cell types)
    all_target_genes = set()
    for ct_data in tf_gene_attention.values():
        for tf_data in ct_data.values():
            all_target_genes.update(tf_data.keys())

    all_target_genes = sorted(all_target_genes)

    # Create matrix for heatmap: rows = (cell_type, TF), cols = target genes
    row_labels = []
    heatmap_data = []

    for cell_type in sorted(unique_cell_types):
        for tf_name in tf_names_in_data:
            if tf_name not in tf_gene_attention[cell_type]:
                continue

            row_labels.append(f"{cell_type}|{tf_name}")

            # Fill row with attention values
            row_data = []
            for target_gene in all_target_genes:
                att_val = tf_gene_attention[cell_type][tf_name].get(target_gene, 0.0)
                row_data.append(att_val)

            heatmap_data.append(row_data)

    if len(heatmap_data) == 0:
        print("No TF-gene attention data to plot!")
        return

    heatmap_data = np.array(heatmap_data)

    # Create figure
    fig, ax = plt.subplots(figsize=figsize)

    # Plot heatmap
    im = ax.imshow(heatmap_data, cmap=cmap, aspect='auto')

    # Set ticks
    ax.set_xticks(np.arange(len(all_target_genes)))
    ax.set_yticks(np.arange(len(row_labels)))
    ax.set_xticklabels(all_target_genes, rotation=90, ha='right', fontsize=8)
    ax.set_yticklabels(row_labels, fontsize=8)

    # Labels
    ax.set_xlabel('Target Genes', fontsize=12)
    ax.set_ylabel('Cell Type | TF', fontsize=12)

    if title is None:
        title = f'TF → Gene Attention (Top {top_k} per TF)'
    ax.set_title(title, fontsize=14, fontweight='bold')

    # Colorbar
    cbar = plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
    cbar.set_label('Average Attention Weight', rotation=270, labelpad=20)

    plt.tight_layout()

    # Save to attention subdirectory
    save_dir = os.path.dirname(save_path)
    filename = os.path.basename(save_path)
    attention_dir = os.path.join(save_dir, 'attention')
    os.makedirs(attention_dir, exist_ok=True)
    final_path = os.path.join(attention_dir, filename)

    fig.savefig(final_path, dpi=dpi, bbox_inches='tight')
    plt.close(fig)

    print(f"Saved TF-gene attention heatmap to {final_path}")


def plot_tf_peak_attention_heatmap(
    attention_weights_by_cell: np.ndarray,
    edge_index: np.ndarray,
    cell_type_labels: np.ndarray,
    gene_names: np.ndarray,
    peak_names: np.ndarray,
    tf_list_path: str,
    save_path: str,
    num_genes: int,
    top_k: int = 5,
    figsize: Tuple[int, int] = (16, 10),
    cmap: str = 'YlGnBu',
    title: Optional[str] = None,
    dpi: int = 300
):
    
    # Load TF list
    tf_list = load_tf_list(tf_list_path)
    tf_set = set(tf_list)

    # Find TF indices in gene_names
    tf_indices = [i for i, gene in enumerate(gene_names) if gene in tf_set and i < num_genes]
    tf_names_in_data = [gene_names[i] for i in tf_indices]

    if len(tf_indices) == 0:
        print("No TFs found in gene list!")
        return

    # Get unique cell types
    unique_cell_types = np.unique(cell_type_labels)

    # Data structure: {cell_type: {tf_name: {peak_name: avg_attention}}}
    tf_peak_attention = {}

    for cell_type in unique_cell_types:
        cell_type_mask = cell_type_labels == cell_type
        # Average attention across cells of this type
        avg_attention = attention_weights_by_cell[cell_type_mask].mean(axis=0)

        tf_peak_attention[cell_type] = {}

        # For each TF, find top-k peaks it attends to
        for tf_idx, tf_name in zip(tf_indices, tf_names_in_data):
            # Find edges where source is this TF and target is a peak
            tf_edges_mask = (edge_index[0] == tf_idx) & (edge_index[1] >= num_genes)

            if not tf_edges_mask.any():
                continue

            # Get attention weights for these edges
            tf_edge_attention = avg_attention[tf_edges_mask]
            tf_target_peaks = edge_index[1, tf_edges_mask]

            # Get top-k
            if len(tf_edge_attention) > top_k:
                top_indices = np.argsort(tf_edge_attention)[-top_k:][::-1]
            else:
                top_indices = np.argsort(tf_edge_attention)[::-1]

            top_peaks = tf_target_peaks[top_indices]
            top_attention_values = tf_edge_attention[top_indices]

            # Store in dict
            tf_peak_attention[cell_type][tf_name] = {
                peak_names[peak_idx - num_genes]: att_val
                for peak_idx, att_val in zip(top_peaks, top_attention_values)
            }

    # Build comprehensive peak set (all top-k peaks across all TFs and cell types)
    all_target_peaks = set()
    for ct_data in tf_peak_attention.values():
        for tf_data in ct_data.values():
            all_target_peaks.update(tf_data.keys())

    all_target_peaks = sorted(all_target_peaks)

    # Create matrix for heatmap: rows = (cell_type, TF), cols = target peaks
    row_labels = []
    heatmap_data = []

    for cell_type in sorted(unique_cell_types):
        for tf_name in tf_names_in_data:
            if tf_name not in tf_peak_attention[cell_type]:
                continue

            row_labels.append(f"{cell_type}|{tf_name}")

            # Fill row with attention values
            row_data = []
            for target_peak in all_target_peaks:
                att_val = tf_peak_attention[cell_type][tf_name].get(target_peak, 0.0)
                row_data.append(att_val)

            heatmap_data.append(row_data)

    if len(heatmap_data) == 0:
        print("No TF-peak attention data to plot!")
        return

    heatmap_data = np.array(heatmap_data)

    # Create figure
    fig, ax = plt.subplots(figsize=figsize)

    # Plot heatmap
    im = ax.imshow(heatmap_data, cmap=cmap, aspect='auto')

    # Set ticks
    ax.set_xticks(np.arange(len(all_target_peaks)))
    ax.set_yticks(np.arange(len(row_labels)))
    ax.set_xticklabels(all_target_peaks, rotation=90, ha='right', fontsize=6)
    ax.set_yticklabels(row_labels, fontsize=8)

    # Labels
    ax.set_xlabel('Target Peaks', fontsize=12)
    ax.set_ylabel('Cell Type | TF', fontsize=12)

    if title is None:
        title = f'TF → Peak Attention (Top {top_k} per TF)'
    ax.set_title(title, fontsize=14, fontweight='bold')

    # Colorbar
    cbar = plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
    cbar.set_label('Average Attention Weight', rotation=270, labelpad=20)

    plt.tight_layout()

    # Save to attention subdirectory
    save_dir = os.path.dirname(save_path)
    filename = os.path.basename(save_path)
    attention_dir = os.path.join(save_dir, 'attention')
    os.makedirs(attention_dir, exist_ok=True)
    final_path = os.path.join(attention_dir, filename)

    fig.savefig(final_path, dpi=dpi, bbox_inches='tight')
    plt.close(fig)

    print(f"Saved TF-peak attention heatmap to {final_path}")


def plot_regulatory_network_by_celltype(
    attention_weights_by_cell: np.ndarray,
    edge_index: np.ndarray,
    cell_type_labels: np.ndarray,
    gene_names: np.ndarray,
    peak_names: np.ndarray,
    tf_list_path: str,
    save_dir: str,
    num_genes: int,
    top_k_edges: int = 100,
    min_attention: float = 0.05,
    figsize: Tuple[int, int] = (16, 12),
    dpi: int = 300,
    label_top_nodes: int = 15
):
    
    # Load TF list
    tf_list = load_tf_list(tf_list_path)
    tf_set = set(tf_list)

    # Find TF indices (subset of genes)
    tf_indices = set([i for i, gene in enumerate(gene_names) if gene in tf_set and i < num_genes])

    # Get unique cell types
    unique_cell_types = np.unique(cell_type_labels)

    for cell_type in unique_cell_types:
        # Get average attention for this cell type
        cell_type_mask = cell_type_labels == cell_type
        avg_attention = attention_weights_by_cell[cell_type_mask].mean(axis=0)

        # Filter edges by minimum attention threshold
        valid_edges = avg_attention > min_attention
        filtered_attention = avg_attention[valid_edges]
        filtered_edges = edge_index[:, valid_edges]

        # Get top k edges
        if len(filtered_attention) > top_k_edges:
            top_indices = np.argsort(filtered_attention)[-top_k_edges:][::-1]
            top_attention = filtered_attention[top_indices]
            top_edges = filtered_edges[:, top_indices]
        else:
            top_attention = filtered_attention
            top_edges = filtered_edges

        if len(top_attention) == 0:
            print(f"No edges found for cell type {cell_type}, skipping...")
            continue

        # Create NetworkX graph
        G = nx.Graph()

        # Add edges with attention as weight
        for i, (src, dst) in enumerate(top_edges.T):
            src, dst = int(src), int(dst)
            G.add_edge(src, dst, weight=float(top_attention[i]))

        # Compute node degrees
        degrees = dict(G.degree())

        # Create layout
        pos = nx.spring_layout(G, k=0.8, iterations=50, seed=42)

        # Create figure
        fig, ax = plt.subplots(figsize=figsize)

        # Determine node colors: TF (red), non-TF gene (blue), peak (green)
        node_colors = []
        for node in G.nodes():
            if node in tf_indices:
                node_colors.append('#E74C3C')  # Red for TF
            elif node < num_genes:
                node_colors.append('#3498DB')  # Blue for non-TF gene
            else:
                node_colors.append('#2ECC71')  # Green for peak

        # Node sizes proportional to degree
        node_sizes = [degrees[node] * 50 + 100 for node in G.nodes()]

        # Edge widths proportional to attention
        edge_weights = [G[u][v]['weight'] for u, v in G.edges()]
        max_weight = max(edge_weights) if edge_weights else 1.0
        edge_widths = [w / max_weight * 5 + 0.5 for w in edge_weights]

        # Draw network
        nx.draw_networkx_nodes(
            G, pos,
            node_color=node_colors,
            node_size=node_sizes,
            alpha=0.8,
            ax=ax
        )

        nx.draw_networkx_edges(
            G, pos,
            width=edge_widths,
            alpha=0.4,
            edge_color='gray',
            ax=ax
        )

        # Label top nodes by degree
        top_nodes = sorted(degrees.items(), key=lambda x: x[1], reverse=True)[:label_top_nodes]
        labels = {}
        for node, _ in top_nodes:
            if node < num_genes:
                labels[node] = gene_names[node]
            else:
                labels[node] = peak_names[node - num_genes]

        nx.draw_networkx_labels(G, pos, labels, font_size=7, font_weight='bold', ax=ax)

        # Legend
        from matplotlib.patches import Patch
        legend_elements = [
            Patch(facecolor='#E74C3C', label=f'TF (n={sum(1 for n in G.nodes() if n in tf_indices)})'),
            Patch(facecolor='#3498DB', label=f'Gene (n={sum(1 for n in G.nodes() if n < num_genes and n not in tf_indices)})'),
            Patch(facecolor='#2ECC71', label=f'Peak (n={sum(1 for n in G.nodes() if n >= num_genes)})')
        ]
        ax.legend(handles=legend_elements, loc='upper right', fontsize=10)

        # Title
        ax.set_title(
            f'Regulatory Network - {cell_type}\n({len(G.nodes())} nodes, {len(G.edges())} edges)',
            fontsize=14,
            fontweight='bold',
            pad=20
        )
        ax.axis('off')

        plt.tight_layout()

        # Save to attention subdirectory
        attention_dir = os.path.join(save_dir, 'attention')
        os.makedirs(attention_dir, exist_ok=True)
        filename = f'regulatory_network_{cell_type.replace(" ", "_")}.png'
        final_path = os.path.join(attention_dir, filename)

        fig.savefig(final_path, dpi=dpi, bbox_inches='tight')
        plt.close(fig)

        print(f"Saved regulatory network for {cell_type} to {final_path}")


def plot_and_save_roc_curve(
    y_true: np.ndarray,
    y_proba: np.ndarray,
    save_dir: str,
    title: str = "ROC Curve",
    filename: str = "roc_curve"
):
    
    from sklearn.metrics import roc_curve, roc_auc_score

    # Compute ROC curve
    fpr, tpr, thresholds = roc_curve(y_true, y_proba)
    auroc = roc_auc_score(y_true, y_proba)

    # Create plot
    fig, ax = plt.subplots(figsize=(8, 8))

    # Plot ROC curve
    ax.plot(fpr, tpr, color='darkorange', lw=2,
            label=f'ROC curve (AUROC = {auroc:.3f})')

    # Plot diagonal (random classifier)
    ax.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--',
            label='Random classifier')

    # Formatting
    ax.set_xlim([0.0, 1.0])
    ax.set_ylim([0.0, 1.05])
    ax.set_xlabel('False Positive Rate', fontsize=12)
    ax.set_ylabel('True Positive Rate', fontsize=12)
    ax.set_title(title, fontsize=14, fontweight='bold')
    ax.legend(loc="lower right", fontsize=11)
    ax.grid(True, alpha=0.3)
    ax.set_aspect('equal')

    # Save plot
    os.makedirs(save_dir, exist_ok=True)
    plot_path = os.path.join(save_dir, f"{filename}.png")
    fig.savefig(plot_path, dpi=300, bbox_inches='tight')
    plt.close(fig)

    # Save data
    data_path = os.path.join(save_dir, f"{filename}_data.npz")
    np.savez(
        data_path,
        y_true=y_true,
        y_proba=y_proba,
        fpr=fpr,
        tpr=tpr,
        thresholds=thresholds,
        auroc=auroc
    )

    print(f"Saved ROC curve plot to {plot_path}")
    print(f"Saved ROC data to {data_path}")

    return auroc, fpr, tpr, thresholds