"""Differential Expression Analysis.

This module provides DE analysis methods that work with trained scDistill models:

1. **pseudobulk** (default): Sample-level pseudobulk DE analysis with proper
   statistical testing. Uses corrected expression from the decoder.

2. **bayesian**: Bayesian DE detection using MC Dropout on the encoder/decoder.
   Provides uncertainty quantification through posterior sampling.

Example
-------
>>> # Using Distiller's wrapper (recommended)
>>> de_results = model.differential_expression(
...     groupby="condition",
...     group1="treated",
...     group2="control",
...     sample_key="patient_id",
...     method="pseudobulk",
... )

>>> # With config objects
>>> from scdistill.de import differential_expression, PseudobulkConfig
>>> de_results = differential_expression(
...     idx1=treated_idx, idx2=control_idx,
...     gene_names=gene_names,
...     encoder=model.encoder, decoder=model.decoder,
...     X_input=X_input, sample_labels=sample_labels,
...     how=PseudobulkConfig(method='deseq2'),
... )

References
----------
- DESeq2: Anders & Huber (2010)
- scVI differential expression guide
"""

import logging
from dataclasses import dataclass
from typing import Optional, Literal, Union

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from scipy.stats import ttest_ind, false_discovery_control, ranksums

from ..nn.inference import get_corrected_expression as _get_corrected_expression

logger = logging.getLogger(__name__)


# =============================================================================
# Configuration Classes
# =============================================================================

@dataclass
class PseudobulkConfig:
    """Configuration for pseudobulk DE analysis.

    Parameters
    ----------
    method : str, default='deseq2'
        Statistical test method:
        - 'deseq2': DESeq2-style negative binomial test (recommended, requires pydeseq2)
        - 'wilcox': Wilcoxon rank-sum test (robust, non-parametric)
        - 'ttest': Student's t-test (parametric)
    min_cells_per_sample : int, default=10
        Minimum cells required per sample to include in analysis.
    center_lfc : bool, default=False
        If True, center LFC values by subtracting the median LFC
        of non-significant genes.
    """
    method: Literal['deseq2', 'wilcox', 'ttest'] = 'deseq2'
    min_cells_per_sample: int = 10
    center_lfc: bool = False


@dataclass
class BayesianConfig:
    """Configuration for Bayesian DE analysis using MC Dropout.

    Parameters
    ----------
    n_samples : int, default=100
        Number of MC Dropout samples for posterior estimation.
    delta : float, optional
        LFC threshold for DE probability calculation.
        If None, auto-estimated using GMM.
    use_data_variance : bool, default=True
        If True, incorporate cell-level variance into LFC uncertainty.
    batch_size : int, default=512
        Batch size for processing cells through the model.
    target_batch : int, default=0
        Target batch index for batch correction during DE.
    """
    n_samples: int = 100
    delta: Optional[float] = None
    use_data_variance: bool = True
    batch_size: int = 512
    target_batch: int = 0


# =============================================================================
# Main API Function
# =============================================================================

def differential_expression(
    idx1: np.ndarray,
    idx2: np.ndarray,
    gene_names: np.ndarray,
    encoder: nn.Module,
    decoder: nn.Module,
    X_input: np.ndarray,
    sample_labels: Optional[np.ndarray] = None,
    how: Optional[Union[PseudobulkConfig, BayesianConfig]] = None,
    fdr_threshold: float = 0.05,
    device: str = 'cpu',
    use_batch_conditioning: bool = True,
    target_batch: int = 0,
) -> pd.DataFrame:
    """Perform differential expression analysis using trained model.

    Parameters
    ----------
    idx1 : np.ndarray
        Cell indices for treatment/case group (numerator in fold change).
    idx2 : np.ndarray
        Cell indices for control/reference group (denominator in fold change).
    gene_names : np.ndarray
        Gene names array (n_genes,).
    encoder : nn.Module
        Trained encoder model.
    decoder : nn.Module
        Trained decoder model.
    X_input : np.ndarray
        Normalized expression matrix for encoder input (n_cells, n_genes).
    sample_labels : np.ndarray, optional
        Sample ID for each cell (n_cells,). Required for pseudobulk method.
    how : PseudobulkConfig or BayesianConfig, optional
        DE method configuration. If None, uses PseudobulkConfig() (default).
    fdr_threshold : float, default=0.05
        FDR threshold for calling DE genes.
    device : str, default='cpu'
        Device for model inference.
    use_batch_conditioning : bool, default=True
        Whether decoder uses batch conditioning. Should match decoder's setting.
    target_batch : int, default=0
        Target batch index for batch correction.

    Returns
    -------
    results : pd.DataFrame
        DE results. Columns depend on method:

        Pseudobulk: gene, lfc, lfc_std, pval, fdr, is_de, direction
        Bayesian: gene, lfc_mean, lfc_std, lfc_median, lfc_ci_low, lfc_ci_high,
                  proba_de, is_de, delta, scale1, scale2, direction

    Notes
    -----
    LFC Convention:
        LFC = log2(idx1 / idx2) = log2(treatment / control)
        - Positive LFC: higher in treatment (idx1)
        - Negative LFC: lower in treatment (idx2)
    """
    # Default to pseudobulk
    if how is None:
        how = PseudobulkConfig()

    # Override target_batch from config if using BayesianConfig
    if isinstance(how, BayesianConfig):
        target_batch = how.target_batch

    if isinstance(how, PseudobulkConfig):
        if sample_labels is None:
            raise ValueError(
                "sample_labels is required for pseudobulk DE. "
                "Provide the sample/donor ID for each cell."
            )
        return _differential_expression_pseudobulk(
            idx1=idx1,
            idx2=idx2,
            gene_names=gene_names,
            encoder=encoder,
            decoder=decoder,
            X_input=X_input,
            sample_labels=sample_labels,
            fdr_threshold=fdr_threshold,
            device=device,
            use_batch_conditioning=use_batch_conditioning,
            target_batch=target_batch,
            config=how,
        )

    elif isinstance(how, BayesianConfig):
        return _differential_expression_bayesian(
            idx1=idx1,
            idx2=idx2,
            gene_names=gene_names,
            encoder=encoder,
            decoder=decoder,
            X_input=X_input,
            fdr_threshold=fdr_threshold,
            device=device,
            use_batch_conditioning=use_batch_conditioning,
            config=how,
        )

    else:
        raise ValueError(
            f"how must be PseudobulkConfig, BayesianConfig, or None. "
            f"Got {type(how).__name__}"
        )


# =============================================================================
# Internal: Pseudobulk DE
# =============================================================================

def _differential_expression_pseudobulk(
    idx1: np.ndarray,
    idx2: np.ndarray,
    gene_names: np.ndarray,
    encoder: nn.Module,
    decoder: nn.Module,
    X_input: np.ndarray,
    sample_labels: np.ndarray,
    fdr_threshold: float,
    device: str,
    use_batch_conditioning: bool,
    target_batch: int,
    config: PseudobulkConfig,
) -> pd.DataFrame:
    """Pseudobulk DE using corrected expression from decoder."""
    logger.info("=" * 70)
    logger.info("DIFFERENTIAL EXPRESSION (Pseudobulk)")
    logger.info("=" * 70)

    # Get corrected expression from decoder (log2(CPM+1) scale)
    logger.info("  Computing corrected expression from decoder...")
    X_corrected = _get_corrected_expression(
        encoder=encoder,
        decoder=decoder,
        X_input=X_input,
        use_batch_conditioning=use_batch_conditioning,
        target_batch=target_batch,
        device=device,
    )

    # Convert to boolean masks if needed
    n_cells = len(sample_labels)
    if idx1.dtype != bool:
        mask1 = np.zeros(n_cells, dtype=bool)
        mask1[idx1] = True
        idx1 = mask1
    if idx2.dtype != bool:
        mask2 = np.zeros(n_cells, dtype=bool)
        mask2[idx2] = True
        idx2 = mask2

    # Find samples with enough cells
    unique_samples = np.unique(sample_labels)
    samples_group1 = []
    samples_group2 = []

    for sample in unique_samples:
        sample_mask = sample_labels == sample
        n_cells_g1 = np.sum(sample_mask & idx1)
        n_cells_g2 = np.sum(sample_mask & idx2)

        if n_cells_g1 >= config.min_cells_per_sample:
            samples_group1.append(sample)
        if n_cells_g2 >= config.min_cells_per_sample:
            samples_group2.append(sample)

    n_samples_g1 = len(samples_group1)
    n_samples_g2 = len(samples_group2)

    logger.info(f"  Group 1: {n_samples_g1} samples (min {config.min_cells_per_sample} cells)")
    logger.info(f"  Group 2: {n_samples_g2} samples (min {config.min_cells_per_sample} cells)")

    if n_samples_g1 < 2 or n_samples_g2 < 2:
        raise ValueError(
            f"Insufficient samples for DE. Group 1: {n_samples_g1}, Group 2: {n_samples_g2}. "
            f"Need at least 2 samples per group."
        )

    # Fall back to t-test if too few samples
    method = config.method
    if (n_samples_g1 < 3 or n_samples_g2 < 3) and method == 'wilcox':
        logger.warning("  <3 samples per group, falling back to t-test")
        method = 'ttest'

    # Compute pseudobulk
    def compute_pseudobulk(samples, group_mask):
        pseudobulk = []
        for sample in samples:
            mask = (sample_labels == sample) & group_mask
            if mask.sum() > 0:
                counts = np.power(2, X_corrected[mask]) - 1
                pseudobulk.append(counts.mean(axis=0))
        return np.array(pseudobulk)

    logger.info("  Computing pseudobulk...")
    pb_g1 = compute_pseudobulk(samples_group1, idx1)
    pb_g2 = compute_pseudobulk(samples_group2, idx2)

    # Compute LFC
    eps = 1.0
    mean_g1 = pb_g1.mean(axis=0)
    mean_g2 = pb_g2.mean(axis=0)
    lfc = np.log2((mean_g1 + eps) / (mean_g2 + eps))

    # LFC std
    sample_lfc_g1 = np.log2(pb_g1 + eps)
    sample_lfc_g2 = np.log2(pb_g2 + eps)
    lfc_std = np.sqrt(
        np.var(sample_lfc_g2.mean(axis=0) - sample_lfc_g1, axis=0) / n_samples_g1 +
        np.var(sample_lfc_g2 - sample_lfc_g1.mean(axis=0), axis=0) / n_samples_g2
    )

    # P-values
    n_genes = X_corrected.shape[1]
    logger.info(f"  Computing p-values (method={method})...")
    pvals = np.ones(n_genes)

    if method == 'wilcox':
        for g in range(n_genes):
            try:
                _, pvals[g] = ranksums(pb_g1[:, g], pb_g2[:, g])
            except Exception:
                pass

    elif method == 'ttest':
        for g in range(n_genes):
            try:
                _, pvals[g] = ttest_ind(pb_g1[:, g], pb_g2[:, g])
            except Exception:
                pass

    elif method == 'deseq2':
        pvals, lfc = _run_deseq2(pb_g1, pb_g2, gene_names, n_samples_g1, n_samples_g2)

    pvals = np.nan_to_num(pvals, nan=1.0)

    # FDR correction
    try:
        fdr = false_discovery_control(pvals, method='bh')
    except Exception:
        fdr = _bh_correction(pvals)

    # LFC centering
    if config.center_lfc:
        null_mask = fdr >= 0.5
        if null_mask.sum() > 10:
            valid_null = null_mask & ~np.isnan(lfc)
            if valid_null.sum() > 10:
                offset = np.nanmedian(lfc[null_mask])
                lfc = lfc - offset
                logger.info(f"  LFC centering: offset = {offset:.4f}")

    # DE calling (FDR only)
    is_de = fdr < fdr_threshold

    results = pd.DataFrame({
        'gene': gene_names,
        'lfc': lfc,
        'lfc_std': lfc_std,
        'pval': pvals,
        'fdr': fdr,
        'is_de': is_de,
        'direction': np.where(is_de, np.where(lfc > 0, 'up', 'down'), 'none'),
    })

    results = results.sort_values('fdr').reset_index(drop=True)

    n_de = is_de.sum()
    n_up = (is_de & (lfc > 0)).sum()
    n_down = (is_de & (lfc < 0)).sum()

    logger.info(f"  Threshold: FDR < {fdr_threshold}")
    logger.info(f"  DE genes: {n_de} ({n_up} up, {n_down} down)")
    logger.info("=" * 70)

    return results


def _run_deseq2(pb_g1, pb_g2, gene_names, n_g1, n_g2):
    """Run DESeq2 analysis."""
    try:
        from pydeseq2.dds import DeseqDataSet
        from pydeseq2.ds import DeseqStats

        counts = np.vstack([pb_g1, pb_g2]).astype(int)
        counts = np.maximum(counts, 0)

        conditions = ['group1'] * n_g1 + ['group2'] * n_g2
        metadata = pd.DataFrame({'condition': conditions})
        metadata.index = [f'sample_{i}' for i in range(len(conditions))]

        count_df = pd.DataFrame(counts, columns=gene_names)
        count_df.index = metadata.index

        dds = DeseqDataSet(counts=count_df, metadata=metadata, design_factors='condition')
        dds.deseq2()

        stat_res = DeseqStats(dds, contrast=['condition', 'group1', 'group2'])
        stat_res.summary()

        results = stat_res.results_df
        return results['pvalue'].values, results['log2FoldChange'].values

    except ImportError:
        raise ImportError("pydeseq2 required for method='deseq2'. Install: pip install pydeseq2")


def _bh_correction(pvals):
    """Benjamini-Hochberg FDR correction."""
    n = len(pvals)
    sorted_idx = np.argsort(pvals)
    fdr = np.zeros(n)
    for i, idx in enumerate(sorted_idx):
        fdr[idx] = pvals[idx] * n / (i + 1)
    fdr = np.minimum.accumulate(fdr[np.argsort(sorted_idx)[::-1]])[::-1]
    return np.clip(fdr, 0, 1)


# =============================================================================
# Internal: Bayesian MC Dropout DE
# =============================================================================

def _differential_expression_bayesian(
    idx1: np.ndarray,
    idx2: np.ndarray,
    gene_names: np.ndarray,
    encoder: nn.Module,
    decoder: nn.Module,
    X_input: np.ndarray,
    fdr_threshold: float,
    device: str,
    use_batch_conditioning: bool,
    config: BayesianConfig,
) -> pd.DataFrame:
    """Bayesian DE using MC Dropout."""
    import scipy.sparse as sp

    logger.info("=" * 70)
    logger.info("DIFFERENTIAL EXPRESSION (Bayesian MC Dropout)")
    logger.info("=" * 70)
    logger.info(f"  Group 1: {len(idx1)} cells")
    logger.info(f"  Group 2: {len(idx2)} cells")
    logger.info(f"  MC samples: {config.n_samples}")

    n_mc_samples = config.n_samples
    mc_batch_size = config.batch_size
    target_batch = config.target_batch

    def _to_dense(X):
        if sp.issparse(X):
            return X.toarray()
        return X

    def _sample_cell_level_rates(indices):
        """Sample cell-level rates using MC Dropout."""
        encoder.to(device)
        decoder.to(device)
        encoder.train()
        decoder.train()

        X_subset = X_input[indices]
        n_cells = len(indices)
        n_genes = X_input.shape[1]

        all_rates = np.zeros((n_mc_samples, n_cells, n_genes), dtype=np.float32)

        for s in range(n_mc_samples):
            cell_rates = []
            for i in range(0, n_cells, mc_batch_size):
                end_idx = min(i + mc_batch_size, n_cells)
                X_batch = _to_dense(X_subset[i:end_idx])
                X_tensor = torch.FloatTensor(X_batch).to(device)

                if use_batch_conditioning:
                    batch_override = torch.full(
                        (X_tensor.shape[0],), target_batch, dtype=torch.long, device=device
                    )
                else:
                    batch_override = None

                with torch.no_grad():
                    z = encoder(X_tensor)
                    rho = decoder.get_rho(z, batch=batch_override)
                cell_rates.append(rho.cpu().numpy())

            all_rates[s] = np.vstack(cell_rates)

        encoder.eval()
        decoder.eval()
        return all_rates

    def _sample_posterior_rates(indices):
        """Sample group-level rates using MC Dropout."""
        encoder.to(device)
        decoder.to(device)
        encoder.train()
        decoder.train()

        X_subset = X_input[indices]
        n_cells = len(indices)
        all_scales = []

        for _ in range(n_mc_samples):
            rho_sum = None
            count = 0

            for i in range(0, n_cells, mc_batch_size):
                end_idx = min(i + mc_batch_size, n_cells)
                X_batch = _to_dense(X_subset[i:end_idx])
                X_tensor = torch.FloatTensor(X_batch).to(device)

                if use_batch_conditioning:
                    batch_override = torch.full(
                        (X_tensor.shape[0],), target_batch, dtype=torch.long, device=device
                    )
                else:
                    batch_override = None

                with torch.no_grad():
                    z = encoder(X_tensor)
                    rho = decoder.get_rho(z, batch=batch_override)

                if rho_sum is None:
                    rho_sum = rho.sum(dim=0).cpu().numpy()
                else:
                    rho_sum += rho.sum(dim=0).cpu().numpy()
                count += len(X_batch)

            scale = rho_sum / count
            all_scales.append(scale)

        encoder.eval()
        decoder.eval()
        return np.array(all_scales)

    def _compute_lfc_with_data_variance(indices1, indices2):
        """Compute LFC with combined model and data variance."""
        rates1_samples = _sample_cell_level_rates(indices1)
        rates2_samples = _sample_cell_level_rates(indices2)

        mean1_samples = rates1_samples.mean(axis=1)
        mean2_samples = rates2_samples.mean(axis=1)

        model_var1 = mean1_samples.var(axis=0)
        model_var2 = mean2_samples.var(axis=0)

        data_var1 = rates1_samples.var(axis=1).mean(axis=0)
        data_var2 = rates2_samples.var(axis=1).mean(axis=0)

        n1, n2 = len(indices1), len(indices2)
        total_var1 = data_var1 / n1 + model_var1
        total_var2 = data_var2 / n2 + model_var2

        scale1_mean = mean1_samples.mean(axis=0)
        scale2_mean = mean2_samples.mean(axis=0)

        eps = 1e-8
        lfc_mean = np.log2(scale1_mean + eps) - np.log2(scale2_mean + eps)

        ln2_sq = np.log(2) ** 2
        lfc_var = (
            total_var1 / ((scale1_mean + eps) ** 2 * ln2_sq)
            + total_var2 / ((scale2_mean + eps) ** 2 * ln2_sq)
        )
        lfc_std = np.sqrt(lfc_var)

        return lfc_mean, lfc_std, scale1_mean, scale2_mean

    def _estimate_delta(lfc_means, coef=0.6, min_thres=0.3):
        """Estimate LFC threshold using GMM."""
        from sklearn.mixture import GaussianMixture
        gmm = GaussianMixture(n_components=3, random_state=42)
        gmm.fit(lfc_means.reshape(-1, 1))
        modes = np.sort(gmm.means_.flatten())
        delta_val = coef * np.abs(modes[[0, -1]]).mean()
        return max(min_thres, delta_val)

    def _fdr_de_prediction(proba_de, fdr=0.05):
        """Compute posterior expected FDR."""
        sorted_idx = np.argsort(proba_de)[::-1]
        sorted_proba = proba_de[sorted_idx]
        cumulative_fdr = np.cumsum(1.0 - sorted_proba) / (1.0 + np.arange(len(sorted_proba)))
        n_de = np.sum(cumulative_fdr <= fdr)
        is_de = np.zeros(len(proba_de), dtype=bool)
        if n_de > 0:
            is_de[sorted_idx[:n_de]] = True
        return is_de

    # Main computation
    if config.use_data_variance:
        lfc_mean, lfc_std, scale1_mean, scale2_mean = _compute_lfc_with_data_variance(idx1, idx2)
        n_pseudo_samples = 1000
        lfc_samples = np.random.normal(
            loc=lfc_mean,
            scale=lfc_std,
            size=(n_pseudo_samples, len(lfc_mean))
        )
    else:
        scales1 = _sample_posterior_rates(idx1)
        scales2 = _sample_posterior_rates(idx2)

        eps = 1e-8
        lfc_samples = np.log2(scales1 + eps) - np.log2(scales2 + eps)

        lfc_mean = lfc_samples.mean(axis=0)
        lfc_std = lfc_samples.std(axis=0)
        scale1_mean = scales1.mean(axis=0)
        scale2_mean = scales2.mean(axis=0)

    lfc_median = np.median(lfc_samples, axis=0)

    delta_val = config.delta
    if delta_val is None:
        delta_val = _estimate_delta(lfc_mean)
        logger.info(f"  Auto-estimated delta: {delta_val:.3f}")

    is_de_samples = np.abs(lfc_samples) > delta_val
    proba_de = is_de_samples.mean(axis=0)
    is_de = _fdr_de_prediction(proba_de, fdr=fdr_threshold)

    lfc_ci_low = np.percentile(lfc_samples, 2.5, axis=0)
    lfc_ci_high = np.percentile(lfc_samples, 97.5, axis=0)

    direction = np.where(lfc_mean > 0, "up", "down")
    direction = np.where(is_de, direction, "none")

    results = pd.DataFrame({
        "gene": gene_names,
        "lfc_mean": lfc_mean,
        "lfc_std": lfc_std,
        "lfc_median": lfc_median,
        "lfc_ci_low": lfc_ci_low,
        "lfc_ci_high": lfc_ci_high,
        "proba_de": proba_de,
        "is_de": is_de,
        "delta": delta_val,
        "scale1": scale1_mean,
        "scale2": scale2_mean,
        "direction": direction,
    })

    n_de = is_de.sum()
    n_up = (is_de & (lfc_mean > 0)).sum()
    n_down = (is_de & (lfc_mean < 0)).sum()

    logger.info(f"  Delta (LFC threshold): {delta_val:.3f}")
    logger.info(f"  FDR threshold: {fdr_threshold}")
    logger.info(f"  DE genes: {n_de} ({n_up} up, {n_down} down)")
    logger.info("=" * 70)

    return results.sort_values("proba_de", ascending=False)
