"""
Dimensionality reduction and feature extraction for primordial features analysis.

This module provides tools to reduce the N-dimensional bin parameter space
to a smaller number of interpretable components, revealing the dominant modes
of variation in $\\delta(k)$.

Supports multiple backends:
- NumPy/sklearn (default): CPU-based computation using sklearn's PCA
- JAX: GPU-accelerated, JIT-compiled, autodiff-compatible

Use the `backend` parameter to select:
    >>> results = perform_pca(chains, backend="numpy")  # Default
    >>> results = perform_pca(chains, backend="jax")    # GPU-accelerated
"""

import numpy as np
import matplotlib.pyplot as plt
from sklearn.decomposition import PCA, FastICA
from sklearn.preprocessing import StandardScaler
from typing import Dict, List, Optional, Tuple, Union, Any
from dataclasses import dataclass, field
from rich.panel import Panel
from rich.table import Table
from .console import console
from .io import print_success, print_warning

# Backend imports
from .backends.numpy.pca import get_numpy_pca_backend, NumPyPCABackend

# Check for JAX availability
try:
    from .backends.jax.pca_jax import (
        get_jax_pca_backend,
        JAXPCABackend,
        is_jax_available,
    )

    HAS_JAX = is_jax_available()
except ImportError:
    HAS_JAX = False
    get_jax_pca_backend = None
    JAXPCABackend = None


def _get_pca_backend(backend: Optional[str] = None) -> Any:
    """
    Get PCA backend instance based on specification.

    Args:
        backend: Backend name ("numpy", "jax") or None for auto-detection

    Returns:
        Backend instance (NumPyPCABackend or JAXPCABackend)

    Raises:
        ValueError: If invalid backend specified
        ImportError: If JAX backend requested but not available
    """
    if backend is None:
        # Auto-detect: prefer JAX if available
        if HAS_JAX:
            return get_jax_pca_backend()
        return get_numpy_pca_backend()

    backend = backend.lower()

    if backend == "numpy":
        return get_numpy_pca_backend()
    elif backend == "jax":
        if not HAS_JAX:
            raise ImportError(
                "JAX backend requested but not available. "
                "Install with: pip install jax jaxlib"
            )
        return get_jax_pca_backend()
    else:
        raise ValueError(f"Unknown backend: {backend}. Use 'numpy' or 'jax'.")


@dataclass
class PCAResults:
    """Results from PCA analysis."""

    n_components: int
    explained_variance_ratio: np.ndarray
    cumulative_variance: np.ndarray
    components: np.ndarray  # Principal components (eigenvectors)
    transformed_data: np.ndarray  # Projected data in PC space
    pca_model: Optional[PCA]  # sklearn PCA model (None for JAX backend)
    scaler: StandardScaler
    effective_dim: int  # Number of PCs explaining 95% variance
    dataset_labels: Optional[List[str]] = None
    backend: str = "numpy"  # Backend used ("numpy" or "jax")
    mean: Optional[np.ndarray] = None  # Data mean (for backend-agnostic reconstruction)


def collect_delta_samples(
    chains_dict: Dict, nbins: int = 20, param_pattern: str = "delta_{i}"
) -> Tuple[np.ndarray, List[str]]:
    """
    Collect all $\\delta$ samples from all chains into a single array.

    Args:
        chains_dict: Dictionary mapping dataset labels to MCMC chains,
                     or a single chain object (will be wrapped in dict)
        nbins: Number of bins (default: 20)
        param_pattern: Parameter name pattern (default: "delta_{i}")

    Returns:
        X: Array of shape (n_total_samples, nbins) with all $\\delta$ values
        labels: List of dataset labels for each sample
    """
    # Handle single chain object (wrap in dict)
    if not isinstance(chains_dict, dict):
        chains_dict = {"chain": chains_dict}

    all_deltas = []
    labels = []

    for label, chain in chains_dict.items():
        # Extract delta parameters for this chain
        delta_matrix = np.array(
            [chain[param_pattern.format(i=i)] for i in range(1, nbins + 1)]
        ).T

        all_deltas.append(delta_matrix)
        labels.extend([label] * len(delta_matrix))

    X = np.vstack(all_deltas)
    return X, labels


def _perform_pca_on_data(
    X: np.ndarray,
    labels: List[str],
    n_components: int,
    nbins: int,
    dataset_name: Optional[str] = None,
    show_input_panel: bool = True,
    show_results: bool = True,
    backend: Optional[str] = None,
) -> PCAResults:
    """
    Internal function to perform PCA on a data matrix.
    Perform Principal Component Analysis on $\\delta$ parameters.

    This identifies the dominant modes of variation in the primordial power
    spectrum deviations across all datasets.

    Args:
        X: Data matrix (n_samples, nbins)
        labels: List of labels for each sample
        n_components: Number of PCs to compute
        nbins: Number of bins
        dataset_name: Optional name for console output
        show_input_panel: Whether to show input data panel
        show_results: Whether to print results table
        backend: Backend to use ("numpy", "jax", or None for auto)

    Returns:
        PCAResults object
    """
    # Standardize the data (critical for PCA!)
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)

    # Get backend
    pca_backend = _get_pca_backend(backend)
    backend_name = pca_backend.backend_name

    # Fit and transform using backend
    X_pca, fit_result = pca_backend.fit_transform(X_scaled, n_components)

    # Extract results
    explained_var_ratio = fit_result["explained_variance_ratio"]
    cumulative_var = np.cumsum(explained_var_ratio)
    components = fit_result["components"]
    mean = fit_result["mean"]

    # Find intrinsic dimensionality (95% variance threshold)
    effective_dim = int(np.argmax(cumulative_var >= 0.95) + 1)

    title_suffix = f" ({dataset_name})" if dataset_name else ""
    backend_suffix = f" [{backend_name}]" if backend_name != "numpy" else ""

    # Display input panel if requested
    if show_input_panel:
        input_panel = Panel(
            f"[key]Samples:[/key] [value]{X.shape[0]:,}[/value]\n"
            f"[key]Dimensionality:[/key] [value]{X.shape[1]}[/value] bins\n"
            f"[key]Backend:[/key] [value]{backend_name}[/value]",
            title=f"[header]PCA Input Data{title_suffix}{backend_suffix}[/header]",
            border_style="cyan",
        )
        console.print(input_panel)

    # Display results if requested
    if show_results:
        # Results table
        results_table = Table(
            title=f"PCA Results{title_suffix}{backend_suffix}", show_header=True
        )
        results_table.add_column("PC", justify="right", style="cyan", width=4)
        results_table.add_column("Variance", justify="right", style="yellow", width=10)
        results_table.add_column("Cumulative", justify="right", style="green", width=12)

        for i in range(min(10, n_components)):
            pc_style = "bold" if i < effective_dim else "dim"
            results_table.add_row(
                f"[{pc_style}]{i + 1}[/{pc_style}]",
                f"[{pc_style}]{explained_var_ratio[i]:.2%}[/{pc_style}]",
                f"[{pc_style}]{cumulative_var[i]:.2%}[/{pc_style}]",
            )

        console.print(results_table)
        print_success(
            f"Effective dimensionality: {effective_dim} PCs (explain 95% variance)"
        )
        console.print(
            f"[dim]Top 5 PCs explain {cumulative_var[4]:.1%} of total variance[/dim]\n"
        )

    # Create sklearn PCA model for backward compatibility (numpy backend only)
    pca_model = None
    if backend_name == "numpy":
        # Create a sklearn PCA with pre-computed values for compatibility
        pca_model = PCA(n_components=n_components)
        pca_model.fit(X_scaled)

    return PCAResults(
        n_components=n_components,
        explained_variance_ratio=explained_var_ratio,
        cumulative_variance=cumulative_var,
        components=components,
        transformed_data=X_pca,
        pca_model=pca_model,
        scaler=scaler,
        effective_dim=effective_dim,
        dataset_labels=labels,
        backend=backend_name,
        mean=mean,
    )


def perform_pca(
    chains_dict: Dict,
    nbins: int = 20,
    n_components: Optional[int] = None,
    param_pattern: str = "delta_{i}",
    mode: str = "pooled",
    backend: Optional[str] = None,
) -> Union[PCAResults, Dict[str, PCAResults]]:
    """
    Perform Principal Component Analysis on delta parameters.

    This identifies the dominant modes of variation in the primordial power
    spectrum deviations across datasets.

    Args:
        chains_dict: Dictionary mapping dataset labels to MCMC chains,
                     or a single chain object (will be wrapped in dict)
        nbins: Number of bins (default: 20)
        n_components: Number of components to compute (default: nbins)
        param_pattern: Parameter name pattern (default: "delta_{i}")
        mode: Analysis mode (default: "pooled")
              - "pooled": Perform PCA on pooled samples from all chains
              - "individual": Perform PCA separately on each chain
        backend: Computation backend (default: None for auto-detection)
              - "numpy": Use NumPy/sklearn backend (CPU)
              - "jax": Use JAX backend (GPU-accelerated, autodiff-compatible)
              - None: Auto-detect (prefer JAX if available)

    Returns:
        If mode="pooled":
            PCAResults object containing all analysis results
        If mode="individual":
            Dictionary mapping chain labels to PCAResults objects

    Example:
        >>> # Pooled mode (default) - backward compatible
        >>> results = perform_pca(chains, nbins=20)
        >>> print(f"Effective dimensionality: {results.effective_dim}")
        >>> print(f"Top 5 PCs explain {results.cumulative_variance[4]:.1%} of variance")

        >>> # Individual mode - analyze each chain separately
        >>> results_dict = perform_pca(chains, nbins=20, mode="individual")
        >>> for label, result in results_dict.items():
        ...     print(f"{label}: {result.effective_dim} effective dims")

        >>> # Use JAX backend for GPU acceleration
        >>> results = perform_pca(chains, nbins=20, backend="jax")
    """
    # Validate mode
    if mode not in ["pooled", "individual"]:
        raise ValueError(f"mode must be 'pooled' or 'individual', got '{mode}'")

    # Handle single chain object (wrap in dict)
    if not isinstance(chains_dict, dict):
        chains_dict = {"chain": chains_dict}

    # Set default number of components
    if n_components is None:
        n_components = nbins

    # MODE: POOLED (default, backward compatible)
    if mode == "pooled":
        # Collect all delta samples
        X, labels = collect_delta_samples(chains_dict, nbins, param_pattern)

        # Get backend name for display
        pca_backend = _get_pca_backend(backend)
        backend_name = pca_backend.backend_name
        backend_suffix = f" [{backend_name}]" if backend_name != "numpy" else ""

        # Display input info
        input_panel = Panel(
            f"[key]Samples:[/key] [value]{X.shape[0]:,}[/value] from "
            f"[value]{len(chains_dict)}[/value] dataset(s)\n"
            f"[key]Dimensionality:[/key] [value]{X.shape[1]}[/value] bins\n"
            f"[key]Backend:[/key] [value]{backend_name}[/value]",
            title=f"[header]PCA Input Data (Pooled){backend_suffix}[/header]",
            border_style="cyan",
        )
        console.print(input_panel)

        # Perform PCA on pooled data (input panel already shown above)
        return _perform_pca_on_data(
            X=X,
            labels=labels,
            n_components=n_components,
            nbins=nbins,
            dataset_name=None,
            show_input_panel=False,
            show_results=True,
            backend=backend,
        )

    # MODE: INDIVIDUAL
    else:  # mode == "individual"
        results_dict = {}

        # Get backend name for display
        pca_backend = _get_pca_backend(backend)
        backend_name = pca_backend.backend_name
        backend_suffix = f" [{backend_name}]" if backend_name != "numpy" else ""

        console.print(
            Panel(
                f"[key]Analyzing:[/key] [value]{len(chains_dict)}[/value] "
                f"dataset(s) individually\n"
                f"[key]Backend:[/key] [value]{backend_name}[/value]",
                title=f"[header]PCA Individual Mode{backend_suffix}[/header]",
                border_style="magenta",
            )
        )
        console.print()

        for label, chain in chains_dict.items():
            # Extract delta parameters for this chain
            delta_matrix = np.array(
                [chain[param_pattern.format(i=i)] for i in range(1, nbins + 1)]
            ).T

            # Create labels for this chain
            chain_labels = [label] * len(delta_matrix)

            # Perform PCA on this chain
            result = _perform_pca_on_data(
                X=delta_matrix,
                labels=chain_labels,
                n_components=n_components,
                nbins=nbins,
                dataset_name=label,
                show_input_panel=True,
                show_results=True,
                backend=backend,
            )

            results_dict[label] = result

        # Print summary across all chains
        summary_table = Table(
            title="Summary Across Datasets",
            show_header=True,
        )
        summary_table.add_column("Dataset", style="cyan")
        summary_table.add_column("Samples", justify="right", style="yellow")
        summary_table.add_column("d_eff", justify="right", style="green")
        summary_table.add_column("Top PC Var", justify="right", style="blue")

        for label, result in results_dict.items():
            summary_table.add_row(
                label,
                f"{result.transformed_data.shape[0]:,}",
                str(result.effective_dim),
                f"{result.explained_variance_ratio[0]:.2%}",
            )

        console.print()
        console.print(summary_table)
        console.print()

        return results_dict


def perform_ica(
    chains_dict: Dict,
    nbins: int = 20,
    n_components: int = 10,
    param_pattern: str = "delta_{i}",
    random_state: int = 42,
) -> Tuple[FastICA, np.ndarray, np.ndarray]:
    """
    Perform Independent Component Analysis on $\\delta$ parameters.

    ICA finds statistically independent patterns, which can be better than PCA
    for identifying localized features or non-Gaussian structures.

    Args:
        chains_dict: Dictionary mapping dataset labels to MCMC chains,
                     or a single chain object (will be wrapped in dict)
        nbins: Number of bins (default: 20)
        n_components: Number of ICs to extract (default: 10)
        param_pattern: Parameter name pattern (default: "delta_{i}")
        random_state: Random seed for reproducibility (default: 42)

    Returns:
        ica_model: Fitted FastICA object
        X_ica: Transformed data in IC space
        components: Independent components (mixing matrix)

    Example:

        >>> ica, X_ica, components = perform_ica(chains, nbins=20, n_components=10)
        >>> # Plot components similar to PCA
    """
    # Handle single chain object (wrap in dict)
    if not isinstance(chains_dict, dict):
        chains_dict = {"chain": chains_dict}

    # Collect samples
    X, labels = collect_delta_samples(chains_dict, nbins, param_pattern)

    from rich.status import Status

    with Status("[cyan]Performing ICA decomposition...", console=console) as status:
        # Standardize first (recommended for ICA)
        scaler = StandardScaler()
        X_scaled = scaler.fit_transform(X)

        # Fit ICA
        ica = FastICA(
            n_components=n_components, random_state=random_state, max_iter=500, tol=1e-4
        )
        X_ica = ica.fit_transform(X_scaled)

        # Get components (sources)
        components = ica.components_

        status.update(f"[cyan]ICA converged after {ica.n_iter_} iterations")

    print_success(f"ICA converged after {ica.n_iter_} iterations")

    return ica, X_ica, components


def compute_reconstruction_error(
    results: PCAResults, X_original: np.ndarray, n_components: int
) -> float:
    """
    Compute reconstruction error using only n_components PCs.

    This quantifies how much information is lost by using fewer components.

    Args:
        results: PCAResults from perform_pca()
        X_original: Original (unstandardized) data
        n_components: Number of PCs to use for reconstruction

    Returns:
        rmse: Root mean squared error of reconstruction
    """
    # Standardize
    X_scaled = results.scaler.transform(X_original)

    # Project to PC space and back
    X_proj = results.pca_model.transform(X_scaled)[:, :n_components]
    X_reconstructed = results.pca_model.inverse_transform(
        np.column_stack(
            [X_proj, np.zeros((X_proj.shape[0], results.n_components - n_components))]
        )
    )

    # Compute RMSE
    rmse = np.sqrt(np.mean((X_scaled - X_reconstructed) ** 2))

    return rmse


def RMSE_vs_n_components(chains_dict, nbins=20, max_components=None):
    """
    Compute reconstruction error (RMSE) as a function of number of PCs.

    Returns reconstruction error per sample as function of number of PCs.
    """
    # Perform PCA to get the model
    pca_result = perform_pca(chains_dict, nbins=nbins)

    # Collect all delta samples
    X, labels = collect_delta_samples(chains_dict, nbins=nbins)

    max_components = max_components or nbins
    error_values = []

    # IMPORTANT: PCA uses StandardScaler, so we need to work in scaled space
    # then transform back to original space for interpretable errors

    for k in range(1, max_components + 1):
        # Reconstruct in scaled space using PCA model
        # PCA does: X_scaled = scaler.transform(X)
        #           X_pca = X_scaled @ components.T
        # Inverse: X_scaled_recon = X_pca[:, :k] @ components[:k, :]
        #          X_recon = scaler.inverse_transform(X_scaled_recon)

        X_scaled = pca_result.scaler.transform(X)
        X_pca_k = pca_result.transformed_data[:, :k]  # First k PC scores

        # Reconstruct in scaled space
        X_scaled_recon = X_pca_k @ pca_result.components[:k, :]

        # Transform back to original space
        X_recon = pca_result.scaler.inverse_transform(X_scaled_recon)

        # Compute reconstruction error (RMSE per sample)
        residuals = X - X_recon
        mse_per_sample = np.mean(residuals**2, axis=1)  # Mean over bins
        rmse_per_sample = np.sqrt(mse_per_sample)

        # Also compute in terms of variance explained
        total_variance = np.var(X)
        residual_variance = np.var(residuals)
        variance_explained = 1 - (residual_variance / total_variance)

        error_values.append(
            {
                "n_components": k,
                "mean_rmse": np.mean(rmse_per_sample),
                "median_rmse": np.median(rmse_per_sample),
                "std_rmse": np.std(rmse_per_sample),
                "percentile_95": np.percentile(rmse_per_sample, 95),
                "variance_explained": variance_explained,
                "total_mse": np.mean(mse_per_sample),
            }
        )

    return error_values, pca_result


def reconstruct_delta_from_pcs(
    results: PCAResults,
    pc_indices: List[int],
    sample_indices: Optional[Union[int, List[int]]] = None,
    return_mean: bool = False,
) -> Dict[str, Any]:
    """
    Reconstruct $\\delta(k)$ using only specified principal components.

    Allows selective reconstruction to study impact of individual PCs or groups.
    For example, reconstruct using PCs 2-8 to see the effect of dropping PC1.

    Backend-agnostic: works with both NumPy and JAX PCA results.

    Args:
        results: PCAResults object from perform_pca()
        pc_indices: List of PC indices to use (1-indexed, e.g., [2, 3, 4, 5, 6, 7, 8])
                   Uses 1-indexing to match standard PC naming (PC1, PC2, ...)
        sample_indices: Optional sample index/indices to reconstruct:
                       - None: reconstruct all samples
                       - int: single sample
                       - List[int]: multiple specific samples
        return_mean: If True, average reconstruction across selected samples

    Returns:
        Dictionary containing:
        - 'delta': Reconstructed $\\delta(k)$ values, shape depends on inputs
        - 'pc_indices': PC indices used (1-indexed)
        - 'n_pcs_used': Number of PCs used
        - 'variance_explained': Cumulative variance explained by selected PCs
        - 'n_samples': Number of samples reconstructed
        - 'shape': Shape of delta array

    Examples:
        >>> # Reconstruct using only PCs 2-8 (drop PC1)
        >>> result = reconstruct_delta_from_pcs(pca_results, pc_indices=[2,3,4,5,6,7,8])
        >>> delta_partial = result['delta']  # Shape: (n_samples, nbins)

        >>> # See impact of PC1 alone
        >>> result = reconstruct_delta_from_pcs(pca_results, pc_indices=[1])
        >>> delta_pc1_only = result['delta']

        >>> # Reconstruct single sample using PCs 1-5
        >>> result = reconstruct_delta_from_pcs(pca_results, [1,2,3,4,5], sample_indices=42)
        >>> delta_sample = result['delta']  # Shape: (nbins,)

        >>> # Get mean reconstruction across all samples (for plotting)
        >>> result = reconstruct_delta_from_pcs(pca_results, [2,3,4,5,6,7,8], return_mean=True)
        >>> delta_mean = result['delta']  # Shape: (nbins,)
    """
    # Validate PC indices (1-indexed)
    pc_indices = np.asarray(pc_indices)
    if np.any(pc_indices < 1) or np.any(pc_indices > results.n_components):
        raise ValueError(
            f"PC indices must be between 1 and {results.n_components}, "
            f"got {pc_indices.tolist()}"
        )

    # Convert to 0-indexed for array operations
    pc_indices_0 = pc_indices - 1

    # Get PC scores for selected components
    pc_scores_selected = results.transformed_data[:, pc_indices_0]

    # Get corresponding components
    components_selected = results.components[pc_indices_0, :]

    # Reconstruct in scaled space
    # X_scaled = PC_scores @ Components
    X_scaled_recon = pc_scores_selected @ components_selected

    # Transform back to original space
    # Use scaler.inverse_transform for backend compatibility
    delta_reconstructed = results.scaler.inverse_transform(X_scaled_recon)

    # Handle sample selection
    if sample_indices is not None:
        if isinstance(sample_indices, int):
            delta_reconstructed = delta_reconstructed[sample_indices, :]
        else:
            delta_reconstructed = delta_reconstructed[sample_indices, :]

    # Compute mean if requested
    if return_mean:
        delta_reconstructed = np.mean(delta_reconstructed, axis=0)

    # Compute variance explained by selected PCs
    variance_explained = np.sum(results.explained_variance_ratio[pc_indices_0])

    # Build result dictionary
    result = {
        "delta": delta_reconstructed,
        "pc_indices": pc_indices.tolist(),
        "n_pcs_used": len(pc_indices),
        "variance_explained": float(variance_explained),
        "n_samples": delta_reconstructed.shape[0]
        if delta_reconstructed.ndim > 1
        else 1,
        "shape": delta_reconstructed.shape,
    }

    return result


def compare_pc_reconstructions(
    results: PCAResults,
    pc_sets: Dict[str, List[int]],
    k_values: Optional[np.ndarray] = None,
    sample_indices: Optional[Union[int, List[int]]] = None,
    figsize: Tuple[float, float] = (12, 6),
    title: Optional[str] = None,
) -> plt.Figure:
    """
    Compare $\\delta(k)$ reconstructions using different PC subsets.

    Useful for visualizing the impact of specific PCs or understanding
    how reconstruction quality changes with different PC selections.

    Args:
        results: PCAResults object from perform_pca()
        pc_sets: Dictionary mapping labels to PC index lists
                Example: {
                    'Full (1-10)': [1,2,3,4,5,6,7,8,9,10],
                    'Without PC1': [2,3,4,5,6,7,8,9,10],
                    'PC1 only': [1],
                    'PCs 2-5': [2,3,4,5]
                }
        k_values: Optional k-values for x-axis (Mpc$^{-1}$)
                 If None, uses bin indices
        sample_indices: Which samples to plot (None = mean over all)
        figsize: Figure size (width, height)
        title: Optional plot title

    Returns:
        Figure object

    Example:
        >>> pc_sets = {
        ...     'Full (PCs 1-10)': list(range(1, 11)),
        ...     'Drop PC1 (PCs 2-10)': list(range(2, 11)),
        ...     'PC1 only': [1],
        ...     'PCs 2-8': [2,3,4,5,6,7,8]
        ... }
        >>> fig = compare_pc_reconstructions(pca_results, pc_sets, k_values=k)
        >>> plt.show()
    """
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=figsize)

    # Use bin indices if k_values not provided
    nbins = results.components.shape[1]
    x_values = k_values if k_values is not None else np.arange(1, nbins + 1)
    x_label = r"$k$ [Mpc$^{-1}$]" if k_values is not None else "Bin Index"

    # Colors for different PC sets
    colors = plt.cm.tab10(np.linspace(0, 1, len(pc_sets)))

    # Left panel: Reconstructions
    for (label, pc_list), color in zip(pc_sets.items(), colors):
        recon = reconstruct_delta_from_pcs(
            results, pc_list, sample_indices=sample_indices, return_mean=True
        )
        delta = recon["delta"]
        var_exp = recon["variance_explained"]

        label_with_var = f"{label} ({var_exp:.1%} var)"
        ax1.plot(x_values, delta, label=label_with_var, color=color, linewidth=2)

    ax1.axhline(0, color="k", linestyle="--", alpha=0.3, linewidth=1)
    ax1.set_xlabel(x_label)
    ax1.set_ylabel(r"$\delta(k)$")
    ax1.legend(fontsize=9, framealpha=0.9)
    ax1.grid(True, alpha=0.3)
    if k_values is not None:
        ax1.set_xscale("log")

    # Right panel: Differences from first reconstruction (typically "full")
    reference_label = list(pc_sets.keys())[0]
    reference_recon = reconstruct_delta_from_pcs(
        results,
        pc_sets[reference_label],
        sample_indices=sample_indices,
        return_mean=True,
    )
    reference_delta = reference_recon["delta"]

    for (label, pc_list), color in zip(list(pc_sets.items())[1:], colors[1:]):
        recon = reconstruct_delta_from_pcs(
            results, pc_list, sample_indices=sample_indices, return_mean=True
        )
        delta = recon["delta"]
        difference = delta - reference_delta

        ax2.plot(
            x_values,
            difference,
            label=f"{label} - {reference_label}",
            color=color,
            linewidth=2,
        )

    ax2.axhline(0, color="k", linestyle="--", alpha=0.3, linewidth=1)
    ax2.set_xlabel(x_label)
    ax2.set_ylabel(r"$\Delta\delta(k)$ from reference")
    ax2.legend(fontsize=9, framealpha=0.9)
    ax2.grid(True, alpha=0.3)
    if k_values is not None:
        ax2.set_xscale("log")

    if title:
        fig.suptitle(title, fontsize=14, fontweight="bold")

    plt.tight_layout()
    return fig


def analyze_bin_correlations(chains_dict: Dict, nbins: int = 20) -> np.ndarray:
    """
    Compute correlation matrix between bins.

    This shows which bins are correlated (likely due to smoothness constraints
    or cosmic variance).

    Args:
        chains_dict: Dictionary mapping dataset labels to MCMC chains,
                     or a single chain object (will be wrapped in dict)
        nbins: Number of bins (default: 20)

    Returns:
        corr_matrix: Correlation matrix (nbins x nbins)
    """
    # Handle single chain object (wrap in dict)
    if not isinstance(chains_dict, dict):
        chains_dict = {"chain": chains_dict}

    X, _ = collect_delta_samples(chains_dict, nbins)

    # Compute correlation matrix
    corr_matrix = np.corrcoef(X.T)

    return corr_matrix


def variance_decomposition(pca_pooled, N_pcs=10, chains_dict=None):
    from scipy.stats import f_oneway

    # Get transformed data and labels
    transformed = pca_pooled.transformed_data
    labels_array = np.array(pca_pooled.dataset_labels)

    # Compute average effective sample size across chains
    if chains_dict is not None:
        if not isinstance(chains_dict, dict):
            chains_dict = {"chain": chains_dict}
        n_eff_list = [chain.getEffectiveSamples() for chain in chains_dict.values()]
        avg_n_eff = np.mean(n_eff_list)
    else:
        # Fallback: use total number of samples
        avg_n_eff = len(transformed)

    # Get explained variance ratio from PCA model (fraction of total variance)
    explained_variance_ratio = pca_pooled.pca_model.explained_variance_ratio_

    variance_between = []
    variance_within = []
    variance_total = []
    f_statistics = []
    snr_values = []
    p_values = []

    n_pcs = min(N_pcs, pca_pooled.n_components)

    # Create variance decomposition table
    decomp_table = Table(title="Variance Decomposition", show_header=True)
    decomp_table.add_column("PC", justify="right", style="cyan", width=4)
    decomp_table.add_column("Var(Between)", justify="right", width=13)
    decomp_table.add_column("Var(Within)", justify="right", width=12)
    decomp_table.add_column("F-stat", justify="right", style="yellow", width=10)
    decomp_table.add_column("SNR", justify="right", style="green", width=10)
    decomp_table.add_column("Interpretation", style="dim")

    for i in range(n_pcs):
        pc_scores = transformed[:, i]

        # Between-dataset variance
        dataset_means = np.array(
            [
                pc_scores[labels_array == l].mean()
                for l in np.unique(pca_pooled.dataset_labels)
            ]
        )
        var_between = np.var(dataset_means, ddof=1) if len(dataset_means) > 1 else 0

        # Within-dataset variance
        var_within = np.mean(
            [
                pc_scores[labels_array == l].var(ddof=1)
                for l in np.unique(pca_pooled.dataset_labels)
            ]
        )

        # Total variance
        var_total = np.var(pc_scores, ddof=1)

        # F-statistic (ratio of between to within variance)
        f_stat = var_between / var_within if var_within > 0 else 0

        # Signal-to-noise ratio: fraction of variance explained, scaled by measurement quality
        # Signal: explained variance ratio (0-1, fraction of total variance)
        # Statistical quality factor: sqrt(N_eff / 100) normalized to ~1 for N_eff=100
        signal = explained_variance_ratio[i]
        quality_factor = np.sqrt(avg_n_eff / 100.0) if avg_n_eff > 0 else 1.0
        snr = signal * quality_factor

        # ANOVA p-value to test if means differ across datasets
        dataset_groups = [
            pc_scores[labels_array == l] for l in np.unique(pca_pooled.dataset_labels)
        ]
        _, p_value = f_oneway(*dataset_groups)

        variance_between.append(var_between)
        variance_within.append(var_within)
        variance_total.append(var_total)
        f_statistics.append(f_stat)
        snr_values.append(snr)
        p_values.append(p_value)

        # Interpretation
        if f_stat > 1.0:
            interpretation = "Dataset separation dominates"
            row_style = "green"
        elif f_stat > 0.1:
            interpretation = "Mixed (both sources contribute)"
            row_style = "yellow"
        else:
            interpretation = "Within-dataset variance dominates"
            row_style = "dim"

        decomp_table.add_row(
            f"[{row_style}]{i + 1}[/{row_style}]",
            f"[{row_style}]{var_between:.4f}[/{row_style}]",
            f"[{row_style}]{var_within:.4f}[/{row_style}]",
            f"[{row_style}]{f_stat:.3f}[/{row_style}]",
            f"[{row_style}]{snr:.2f}[/{row_style}]",
            f"[{row_style}]{interpretation}[/{row_style}]",
        )

    console.print(decomp_table)

    # Print interpretation guide
    console.print("\n[header]Interpretation Guide:[/header]")
    console.print(
        "  • [green]High F-stat (>1)[/green]: PC captures differences between datasets"
    )
    console.print("  • [yellow]Mid F-stat (0.1-1)[/yellow]: PC captures both sources")
    console.print(
        "  • [dim]Low F-stat (<0.1)[/dim]: PC captures shared variation within datasets"
    )
    console.print(
        "  • [green]High SNR[/green]: PC is well-measured (signal >> measurement noise)"
    )
    console.print("  • [dim]Low SNR[/dim]: PC is dominated by MCMC sampling noise\n")
    return (
        variance_between,
        variance_within,
        variance_total,
        f_statistics,
        snr_values,
        p_values,
    )
