"""Cross-experiment correlation analysis functions.

This module provides functions for analyzing correlations between traits
measured in different experimental modalities (e.g., cylinder vs turface).
"""

from __future__ import annotations

from pathlib import Path
from typing import Dict, List, Tuple, Optional
import warnings

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from scipy import stats
from scipy.cluster.hierarchy import linkage as hierarchical_linkage, fcluster
from scipy.spatial.distance import squareform

from sleap_roots_analyze.data_utils import sanitize_single_trait_name


def _calculate_correlations(
    values1: np.ndarray, values2: np.ndarray
) -> Tuple[float, float, float, float]:
    """Helper function to calculate both Pearson and Spearman correlations.

    Args:
        values1: First set of values
        values2: Second set of values

    Returns:
        Tuple of (pearson_r, pearson_p, spearman_r, spearman_p)
    """
    pearson_r, pearson_p = stats.pearsonr(values1, values2)
    spearman_r, spearman_p = stats.spearmanr(values1, values2)
    return pearson_r, pearson_p, spearman_r, spearman_p


def _prepare_aligned_values(
    exp1_data: pd.DataFrame,
    exp2_data: pd.DataFrame,
    exp1_col: str,
    exp2_col: str,
    min_samples: int = 0,
) -> Tuple[np.ndarray, np.ndarray, List[str]]:
    """Helper to align values from two experiments.

    Args:
        exp1_data: Data from experiment 1 (with index as genotypes)
        exp2_data: Data from experiment 2 (with index as genotypes)
        exp1_col: Column name from exp1
        exp2_col: Column name from exp2
        min_samples: Minimum samples required per genotype

    Returns:
        Tuple of (values1, values2, valid_genotypes) with NaN removed
    """
    # Find common genotypes
    valid_genotypes = list(set(exp1_data.index) & set(exp2_data.index))

    # Filter by sample count if needed
    if (
        min_samples > 0
        and "n_samples" in exp1_data.columns
        and "n_samples" in exp2_data.columns
    ):
        valid_genotypes = [
            g
            for g in valid_genotypes
            if exp1_data.loc[g, "n_samples"] >= min_samples
            and exp2_data.loc[g, "n_samples"] >= min_samples
        ]

    if len(valid_genotypes) < 3:
        return np.array([]), np.array([]), []

    # Get aligned values
    values1 = exp1_data.loc[valid_genotypes, exp1_col].values
    values2 = exp2_data.loc[valid_genotypes, exp2_col].values

    # Remove NaN pairs
    valid_mask = ~(np.isnan(values1) | np.isnan(values2))
    if valid_mask.sum() < 3:
        return np.array([]), np.array([]), []

    return (
        values1[valid_mask],
        values2[valid_mask],
        [g for i, g in enumerate(valid_genotypes) if valid_mask[i]],
    )


def calculate_correlations(
    x: np.ndarray, y: np.ndarray, method: str = "spearman"
) -> Tuple[float, float]:
    """Calculate correlation between two arrays using specified method.

    This function provides a unified interface for computing correlations with
    proper handling of edge cases (NaN values, insufficient data, etc.).

    Args:
        x: First array of values
        y: Second array of values
        method: Correlation method - one of 'spearman', 'pearson', or 'kendall'

    Returns:
        Tuple of (correlation_coefficient, p_value)
        Returns (NaN, NaN) if:
            - Arrays have different lengths
            - Fewer than 3 valid paired samples after NaN removal
            - Either array has zero variance (all constant values)

    Raises:
        ValueError: If method is not one of 'spearman', 'pearson', or 'kendall'
        ValueError: If x and y have different lengths

    Examples:
        >>> x = np.array([1, 2, 3, 4, 5])
        >>> y = np.array([2, 4, 6, 8, 10])
        >>> r, p = calculate_correlations(x, y, method='pearson')
        >>> r  # Should be 1.0 for perfect linear relationship
        1.0
    """
    # Validate method
    valid_methods = ["spearman", "pearson", "kendall"]
    if method not in valid_methods:
        raise ValueError(f"method must be one of {valid_methods}, got '{method}'")

    # Validate input lengths
    if len(x) != len(y):
        raise ValueError(
            f"x and y must have the same length, got {len(x)} and {len(y)}"
        )

    # Remove NaN pairs
    valid_mask = ~(np.isnan(x) | np.isnan(y))
    x_clean = x[valid_mask]
    y_clean = y[valid_mask]

    # Check for insufficient data
    if len(x_clean) < 3:
        return np.nan, np.nan

    # Check for constant values (zero variance)
    if np.std(x_clean) == 0 or np.std(y_clean) == 0:
        return np.nan, np.nan

    # Calculate correlation based on method
    if method == "spearman":
        r, p = stats.spearmanr(x_clean, y_clean)
    elif method == "pearson":
        r, p = stats.pearsonr(x_clean, y_clean)
    elif method == "kendall":
        r, p = stats.kendalltau(x_clean, y_clean)

    return r, p


def calculate_correlation_ci(
    r: float, n: int, confidence_level: float = 0.95
) -> Tuple[float, float]:
    """Calculate confidence interval for a correlation coefficient using Fisher z-transformation.

    The Fisher z-transformation converts a correlation coefficient to a normally
    distributed variable, enabling straightforward CI calculation:

    1. Transform: z = arctanh(r) = 0.5 × ln((1+r)/(1-r))
    2. Standard error: SE_z = 1 / √(n-3)
    3. CI on z-scale: z ± z_{α/2} × SE_z
    4. Back-transform: r = tanh(z)

    This method is exact for Pearson correlation under bivariate normality.
    For Spearman correlations, this Fisher z-based interval is generally
    considered accurate for sample sizes n >= 10; for 4 <= n < 10, the interval
    is still computed but should be interpreted with caution due to reduced
    approximation accuracy.

    Note:
        For DataFrame-based CI calculation on multiple correlations, see
        :func:`calculate_correlation_confidence_intervals`.

    Args:
        r: Correlation coefficient (Pearson or Spearman), must be in [-1, 1].
            Values outside this range will raise a ValueError. NaN is allowed
            and returns (NaN, NaN).
        n: Sample size (number of paired observations). For Spearman
            correlations, n >= 10 is recommended for accurate approximation;
            for 4 <= n < 10, the CI is approximate and may be less reliable.
        confidence_level: Confidence level for the interval, default 0.95.
            Must be in (0, 1) exclusive range.

    Returns:
        Tuple of (ci_low, ci_high) representing the confidence interval bounds.
        Returns (NaN, NaN) if:
            - r is NaN
            - n < 4 (variance formula undefined with n-3 in denominator)
        Returns (r, r) if r = ±1.0 (point mass at boundary)

    Raises:
        ValueError: If r is not in [-1, 1] (and not NaN), if
            confidence_level is not in (0, 1), or if n <= 0.

    References:
        Fisher, R.A. (1921). On the "probable error" of a coefficient of
        correlation deduced from a small sample. Metron, 1, 3-32.

    Examples:
        >>> ci_low, ci_high = calculate_correlation_ci(r=0.5, n=20, confidence_level=0.95)
        >>> print(f"95% CI: ({ci_low:.3f}, {ci_high:.3f})")
        95% CI: (0.058, 0.776)
    """
    # Handle NaN input first (before validation)
    if np.isnan(r):
        return np.nan, np.nan

    # Validate r is in [-1, 1]
    if abs(r) > 1:
        raise ValueError(f"Correlation coefficient r must be in [-1, 1], got {r}")

    # Validate confidence_level is in (0, 1)
    if not (0 < confidence_level < 1):
        raise ValueError(
            f"confidence_level must be in (0, 1) exclusive range, got {confidence_level}"
        )

    # Validate n is a positive integer
    if n <= 0:
        raise ValueError(f"n must be a positive integer, got {n}")

    # Handle perfect correlations (arctanh undefined at ±1)
    if r == 1.0:
        return 1.0, 1.0
    if r == -1.0:
        return -1.0, -1.0

    # Handle small n (variance undefined when n-3 <= 0)
    if n < 4:
        return np.nan, np.nan

    # Fisher z-transformation
    z = np.arctanh(r)

    # Standard error on z-scale
    se_z = 1.0 / np.sqrt(n - 3)

    # Critical value for the confidence level
    alpha = 1.0 - confidence_level
    z_critical = stats.norm.ppf(1.0 - alpha / 2.0)

    # CI on z-scale
    z_low = z - z_critical * se_z
    z_high = z + z_critical * se_z

    # Back-transform to r-scale
    ci_low = np.tanh(z_low)
    ci_high = np.tanh(z_high)

    # Clamp to [-1, 1] for numerical precision
    ci_low = max(-1.0, min(1.0, ci_low))
    ci_high = max(-1.0, min(1.0, ci_high))

    return ci_low, ci_high


def minimum_detectable_correlation(
    n: int, alpha: float = 0.05, power: float = 0.80
) -> float:
    """Calculate the minimum detectable correlation coefficient given sample size and power.

    Uses the Fisher z-transformation to determine the smallest correlation
    that can be detected with the specified statistical power:

    1. Compute critical z-values: z_α = Φ⁻¹(1 - α/2), z_β = Φ⁻¹(power)
    2. Minimum detectable z: z_r = (z_α + z_β) / √(n-3)
    3. Back-transform: r = tanh(z_r)

    This calculation is valid for both Pearson and Spearman correlations,
    as both use the same Fisher z-transformation framework for power analysis.

    Args:
        n: Sample size (number of paired observations). Must be a positive integer.
            For n < 4, returns NaN (Fisher z variance undefined with n-3 denominator).
        alpha: Significance level (Type I error rate), default 0.05.
            Must be in (0, 1) exclusive range.
        power: Desired statistical power (1 - Type II error rate), default 0.80.
            Must be in (0, 1) exclusive range.

    Returns:
        Minimum detectable correlation coefficient (always positive).
        Returns NaN if n < 4 (variance formula undefined).

    Raises:
        ValueError: If n <= 0, alpha not in (0, 1), or power not in (0, 1).

    References:
        Cohen, J. (1988). Statistical Power Analysis for the Behavioral Sciences.
            Lawrence Erlbaum Associates, 2nd edition.
        Bonett, D.G. & Wright, T.A. (2000). Sample size requirements for estimating
            Pearson, Kendall and Spearman correlations. Psychometrika, 65(1), 23-28.

    Examples:
        >>> mdr = minimum_detectable_correlation(n=20, alpha=0.05, power=0.80)
        >>> print(f"Minimum detectable r with n=20: {mdr:.3f}")
        Minimum detectable r with n=20: 0.582

        >>> # Larger sample can detect smaller effects
        >>> mdr_100 = minimum_detectable_correlation(n=100, alpha=0.05, power=0.80)
        >>> print(f"Minimum detectable r with n=100: {mdr_100:.3f}")
        Minimum detectable r with n=100: 0.277
    """
    # Validate n is a positive integer
    if n <= 0:
        raise ValueError(f"n must be a positive integer, got {n}")

    # Validate alpha is in (0, 1)
    if not (0 < alpha < 1):
        raise ValueError(f"alpha must be in (0, 1) exclusive range, got {alpha}")

    # Validate power is in (0, 1)
    if not (0 < power < 1):
        raise ValueError(f"power must be in (0, 1) exclusive range, got {power}")

    # Handle small n (variance undefined when n-3 <= 0)
    if n < 4:
        return np.nan

    # Critical z-values
    z_alpha = stats.norm.ppf(1 - alpha / 2)  # Two-tailed
    z_beta = stats.norm.ppf(power)

    # Minimum detectable effect on z-scale
    z_r = (z_alpha + z_beta) / np.sqrt(n - 3)

    # Back-transform to r-scale
    mdr = np.tanh(z_r)

    return mdr


def achieved_power(r: float, n: int, alpha: float = 0.05) -> float:
    """Calculate the achieved statistical power for an observed correlation.

    Uses the Fisher z-transformation to compute power:

    1. Transform correlation: z_r = arctanh(|r|)
    2. Effect size on z-scale: λ = z_r × √(n-3)
    3. Critical z-value: z_α = Φ⁻¹(1 - α/2)
    4. Power: Φ(λ - z_α) + Φ(-λ - z_α)

    The second term accounts for detecting effects in either direction
    (two-tailed test). For practical purposes, when the effect is non-zero,
    this term is negligible.

    This calculation is valid for both Pearson and Spearman correlations,
    as both use the same Fisher z-transformation framework.

    Args:
        r: Observed correlation coefficient (Pearson or Spearman), must be in [-1, 1].
            NaN is allowed and returns NaN.
        n: Sample size (number of paired observations). Must be a positive integer.
            For n < 4, returns NaN (Fisher z variance undefined).
        alpha: Significance level (Type I error rate), default 0.05.
            Must be in (0, 1) exclusive range.

    Returns:
        Achieved statistical power (probability of detecting the effect).
        Returns NaN if:
            - r is NaN
            - n < 4 (variance formula undefined)
        Returns approximately alpha if r = 0 (no effect to detect).

    Raises:
        ValueError: If r is not in [-1, 1] (and not NaN), n <= 0,
            or alpha not in (0, 1).

    References:
        Cohen, J. (1988). Statistical Power Analysis for the Behavioral Sciences.
            Lawrence Erlbaum Associates, 2nd edition.

    Examples:
        >>> power = achieved_power(r=0.5, n=20, alpha=0.05)
        >>> print(f"Power for r=0.5, n=20: {power:.3f}")
        Power for r=0.5, n=20: 0.654

        >>> # Larger sample gives higher power
        >>> power_100 = achieved_power(r=0.5, n=100, alpha=0.05)
        >>> print(f"Power for r=0.5, n=100: {power_100:.3f}")
        Power for r=0.5, n=100: 1.000
    """
    # Handle NaN input first (before validation)
    if np.isnan(r):
        return np.nan

    # Validate r is in [-1, 1]
    if abs(r) > 1:
        raise ValueError(f"Correlation coefficient r must be in [-1, 1], got {r}")

    # Validate n is a positive integer
    if n <= 0:
        raise ValueError(f"n must be a positive integer, got {n}")

    # Validate alpha is in (0, 1)
    if not (0 < alpha < 1):
        raise ValueError(f"alpha must be in (0, 1) exclusive range, got {alpha}")

    # Handle small n (variance undefined when n-3 <= 0)
    if n < 4:
        return np.nan

    # Use absolute value of r (power is symmetric)
    abs_r = abs(r)

    # Handle perfect correlation (arctanh undefined at ±1)
    if abs_r == 1.0:
        return 1.0

    # Fisher z-transformation of the effect
    z_r = np.arctanh(abs_r)

    # Effect size (non-centrality parameter)
    lambda_nc = z_r * np.sqrt(n - 3)

    # Critical z-value (two-tailed)
    z_alpha = stats.norm.ppf(1 - alpha / 2)

    # Power calculation (two-tailed test)
    # P(reject H0 | H1 true) = P(Z > z_α - λ) + P(Z < -z_α - λ)
    power = stats.norm.cdf(lambda_nc - z_alpha) + stats.norm.cdf(-lambda_nc - z_alpha)

    return power


def create_correlation_summary_plot(
    correlation_df: pd.DataFrame,
    correlation_col: str = "correlation",
    pvalue_col: str = "p_value",
    exp1_trait_col: str = "exp1_trait",
    exp2_trait_col: str = "exp2_trait",
    figsize: tuple = (14, 12),
    significance_threshold: float = 0.05,
    top_n: int = 15,
    significant_col: Optional[str] = None,
) -> plt.Figure:
    """Create a 4-panel summary visualization of cross-platform correlations.

    Creates a comprehensive summary figure with:
    - Panel 1 (top-left): Distribution histogram of correlations
    - Panel 2 (top-right): Volcano plot (correlation vs -log10(p-value))
    - Panel 3 (bottom-left): Top positive correlations bar chart
    - Panel 4 (bottom-right): Top negative correlations bar chart

    Args:
        correlation_df: DataFrame with correlation results
        correlation_col: Column name containing correlation coefficients
        pvalue_col: Column name containing p-values
        exp1_trait_col: Column name for experiment 1 trait names
        exp2_trait_col: Column name for experiment 2 trait names
        figsize: Figure size as (width, height) tuple
        significance_threshold: P-value threshold for significance lines
        top_n: Number of top correlations to show in bar charts
        significant_col: Optional column name containing pre-computed significance
            flags (e.g., "significant_fdr"). If provided, uses this column for
            counting significant correlations instead of p-value thresholding.

    Returns:
        matplotlib Figure object

    Raises:
        ValueError: If correlation_df is empty or missing required columns

    Examples:
        >>> correlation_df = pd.DataFrame({
        ...     'correlation': [0.8, -0.6, 0.4],
        ...     'p_value': [0.01, 0.05, 0.10],
        ...     'exp1_trait': ['trait1', 'trait2', 'trait3'],
        ...     'exp2_trait': ['traitA', 'traitB', 'traitC']
        ... })
        >>> fig = create_correlation_summary_plot(correlation_df)
    """
    # Validate inputs
    if len(correlation_df) == 0:
        raise ValueError("correlation_df cannot be empty")

    required_cols = [correlation_col, pvalue_col, exp1_trait_col, exp2_trait_col]
    missing_cols = [col for col in required_cols if col not in correlation_df.columns]
    if missing_cols:
        raise ValueError(f"Column(s) {missing_cols} not found in correlation_df")

    # Create figure with 2x2 subplots
    fig, axes = plt.subplots(2, 2, figsize=figsize)

    # Panel 1: Distribution of correlations
    ax = axes[0, 0]
    ax.hist(
        correlation_df[correlation_col],
        bins=50,
        edgecolor="black",
        alpha=0.7,
        color="#2E6E73",
    )
    ax.axvline(0, color="gray", linestyle="-", alpha=0.3, linewidth=1)
    ax.set_xlabel("Correlation Coefficient")
    ax.set_ylabel("Count")
    ax.set_title("Distribution of Correlations")

    # Count significant correlations
    if significant_col is not None and significant_col in correlation_df.columns:
        # Use pre-computed significance column (e.g., FDR-corrected)
        n_significant = correlation_df[significant_col].sum()
        sig_label = "Significant (FDR)"
    else:
        # Fall back to p-value thresholding
        n_significant = (correlation_df[pvalue_col] < significance_threshold).sum()
        sig_label = "Significant"
    ax.text(
        0.05,
        0.95,
        f"n = {len(correlation_df):,}\n{sig_label}: {n_significant}",
        transform=ax.transAxes,
        fontsize=10,
        verticalalignment="top",
        bbox=dict(boxstyle="round", facecolor="white", alpha=0.8),
    )
    ax.grid(True, alpha=0.3)

    # Panel 2: Volcano plot
    ax = axes[0, 1]
    colors = [
        "red" if p < 0.01 else "orange" if p < 0.05 else "gray"
        for p in correlation_df[pvalue_col]
    ]
    ax.scatter(
        correlation_df[correlation_col],
        -np.log10(correlation_df[pvalue_col]),
        c=colors,
        alpha=0.5,
        s=10,
    )
    ax.axhline(
        -np.log10(0.05), color="orange", linestyle="--", alpha=0.5, label="p=0.05"
    )
    ax.axhline(-np.log10(0.01), color="red", linestyle="--", alpha=0.5, label="p=0.01")
    ax.axvline(0, color="black", linestyle="-", alpha=0.3)
    ax.set_xlabel("Correlation Coefficient")
    ax.set_ylabel("-log10(p-value)")
    ax.set_title("Volcano Plot")
    ax.legend(loc="upper right")
    ax.grid(True, alpha=0.3)

    # Panel 3: Top positive correlations
    ax = axes[1, 0]
    top_positive = correlation_df[correlation_df[correlation_col] > 0].head(top_n)

    if len(top_positive) > 0:
        y_pos = np.arange(len(top_positive))
        ax.barh(y_pos, top_positive[correlation_col].values, color="#2E6E73", alpha=0.7)
        ax.set_yticks(y_pos)
        labels = [
            f"{row[exp1_trait_col][:15]}...{row[exp2_trait_col][:15]}"
            for _, row in top_positive.iterrows()
        ]
        ax.set_yticklabels(labels, fontsize=7)
        ax.set_xlabel("Correlation Coefficient")
        ax.set_title("Top Positive Correlations")
        ax.invert_yaxis()
        ax.grid(axis="x", alpha=0.3)
    else:
        ax.text(
            0.5,
            0.5,
            "No positive correlations found",
            ha="center",
            va="center",
            transform=ax.transAxes,
        )
        ax.set_title("Top Positive Correlations")

    # Panel 4: Top negative correlations
    ax = axes[1, 1]
    top_negative = correlation_df[correlation_df[correlation_col] < 0].head(top_n)

    if len(top_negative) > 0:
        y_pos = np.arange(len(top_negative))
        ax.barh(y_pos, top_negative[correlation_col].values, color="#E8998D", alpha=0.7)
        ax.set_yticks(y_pos)
        labels = [
            f"{row[exp1_trait_col][:15]}...{row[exp2_trait_col][:15]}"
            for _, row in top_negative.iterrows()
        ]
        ax.set_yticklabels(labels, fontsize=7)
        ax.set_xlabel("Correlation Coefficient")
        ax.set_title("Top Negative Correlations")
        ax.invert_yaxis()
        ax.grid(axis="x", alpha=0.3)
    else:
        ax.text(
            0.5,
            0.5,
            "No negative correlations found",
            ha="center",
            va="center",
            transform=ax.transAxes,
        )
        ax.set_title("Top Negative Correlations")

    plt.tight_layout()

    return fig


def load_and_align_experiments(
    exp1_path: Path | str,
    exp2_path: Path | str,
    genotype_col1: str = "Geno",
    genotype_col2: str = "geno",
    rep_col1: str = "Rep",
    rep_col2: str = "rep",
) -> Tuple[pd.DataFrame, pd.DataFrame, List[str]]:
    """Load and align two experimental datasets by genotype.

    Args:
        exp1_path: Path to first experiment data
        exp2_path: Path to second experiment data
        genotype_col1: Genotype column name in first dataset
        genotype_col2: Genotype column name in second dataset
        rep_col1: Replicate column name in first dataset
        rep_col2: Replicate column name in second dataset

    Returns:
        Tuple of (exp1_data, exp2_data, common_genotypes)
    """
    exp1_path = Path(exp1_path)
    exp2_path = Path(exp2_path)

    # Load data
    exp1_df = pd.read_csv(exp1_path)
    exp2_df = pd.read_csv(exp2_path)

    # Standardize column names
    # Handle both provided rep column names and common variations
    rep_renames_exp1 = {genotype_col1: "genotype"}
    rep_renames_exp2 = {genotype_col2: "genotype"}

    # Find and rename replicate column (try multiple common names)
    # Warn if multiple variants exist to avoid ambiguity
    rep_candidates = [rep_col1, "Replicate", "replicate", "Rep", "rep"]
    found_rep_cols_exp1 = [
        c for c in rep_candidates if c in exp1_df.columns and c != "replicate"
    ]
    if len(found_rep_cols_exp1) > 1:
        warnings.warn(
            f"Experiment 1: Multiple replicate column variants found: {found_rep_cols_exp1}. "
            f"Using '{found_rep_cols_exp1[0]}'.",
            UserWarning,
        )
    if found_rep_cols_exp1:
        rep_renames_exp1[found_rep_cols_exp1[0]] = "replicate"

    found_rep_cols_exp2 = [
        c
        for c in [rep_col2, "Replicate", "replicate", "Rep", "rep"]
        if c in exp2_df.columns and c != "replicate"
    ]
    if len(found_rep_cols_exp2) > 1:
        warnings.warn(
            f"Experiment 2: Multiple replicate column variants found: {found_rep_cols_exp2}. "
            f"Using '{found_rep_cols_exp2[0]}'.",
            UserWarning,
        )
    if found_rep_cols_exp2:
        rep_renames_exp2[found_rep_cols_exp2[0]] = "replicate"

    exp1_df = exp1_df.rename(columns=rep_renames_exp1)
    exp2_df = exp2_df.rename(columns=rep_renames_exp2)

    # Find common genotypes
    genotypes1 = set(exp1_df["genotype"].unique())
    genotypes2 = set(exp2_df["genotype"].unique())
    common_genotypes = sorted(list(genotypes1.intersection(genotypes2)))

    # Filter to common genotypes
    exp1_df = exp1_df[exp1_df["genotype"].isin(common_genotypes)]
    exp2_df = exp2_df[exp2_df["genotype"].isin(common_genotypes)]

    print(f"Found {len(common_genotypes)} common genotypes between experiments")
    print(f"Experiment 1: {len(exp1_df)} samples, {len(exp1_df.columns)} traits")
    print(f"Experiment 2: {len(exp2_df)} samples, {len(exp2_df.columns)} traits")

    return exp1_df, exp2_df, common_genotypes


def calculate_genotype_means(
    df: pd.DataFrame, trait_cols: List[str], genotype_col: str = "genotype"
) -> pd.DataFrame:
    """Calculate mean trait values per genotype.

    Args:
        df: DataFrame with trait data
        trait_cols: List of trait columns to aggregate
        genotype_col: Name of genotype column

    Returns:
        DataFrame with mean trait values per genotype
    """
    # Group by genotype and calculate means
    genotype_means = df.groupby(genotype_col)[trait_cols].mean()

    # Add sample counts
    genotype_means["n_samples"] = df.groupby(genotype_col).size()

    return genotype_means


def calculate_genotype_statistics(
    df: pd.DataFrame,
    trait_cols: List[str],
    genotype_col: str = "genotype",
    statistics: Optional[List[str]] = None,
) -> Dict[str, pd.DataFrame]:
    """Calculate multiple statistics per genotype for all traits.

    Args:
        df: DataFrame with trait data
        trait_cols: List of trait columns to aggregate
        genotype_col: Name of genotype column
        statistics: List of statistics to compute.
            Default: ["mean", "median", "min", "max", "std"]

    Returns:
        Dictionary with statistic name as key and DataFrame as value
    """
    if statistics is None:
        statistics = ["mean", "median", "min", "max", "std"]

    results = {}

    # Group by genotype
    grouped = df.groupby(genotype_col)[trait_cols]

    # Calculate each statistic
    for stat in statistics:
        if stat == "mean":
            stat_df = grouped.mean()
        elif stat == "median":
            stat_df = grouped.median()
        elif stat == "min":
            stat_df = grouped.min()
        elif stat == "max":
            stat_df = grouped.max()
        elif stat == "std":
            stat_df = grouped.std()
        elif stat == "count":
            stat_df = grouped.count()
        else:
            continue

        # Add sample counts
        stat_df["n_samples"] = df.groupby(genotype_col).size()
        results[stat] = stat_df

    return results


def calculate_cross_experiment_correlations_extended(
    exp1_stats: Dict[str, pd.DataFrame],
    exp2_stats: Dict[str, pd.DataFrame],
    exp1_traits: List[str],
    exp2_traits: List[str],
    min_samples: int = 3,
    top_n: Optional[int] = None,
) -> pd.DataFrame:
    """Calculate correlations between traits using all statistic combinations.

    Args:
        exp1_stats: Dictionary of statistics DataFrames for experiment 1
        exp2_stats: Dictionary of statistics DataFrames for experiment 2
        exp1_traits: Trait columns from experiment 1
        exp2_traits: Trait columns from experiment 2
        min_samples: Minimum samples required per genotype
        top_n: If specified, only return top N correlations

    Returns:
        DataFrame with correlation results for all statistic combinations
    """
    results = []

    # Iterate through all statistic combinations
    for stat1_name, exp1_df in exp1_stats.items():
        for stat2_name, exp2_df in exp2_stats.items():
            # Find common genotypes with sufficient samples
            valid_genotypes = list(set(exp1_df.index) & set(exp2_df.index))

            if min_samples > 0:
                valid_genotypes = [
                    g
                    for g in valid_genotypes
                    if exp1_df.loc[g, "n_samples"] >= min_samples
                    and exp2_df.loc[g, "n_samples"] >= min_samples
                ]

            if len(valid_genotypes) < 3:
                continue

            # Calculate correlations for this statistic combination
            for trait1 in exp1_traits:
                for trait2 in exp2_traits:
                    # Get aligned and cleaned values
                    values1_clean, values2_clean, _ = _prepare_aligned_values(
                        exp1_df, exp2_df, trait1, trait2, min_samples=0
                    )

                    if len(values1_clean) < 3:
                        continue

                    # Calculate correlations
                    pearson_corr, pearson_p, spearman_corr, spearman_p = (
                        _calculate_correlations(values1_clean, values2_clean)
                    )

                    results.append(
                        {
                            "exp1_trait": trait1,
                            "exp2_trait": trait2,
                            "exp1_statistic": stat1_name,
                            "exp2_statistic": stat2_name,
                            "pearson_r": pearson_corr,
                            "pearson_p": pearson_p,
                            "spearman_r": spearman_corr,
                            "spearman_p": spearman_p,
                            "n_genotypes": len(values1_clean),
                            "abs_pearson": abs(pearson_corr),
                            "abs_spearman": abs(spearman_corr),
                        }
                    )

    results_df = pd.DataFrame(results)

    if len(results_df) == 0:
        return results_df

    # Add significance flags
    results_df["pearson_significant"] = results_df["pearson_p"] < 0.05
    results_df["spearman_significant"] = results_df["spearman_p"] < 0.05

    # Sort by absolute correlation
    results_df = results_df.sort_values("abs_pearson", ascending=False)

    # Return top N if specified
    if top_n is not None and len(results_df) > top_n:
        return results_df.head(top_n)

    return results_df


def summarize_statistic_combinations(
    extended_corr_df: pd.DataFrame, metric: str = "pearson_r"
) -> pd.DataFrame:
    """Summarize which statistic combinations yield the strongest correlations.

    Args:
        extended_corr_df: DataFrame from calculate_cross_experiment_correlations_extended
        metric: Which correlation metric to use ('pearson_r' or 'spearman_r')

    Returns:
        DataFrame summarizing best statistic combinations
    """
    if len(extended_corr_df) == 0:
        return pd.DataFrame()

    abs_metric = f"abs_{metric.split('_')[0]}"

    # Group by statistic combination
    stat_summary = (
        extended_corr_df.groupby(["exp1_statistic", "exp2_statistic"])
        .agg(
            {
                abs_metric: ["mean", "max", "std", "count"],
                f"{metric.split('_')[0]}_significant": "sum",
            }
        )
        .round(3)
    )

    # Flatten column names
    stat_summary.columns = [
        "_".join(col).strip() for col in stat_summary.columns.values
    ]
    stat_summary = stat_summary.reset_index()

    # Sort by mean absolute correlation
    stat_summary = stat_summary.sort_values(f"{abs_metric}_mean", ascending=False)

    return stat_summary


def create_statistic_combination_heatmap(
    extended_corr_df: pd.DataFrame,
    trait1: str,
    trait2: str,
    metric: str = "pearson_r",
    figsize: Tuple[int, int] = (10, 8),
    cmap: str = "RdBu_r",
) -> plt.Figure:
    """Create heatmap showing correlations across all statistic combinations for a trait pair.

    Args:
        extended_corr_df: DataFrame from calculate_cross_experiment_correlations_extended
        trait1: Trait from experiment 1
        trait2: Trait from experiment 2
        metric: Which correlation metric to plot
        figsize: Figure size
        cmap: Colormap

    Returns:
        Matplotlib figure
    """
    # Handle empty dataframe
    if len(extended_corr_df) == 0 or "exp1_trait" not in extended_corr_df.columns:
        fig, ax = plt.subplots(figsize=figsize)
        ax.text(
            0.5,
            0.5,
            "No data available",
            ha="center",
            va="center",
            fontsize=14,
        )
        return fig

    # Filter to specific trait pair
    trait_data = extended_corr_df[
        (extended_corr_df["exp1_trait"] == trait1)
        & (extended_corr_df["exp2_trait"] == trait2)
    ]

    if len(trait_data) == 0:
        fig, ax = plt.subplots(figsize=figsize)
        ax.text(
            0.5,
            0.5,
            "No data for this trait pair",
            ha="center",
            va="center",
            fontsize=14,
        )
        return fig

    # Pivot for heatmap
    heatmap_data = trait_data.pivot(
        index="exp1_statistic", columns="exp2_statistic", values=metric
    )

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

    # Create heatmap
    sns.heatmap(
        heatmap_data,
        annot=True,
        fmt=".3f",
        cmap=cmap,
        center=0,
        vmin=-1,
        vmax=1,
        square=True,
        ax=ax,
        cbar_kws={"label": f"{metric.replace('_', ' ').title()}"},
    )

    ax.set_title(
        f"Correlation Across Statistic Combinations\n{trait1} vs {trait2}", fontsize=12
    )
    ax.set_xlabel("Experiment 2 Statistic", fontsize=10)
    ax.set_ylabel("Experiment 1 Statistic", fontsize=10)

    plt.tight_layout()

    return fig


def calculate_cross_experiment_correlations(
    exp1_means: pd.DataFrame,
    exp2_means: pd.DataFrame,
    exp1_traits: List[str],
    exp2_traits: List[str],
    min_samples: int = 3,
) -> pd.DataFrame:
    """Calculate correlations between traits from two experiments.

    Args:
        exp1_means: Mean trait values from experiment 1
        exp2_means: Mean trait values from experiment 2
        exp1_traits: Trait columns from experiment 1
        exp2_traits: Trait columns from experiment 2
        min_samples: Minimum samples required per genotype

    Returns:
        DataFrame with correlation results
    """
    # Find common genotypes with sufficient samples
    valid_genotypes = list(set(exp1_means.index) & set(exp2_means.index))

    if min_samples > 0:
        valid_genotypes = [
            g
            for g in valid_genotypes
            if exp1_means.loc[g, "n_samples"] >= min_samples
            and exp2_means.loc[g, "n_samples"] >= min_samples
        ]

    print(f"Using {len(valid_genotypes)} genotypes with >= {min_samples} samples each")

    # Calculate correlations
    results = []
    for trait1 in exp1_traits:
        for trait2 in exp2_traits:
            # Get aligned and cleaned values
            values1_clean, values2_clean, _ = _prepare_aligned_values(
                exp1_means, exp2_means, trait1, trait2, min_samples=0
            )

            if len(values1_clean) < 3:
                continue

            # Calculate correlation (only Pearson for backward compatibility)
            corr, pval, _, _ = _calculate_correlations(values1_clean, values2_clean)

            results.append(
                {
                    "exp1_trait": trait1,
                    "exp2_trait": trait2,
                    "correlation": corr,
                    "p_value": pval,
                    "n_genotypes": len(values1_clean),
                    "abs_correlation": abs(corr),
                }
            )

    results_df = pd.DataFrame(results)

    # Add significance flags if there are results
    if len(results_df) > 0:
        results_df["significant"] = results_df["p_value"] < 0.05
        results_df["highly_significant"] = results_df["p_value"] < 0.01
        # Sort by absolute correlation
        results_df = results_df.sort_values("abs_correlation", ascending=False)
    else:
        # Create empty dataframe with expected columns
        results_df = pd.DataFrame(
            columns=[
                "exp1_trait",
                "exp2_trait",
                "correlation",
                "p_value",
                "n_genotypes",
                "abs_correlation",
                "significant",
                "highly_significant",
            ]
        )

    return results_df


def create_cross_experiment_heatmap(
    correlation_df: pd.DataFrame,
    top_n_traits: int = 20,
    figsize: Tuple[int, int] = (14, 12),
    cmap: str = "RdBu_r",
    vmin: float = -1.0,
    vmax: float = 1.0,
) -> plt.Figure:
    """Create heatmap of cross-experiment correlations.

    Args:
        correlation_df: DataFrame with correlation results
        top_n_traits: Number of top traits to show from each experiment
        figsize: Figure size
        cmap: Colormap for heatmap
        vmin: Minimum value for color scale
        vmax: Maximum value for color scale

    Returns:
        Matplotlib figure
    """
    # Handle empty dataframe
    if len(correlation_df) == 0:
        fig, ax = plt.subplots(figsize=figsize)
        ax.text(
            0.5,
            0.5,
            "No correlations to display",
            ha="center",
            va="center",
            fontsize=14,
        )
        ax.set_title("Cross-Experiment Trait Correlations")
        return fig

    # Get top traits from each experiment
    top_exp1_traits = (
        correlation_df.groupby("exp1_trait")["abs_correlation"]
        .max()
        .nlargest(top_n_traits)
        .index.tolist()
    )
    top_exp2_traits = (
        correlation_df.groupby("exp2_trait")["abs_correlation"]
        .max()
        .nlargest(top_n_traits)
        .index.tolist()
    )

    # Filter to top traits
    filtered_df = correlation_df[
        (correlation_df["exp1_trait"].isin(top_exp1_traits))
        & (correlation_df["exp2_trait"].isin(top_exp2_traits))
    ]

    # Pivot to matrix form
    corr_matrix = filtered_df.pivot(
        index="exp1_trait", columns="exp2_trait", values="correlation"
    )

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

    # Create heatmap
    sns.heatmap(
        corr_matrix,
        annot=True,
        fmt=".2f",
        cmap=cmap,
        center=0,
        vmin=vmin,
        vmax=vmax,
        square=False,
        ax=ax,
        cbar_kws={"label": "Pearson Correlation", "shrink": 0.8},
    )

    ax.set_title(
        "Cross-Experiment Trait Correlations\n(Top Correlated Traits)", fontsize=14
    )
    ax.set_xlabel("Experiment 2 Traits", fontsize=12)
    ax.set_ylabel("Experiment 1 Traits", fontsize=12)

    # Rotate labels
    plt.setp(ax.get_xticklabels(), rotation=45, ha="right", fontsize=9)
    plt.setp(ax.get_yticklabels(), rotation=0, fontsize=9)

    plt.tight_layout()

    return fig


def create_top_correlations_plot(
    correlation_df: pd.DataFrame, top_n: int = 30, figsize: Tuple[int, int] = (14, 8)
) -> plt.Figure:
    """Create bar plot of top correlations between experiments.

    Args:
        correlation_df: DataFrame with correlation results
        top_n: Number of top correlations to show
        figsize: Figure size

    Returns:
        Matplotlib figure
    """
    # Get top correlations
    top_corr = correlation_df.nlargest(top_n, "abs_correlation").copy()

    # Create trait pair labels
    top_corr["trait_pair"] = (
        top_corr["exp1_trait"].str[:20] + " vs " + top_corr["exp2_trait"].str[:20]
    )

    # Create figure
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=figsize)

    # Plot absolute correlations
    colors = ["red" if x < 0 else "blue" for x in top_corr["correlation"].values]
    bars1 = ax1.barh(
        range(len(top_corr)),
        top_corr["abs_correlation"].values,
        color=colors,
        alpha=0.7,
    )
    ax1.set_yticks(range(len(top_corr)))
    ax1.set_yticklabels(top_corr["trait_pair"].values, fontsize=8)
    ax1.set_xlabel("Absolute Correlation", fontsize=10)
    ax1.set_title(f"Top {top_n} Cross-Experiment Correlations", fontsize=12)
    ax1.invert_yaxis()
    ax1.grid(axis="x", alpha=0.3)

    # Add significance markers
    for i, (corr, sig) in enumerate(
        zip(top_corr["abs_correlation"], top_corr["highly_significant"])
    ):
        if sig:
            ax1.text(corr + 0.01, i, "**", va="center", fontsize=8)
        elif top_corr.iloc[i]["significant"]:
            ax1.text(corr + 0.01, i, "*", va="center", fontsize=8)

    # Plot signed correlations
    bars2 = ax2.barh(
        range(len(top_corr)), top_corr["correlation"].values, color=colors, alpha=0.7
    )
    ax2.set_yticks(range(len(top_corr)))
    ax2.set_yticklabels([])
    ax2.set_xlabel("Correlation (Pearson r)", fontsize=10)
    ax2.set_title("Signed Correlations", fontsize=12)
    ax2.invert_yaxis()
    ax2.grid(axis="x", alpha=0.3)
    ax2.axvline(x=0, color="black", linestyle="-", linewidth=0.5)

    # Add legend
    from matplotlib.patches import Patch

    legend_elements = [
        Patch(facecolor="blue", alpha=0.7, label="Positive"),
        Patch(facecolor="red", alpha=0.7, label="Negative"),
    ]
    ax2.legend(handles=legend_elements, loc="lower right", fontsize=9)

    plt.suptitle("Cross-Experiment Trait Correlations Analysis", fontsize=14, y=1.02)
    plt.tight_layout()

    return fig


def create_scatter_plot_grid(
    exp1_means: pd.DataFrame,
    exp2_means: pd.DataFrame,
    top_pairs: List[Tuple[str, str]],
    figsize: Tuple[int, int] = (16, 12),
    n_cols: int = 4,
) -> plt.Figure:
    """Create grid of scatter plots for top correlated trait pairs.

    Args:
        exp1_means: Mean trait values from experiment 1
        exp2_means: Mean trait values from experiment 2
        top_pairs: List of (exp1_trait, exp2_trait) tuples
        figsize: Figure size
        n_cols: Number of columns in grid

    Returns:
        Matplotlib figure
    """
    n_pairs = len(top_pairs)
    n_rows = (n_pairs + n_cols - 1) // n_cols

    fig, axes = plt.subplots(n_rows, n_cols, figsize=figsize)
    axes = axes.flatten() if n_pairs > 1 else [axes]

    # Find common genotypes
    common_genotypes = list(set(exp1_means.index) & set(exp2_means.index))

    for idx, (trait1, trait2) in enumerate(top_pairs):
        ax = axes[idx]

        # Get aligned values
        x_vals = exp1_means.loc[common_genotypes, trait1].values
        y_vals = exp2_means.loc[common_genotypes, trait2].values

        # Remove NaNs
        valid_mask = ~(np.isnan(x_vals) | np.isnan(y_vals))
        x_vals = x_vals[valid_mask]
        y_vals = y_vals[valid_mask]

        # Create scatter plot
        ax.scatter(x_vals, y_vals, alpha=0.6, s=30)

        # Add regression line
        if len(x_vals) > 2:
            z = np.polyfit(x_vals, y_vals, 1)
            p = np.poly1d(z)
            x_line = np.linspace(x_vals.min(), x_vals.max(), 100)
            ax.plot(x_line, p(x_line), "r-", alpha=0.5, linewidth=2)

            # Calculate correlation
            corr, pval, _, _ = _calculate_correlations(x_vals, y_vals)
            ax.text(
                0.05,
                0.95,
                f"r={corr:.3f}\np={pval:.3e}",
                transform=ax.transAxes,
                fontsize=8,
                verticalalignment="top",
                bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.5),
            )

        # Labels
        ax.set_xlabel(f"Exp1: {trait1[:25]}", fontsize=8)
        ax.set_ylabel(f"Exp2: {trait2[:25]}", fontsize=8)
        ax.tick_params(labelsize=7)
        ax.grid(True, alpha=0.3)

    # Hide unused subplots
    for idx in range(n_pairs, len(axes)):
        axes[idx].set_visible(False)

    plt.suptitle("Scatter Plots of Top Correlated Trait Pairs", fontsize=14)
    plt.tight_layout()

    return fig


def calculate_per_trait_correlations(
    exp1_df: pd.DataFrame,
    exp2_df: pd.DataFrame,
    exp1_trait: str,
    exp2_trait: str,
    genotype_col: str = "genotype",
) -> Dict:
    """Calculate correlation between specific traits using individual samples.

    Args:
        exp1_df: DataFrame from experiment 1 with individual samples
        exp2_df: DataFrame from experiment 2 with individual samples
        exp1_trait: Trait column from experiment 1
        exp2_trait: Trait column from experiment 2
        genotype_col: Name of genotype column

    Returns:
        Dictionary with correlation statistics
    """
    # Get genotype means for each trait
    exp1_means = exp1_df.groupby(genotype_col)[exp1_trait].mean()
    exp2_means = exp2_df.groupby(genotype_col)[exp2_trait].mean()

    # Find common genotypes
    common_genos = list(set(exp1_means.index) & set(exp2_means.index))

    if len(common_genos) < 3:
        return {
            "n_genotypes": len(common_genos),
            "pearson_r": np.nan,
            "pearson_p": np.nan,
            "spearman_r": np.nan,
            "spearman_p": np.nan,
            "valid": False,
        }

    # Align values
    values1 = exp1_means.loc[common_genos].values
    values2 = exp2_means.loc[common_genos].values

    # Remove NaN pairs
    valid_mask = ~(np.isnan(values1) | np.isnan(values2))
    if valid_mask.sum() < 3:
        return {
            "n_genotypes": valid_mask.sum(),
            "pearson_r": np.nan,
            "pearson_p": np.nan,
            "spearman_r": np.nan,
            "spearman_p": np.nan,
            "valid": False,
        }

    values1_clean = values1[valid_mask]
    values2_clean = values2[valid_mask]

    # Calculate correlations
    pearson_r, pearson_p, spearman_r, spearman_p = _calculate_correlations(
        values1_clean, values2_clean
    )

    return {
        "n_genotypes": len(values1_clean),
        "pearson_r": pearson_r,
        "pearson_p": pearson_p,
        "spearman_r": spearman_r,
        "spearman_p": spearman_p,
        "values_exp1": values1_clean,
        "values_exp2": values2_clean,
        "genotypes": [
            common_genos[i] for i in range(len(common_genos)) if valid_mask[i]
        ],
        "valid": True,
    }


def create_joint_plot(
    exp1_means: pd.DataFrame,
    exp2_means: pd.DataFrame,
    exp1_trait: str,
    exp2_trait: str,
    exp1_name: str = "Experiment 1",
    exp2_name: str = "Experiment 2",
    figsize: Tuple[int, int] = (10, 10),
    color: str = "#4CB391",
    line_color: str = "#2E6E73",
    # Pre-computed correlation values (single source of truth - DRY principle)
    correlation: Optional[float] = None,
    p_value: Optional[float] = None,
    n_genotypes: Optional[int] = None,
    # Pre-computed Pearson values (for consistency with CSV output)
    pearson_r: Optional[float] = None,
    pearson_p: Optional[float] = None,
) -> plt.Figure:
    """Create joint plot for two traits with regression line and marginal distributions.

    Args:
        exp1_means: Genotype means from experiment 1
        exp2_means: Genotype means from experiment 2
        exp1_trait: Trait column from experiment 1
        exp2_trait: Trait column from experiment 2
        exp1_name: Name of experiment 1 for labeling
        exp2_name: Name of experiment 2 for labeling
        figsize: Figure size
        color: Color for scatter points
        line_color: Color for regression line
        correlation: Pre-computed Spearman correlation coefficient. When provided,
            this value is displayed instead of recalculating. Use this to ensure
            consistency with CSV output (single source of truth).
        p_value: Pre-computed Spearman p-value corresponding to the correlation.
        n_genotypes: Pre-computed number of genotypes used in correlation calculation.
            May differ from len(plot_df) if min_samples_per_genotype filtering was applied.
        pearson_r: Pre-computed Pearson correlation coefficient. When provided along
            with pearson_p, this value is displayed instead of recalculating.
        pearson_p: Pre-computed Pearson p-value.

    Returns:
        Matplotlib figure

    Note:
        When all pre-computed values are provided (correlation, p_value, n_genotypes,
        pearson_r, pearson_p), these values are displayed in the annotation without
        recalculation. This ensures the plot matches the CSV output exactly (DRY
        principle). When not provided, correlations are calculated from the data
        (backward compatible fallback).

    See Also:
        identify_significant_correlations: Computes cross-experiment correlations
            and exports CSV output. The correlation values from that function
            should be passed here to ensure consistency.
    """
    import seaborn as sns

    # Get common genotypes
    common_genos = list(set(exp1_means.index) & set(exp2_means.index))

    # Create DataFrame for plotting
    plot_df = pd.DataFrame(
        {
            exp1_trait: exp1_means.loc[common_genos, exp1_trait],
            exp2_trait: exp2_means.loc[common_genos, exp2_trait],
        }
    )

    # Remove NaN pairs
    plot_df = plot_df.dropna()

    if len(plot_df) < 3:
        fig, ax = plt.subplots(figsize=figsize)
        ax.text(
            0.5,
            0.5,
            "Insufficient data for plot",
            ha="center",
            va="center",
            fontsize=14,
        )
        return fig

    # Use pre-computed values if provided (single source of truth)
    # Otherwise fall back to calculating from data (backward compatibility)
    if correlation is not None and p_value is not None and n_genotypes is not None:
        # Use pre-computed Spearman values - ensures consistency with CSV output
        spearman_r_val = correlation
        spearman_p_val = p_value
        n_genotypes_for_annotation = n_genotypes

        # Use pre-computed Pearson if provided, otherwise calculate from plot data
        if pearson_r is not None and pearson_p is not None:
            # Both Spearman and Pearson from CSV (single source of truth)
            pearson_r_val = pearson_r
            pearson_p_val = pearson_p
        else:
            # Legacy fallback: only Spearman was pre-computed, calculate Pearson
            pearson_r_val, pearson_p_val, _, _ = _calculate_correlations(
                plot_df[exp1_trait].values, plot_df[exp2_trait].values
            )
    else:
        # Fallback: calculate from data (backward compatible for direct API usage)
        pearson_r_val, pearson_p_val, spearman_r_val, spearman_p_val = (
            _calculate_correlations(
                plot_df[exp1_trait].values, plot_df[exp2_trait].values
            )
        )
        n_genotypes_for_annotation = len(plot_df)

    # Create joint plot
    g = sns.jointplot(
        data=plot_df,
        x=exp1_trait,
        y=exp2_trait,
        kind="reg",
        scatter_kws={"s": 50, "alpha": 0.8, "color": color},
        line_kws={"color": line_color},
        height=figsize[0],
    )

    # Add labels - use sanitize_single_trait_name for consistent formatting
    g.set_axis_labels(
        f"{exp1_name}: {sanitize_single_trait_name(exp1_trait, abbreviate=False)}",
        f"{exp2_name}: {sanitize_single_trait_name(exp2_trait, abbreviate=False)}",
        fontsize=12,
    )

    # Add title
    g.figure.suptitle(
        f"Cross-Experiment Trait Correlation\n{exp1_name} vs {exp2_name}", fontsize=14
    )

    # Add correlation text
    corr_text = f"Pearson r = {pearson_r_val:.3f} (p = {pearson_p_val:.3g})\n"
    corr_text += f"Spearman ρ = {spearman_r_val:.3f} (p = {spearman_p_val:.3g})\n"
    corr_text += f"n = {n_genotypes_for_annotation} genotypes"

    g.ax_joint.annotate(
        corr_text,
        xy=(0.05, 0.95),
        xycoords="axes fraction",
        ha="left",
        va="top",
        fontsize=10,
        bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="gray", lw=1),
    )

    g.figure.tight_layout()
    g.figure.subplots_adjust(top=0.93)

    return g.figure


def identify_significant_correlations(
    correlation_df: pd.DataFrame,
    p_threshold: float = 0.05,
    r_threshold: float = 0.5,
    use_fdr: bool = True,
) -> pd.DataFrame:
    """Identify statistically significant correlations with FDR correction.

    Args:
        correlation_df: DataFrame with correlation results
        p_threshold: P-value threshold for significance
        r_threshold: Minimum absolute correlation threshold
        use_fdr: Whether to use FDR correction for multiple testing

    Returns:
        DataFrame with significant correlations only
    """
    from statsmodels.stats.multitest import multipletests

    # Filter by absolute correlation threshold
    strong_corr = correlation_df[
        correlation_df["abs_correlation"] >= r_threshold
    ].copy()

    if len(strong_corr) == 0:
        return pd.DataFrame()

    # Apply multiple testing correction if requested
    if use_fdr and len(strong_corr) > 1:
        # Benjamini-Hochberg FDR correction
        reject, pvals_corrected, _, _ = multipletests(
            strong_corr["p_value"], alpha=p_threshold, method="fdr_bh"
        )
        strong_corr["p_value_corrected"] = pvals_corrected
        strong_corr["significant_fdr"] = reject

        # Filter to significant after correction
        significant = strong_corr[strong_corr["significant_fdr"]].copy()
    else:
        # Simple p-value threshold
        significant = strong_corr[strong_corr["p_value"] < p_threshold].copy()
        significant["p_value_corrected"] = significant["p_value"]
        significant["significant_fdr"] = True

    return significant.sort_values("abs_correlation", ascending=False)


def calculate_correlation_confidence_intervals(
    correlation_df: pd.DataFrame, n_genotypes: int, confidence: float = 0.95
) -> pd.DataFrame:
    """Calculate confidence intervals for correlations using Fisher's z-transformation.

    This is a DataFrame-based wrapper around :func:`calculate_correlation_ci` that
    applies CI calculation to all rows in a correlation DataFrame. It provides
    consistent edge case handling (perfect correlations, small n, NaN values).

    Args:
        correlation_df: DataFrame with correlation results. Must have a
            "correlation" column containing correlation coefficients.
        n_genotypes: Number of genotypes used in correlation.
        confidence: Confidence level (default 0.95 for 95% CI).
            Must be in (0, 1) exclusive range.

    Returns:
        DataFrame with confidence intervals added as columns:
            - ci_lower: Lower bound of confidence interval
            - ci_upper: Upper bound of confidence interval
            - ci_width: Width of confidence interval (ci_upper - ci_lower)

    Note:
        For scalar CI calculation, see :func:`calculate_correlation_ci`.
    """
    df = correlation_df.copy()

    # Apply calculate_correlation_ci to each row for consistent edge case handling
    ci_results = df["correlation"].apply(
        lambda r: calculate_correlation_ci(r, n_genotypes, confidence)
    )

    # Unpack CI tuples in a single pass for efficiency
    ci_bounds = ci_results.apply(
        lambda x: pd.Series({"ci_lower": x[0], "ci_upper": x[1]})
    )
    df[["ci_lower", "ci_upper"]] = ci_bounds
    df["ci_width"] = df["ci_upper"] - df["ci_lower"]

    return df


def summarize_correlation_results(
    correlation_df: pd.DataFrame,
    exp1_name: str = "Experiment 1",
    exp2_name: str = "Experiment 2",
) -> Dict:
    """Summarize cross-experiment correlation results.

    Args:
        correlation_df: DataFrame with correlation results
        exp1_name: Name of first experiment
        exp2_name: Name of second experiment

    Returns:
        Dictionary with summary statistics
    """
    summary = {
        "experiment_1": exp1_name,
        "experiment_2": exp2_name,
        "total_correlations": len(correlation_df),
        "significant_correlations": correlation_df["significant"].sum(),
        "highly_significant_correlations": correlation_df["highly_significant"].sum(),
        "max_positive_correlation": correlation_df["correlation"].max(),
        "max_negative_correlation": correlation_df["correlation"].min(),
        "mean_abs_correlation": correlation_df["abs_correlation"].mean(),
        "median_abs_correlation": correlation_df["abs_correlation"].median(),
        "n_exp1_traits": correlation_df["exp1_trait"].nunique(),
        "n_exp2_traits": correlation_df["exp2_trait"].nunique(),
    }

    # Get top correlated traits from each experiment
    top_exp1 = (
        correlation_df.groupby("exp1_trait")["abs_correlation"]
        .max()
        .nlargest(5)
        .to_dict()
    )
    top_exp2 = (
        correlation_df.groupby("exp2_trait")["abs_correlation"]
        .max()
        .nlargest(5)
        .to_dict()
    )

    summary["top_exp1_traits"] = top_exp1
    summary["top_exp2_traits"] = top_exp2

    # Get top trait pairs
    top_pairs = correlation_df.nlargest(10, "abs_correlation")[
        ["exp1_trait", "exp2_trait", "correlation"]
    ].to_dict("records")
    summary["top_trait_pairs"] = top_pairs

    return summary


def create_genotype_boxplots(
    exp1_df: pd.DataFrame,
    exp2_df: pd.DataFrame,
    exp1_trait: str,
    exp2_trait: str,
    genotype_col: str = "genotype",
    exp1_name: str = "Experiment 1",
    exp2_name: str = "Experiment 2",
    figsize: Tuple[float, float] = (14, 6),
    colors: Tuple[str, str] = ("#E8998D", "#2E6E73"),
) -> plt.Figure:
    """Create side-by-side boxplots grouped by genotype for two traits.

    Args:
        exp1_df: First experiment data with genotype and trait columns
        exp2_df: Second experiment data with genotype and trait columns
        exp1_trait: Trait name from first experiment
        exp2_trait: Trait name from second experiment
        genotype_col: Column name for genotype
        exp1_name: Name of first experiment for labeling
        exp2_name: Name of second experiment for labeling
        figsize: Figure size (width, height)
        colors: Colors for exp1 and exp2 boxplots

    Returns:
        Figure with genotype-grouped boxplots

    Raises:
        ValueError: If no common genotypes found or no data available
    """
    # Get common genotypes
    common_genotypes = sorted(
        set(exp1_df[genotype_col].unique()) & set(exp2_df[genotype_col].unique())
    )

    if len(common_genotypes) == 0:
        raise ValueError("No common genotypes found between experiments")

    # Create figure
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=figsize)

    # Prepare data for boxplots
    exp1_data = [
        exp1_df[exp1_df[genotype_col] == g][exp1_trait].dropna().values
        for g in common_genotypes
    ]
    exp2_data = [
        exp2_df[exp2_df[genotype_col] == g][exp2_trait].dropna().values
        for g in common_genotypes
    ]

    # Filter out empty groups
    valid_indices = [
        i
        for i, (d1, d2) in enumerate(zip(exp1_data, exp2_data))
        if len(d1) > 0 and len(d2) > 0
    ]

    if len(valid_indices) == 0:
        raise ValueError("No genotypes with data in both experiments")

    exp1_data = [exp1_data[i] for i in valid_indices]
    exp2_data = [exp2_data[i] for i in valid_indices]
    valid_genotypes = [common_genotypes[i] for i in valid_indices]

    # Sort genotypes by median of first experiment
    medians = [np.median(d) for d in exp1_data]
    sorted_indices = np.argsort(medians)
    exp1_data = [exp1_data[i] for i in sorted_indices]
    exp2_data = [exp2_data[i] for i in sorted_indices]
    valid_genotypes = [valid_genotypes[i] for i in sorted_indices]

    # Create boxplots for experiment 1
    bp1 = ax1.boxplot(
        exp1_data,
        tick_labels=valid_genotypes,
        patch_artist=True,
        medianprops=dict(color="red", linewidth=2),
        flierprops=dict(marker="o", markerfacecolor="gray", markersize=4, alpha=0.5),
    )

    for patch in bp1["boxes"]:
        patch.set_facecolor(colors[0])
        patch.set_alpha(0.7)

    ax1.set_xlabel("Genotype", fontsize=10)
    sanitized_exp1_trait = sanitize_single_trait_name(exp1_trait, abbreviate=False)
    ax1.set_ylabel(sanitized_exp1_trait[:50], fontsize=9)
    ax1.set_title(f"{exp1_name}: {sanitized_exp1_trait[:60]}", fontsize=11)
    ax1.tick_params(axis="x", rotation=45, labelsize=8)
    ax1.grid(True, alpha=0.3, axis="y")

    # Add sample sizes
    sample_sizes = [len(d) for d in exp1_data]
    ax1.set_xticklabels(
        [f"{g}\n(n={n})" for g, n in zip(valid_genotypes, sample_sizes)],
        rotation=45,
        ha="right",
        fontsize=7,
    )

    # Create boxplots for experiment 2
    bp2 = ax2.boxplot(
        exp2_data,
        tick_labels=valid_genotypes,
        patch_artist=True,
        medianprops=dict(color="red", linewidth=2),
        flierprops=dict(marker="o", markerfacecolor="gray", markersize=4, alpha=0.5),
    )

    for patch in bp2["boxes"]:
        patch.set_facecolor(colors[1])
        patch.set_alpha(0.7)

    ax2.set_xlabel("Genotype", fontsize=10)
    sanitized_exp2_trait = sanitize_single_trait_name(exp2_trait, abbreviate=False)
    ax2.set_ylabel(sanitized_exp2_trait[:50], fontsize=9)
    ax2.set_title(f"{exp2_name}: {sanitized_exp2_trait[:60]}", fontsize=11)
    ax2.tick_params(axis="x", rotation=45, labelsize=8)
    ax2.grid(True, alpha=0.3, axis="y")

    # Add sample sizes
    sample_sizes = [len(d) for d in exp2_data]
    ax2.set_xticklabels(
        [f"{g}\n(n={n})" for g, n in zip(valid_genotypes, sample_sizes)],
        rotation=45,
        ha="right",
        fontsize=7,
    )

    plt.tight_layout()
    return fig


# ============================================================================
# TRAIT REDUNDANCY REDUCTION FUNCTIONS
# ============================================================================


def cluster_correlated_traits(
    df: pd.DataFrame,
    threshold: float = 0.8,
    linkage: str = "complete",
) -> Dict[int, List[str]]:
    """Cluster traits based on pairwise correlations.

    Uses hierarchical clustering with the specified linkage method to group
    traits that are highly correlated. Traits with |r| >= threshold are
    considered redundant and grouped together.

    Algorithm:
        1. Compute Spearman correlation matrix between all trait pairs
        2. Convert to distance: d = 1 - |r| (range [0, 1])
        3. Apply hierarchical clustering with specified linkage
        4. Cut dendrogram at distance t = 1 - threshold
        5. Return cluster assignments

    Edge cases:
        - Empty DataFrame: Returns empty dict
        - Single trait: Returns {0: [trait_name]}
        - All NaN trait: Placed in singleton cluster
        - Constant trait (zero variance): Placed in singleton cluster

    Args:
        df: DataFrame with traits as columns and genotypes as rows.
            Index should be genotype identifiers.
        threshold: Minimum |r| for traits to be considered redundant.
            Must be in (0, 1]. Default 0.8.
        linkage: Linkage method for hierarchical clustering.
            One of "complete", "average", "single". Default "complete".

    Returns:
        Dict mapping cluster_id (int) to list of trait names (sorted alphabetically).

    Example:
        >>> df = pd.DataFrame({'A': [1, 2, 3], 'B': [1, 2, 3], 'C': [3, 2, 1]})
        >>> clusters = cluster_correlated_traits(df, threshold=0.8)
        >>> # A and B are perfectly correlated, C is anti-correlated
        >>> # Result: {0: ['A', 'B'], 1: ['C']}
    """
    # Handle empty DataFrame
    if df.empty or len(df.columns) == 0:
        return {}

    traits = list(df.columns)
    n_traits = len(traits)

    # Handle single trait
    if n_traits == 1:
        return {0: traits}

    # Compute pairwise Spearman correlations
    # Use min_periods to handle NaN values
    corr_matrix = df.corr(method="spearman")

    # Identify problematic traits (constant or all NaN)
    # These will be placed in singleton clusters
    problematic_traits = []
    valid_traits = []

    for trait in traits:
        # Check if trait is constant (zero variance) or all NaN
        trait_values = df[trait].dropna()
        if len(trait_values) == 0 or trait_values.std() == 0:
            problematic_traits.append(trait)
        else:
            valid_traits.append(trait)

    # If no valid traits, all are singletons
    if len(valid_traits) == 0:
        return {i: [trait] for i, trait in enumerate(sorted(traits))}

    # If only one valid trait, it's a singleton
    if len(valid_traits) == 1:
        clusters = {0: valid_traits}
        # Add problematic traits as singletons
        for i, trait in enumerate(sorted(problematic_traits), start=1):
            clusters[i] = [trait]
        return clusters

    # Extract correlation submatrix for valid traits
    valid_corr = corr_matrix.loc[valid_traits, valid_traits]

    # Convert to distance matrix: d = 1 - |r|
    # Handle any NaN correlations by treating them as maximum distance
    distance_matrix = 1 - valid_corr.abs().fillna(0)

    # Ensure diagonal is 0 (no self-distance)
    np.fill_diagonal(distance_matrix.values, 0)

    # Ensure symmetry (should already be, but just in case of floating point issues)
    distance_matrix = (distance_matrix + distance_matrix.T) / 2

    # Convert to condensed distance matrix for scipy
    condensed_dist = squareform(distance_matrix.values, checks=False)

    # Perform hierarchical clustering
    Z = hierarchical_linkage(condensed_dist, method=linkage)

    # Cut dendrogram at threshold distance
    # distance = 1 - threshold means: cluster traits with |r| >= threshold
    cut_distance = 1 - threshold
    cluster_labels = fcluster(Z, t=cut_distance, criterion="distance")

    # Build cluster dict from valid traits
    clusters_dict: Dict[int, List[str]] = {}
    for trait, label in zip(valid_traits, cluster_labels):
        if label not in clusters_dict:
            clusters_dict[label] = []
        clusters_dict[label].append(trait)

    # Sort traits within each cluster alphabetically
    for label in clusters_dict:
        clusters_dict[label] = sorted(clusters_dict[label])

    # Add problematic traits as singleton clusters
    next_cluster_id = max(clusters_dict.keys()) + 1 if clusters_dict else 0
    for trait in sorted(problematic_traits):
        clusters_dict[next_cluster_id] = [trait]
        next_cluster_id += 1

    # Renumber clusters to be consecutive starting from 0
    old_to_new = {old: new for new, old in enumerate(sorted(clusters_dict.keys()))}
    result = {old_to_new[old]: traits for old, traits in clusters_dict.items()}

    return result


def select_cluster_representatives(
    df: pd.DataFrame,
    clusters: Dict[int, List[str]],
) -> List[str]:
    """Select one representative trait per cluster.

    Selection criterion: trait with highest variance across genotypes.
    Tie-breaking: alphabetically first trait name.

    Args:
        df: DataFrame with traits as columns and genotypes as rows.
        clusters: Dict mapping cluster_id to list of trait names.
            Output from cluster_correlated_traits().

    Returns:
        List of representative trait names, one per cluster.
        Order corresponds to cluster IDs sorted ascending.

    Example:
        >>> df = pd.DataFrame({'A': [1, 2, 3], 'B': [10, 20, 30], 'C': [5, 5, 5]})
        >>> clusters = {0: ['A', 'B'], 1: ['C']}
        >>> reps = select_cluster_representatives(df, clusters)
        >>> reps  # ['B', 'C'] - B has higher variance than A
    """
    if not clusters:
        return []

    representatives = []

    for cluster_id in sorted(clusters.keys()):
        traits = clusters[cluster_id]

        if len(traits) == 1:
            representatives.append(traits[0])
            continue

        # Calculate variance for each trait, handling NaN
        variances = {}
        for trait in traits:
            var = df[trait].var(skipna=True)
            # If variance is NaN (e.g., all values are NaN), use -inf
            variances[trait] = var if not np.isnan(var) else float("-inf")

        # Find trait(s) with maximum variance
        max_var = max(variances.values())

        # Get all traits with max variance (for tie-breaking)
        max_var_traits = [t for t, v in variances.items() if v == max_var]

        # Tie-break alphabetically
        representative = sorted(max_var_traits)[0]
        representatives.append(representative)

    return representatives
