"""scDistill: Two-phase batch correction via knowledge distillation.

Architecture:
    Phase 1: PCA + Teacher (Harmony) → Z* (batch-corrected teacher signal)
    Phase 2a: Train Encoder: X → Z*
    Phase 2b: Train Decoder: Z* → X (with batch conditioning)

Example
-------
>>> from scdistill import Distiller
>>> model = Distiller(adata, batch_key="batch")
>>> model.train()
>>> Z = model.get_latent_representation()
>>> X_corrected = model.get_corrected_expression()

>>> # With different teacher
>>> from scdistill.teachers import CombatTeacher
>>> model = Distiller(adata, batch_key="batch", teacher=CombatTeacher())
>>> model.train()
"""

import logging
import torch
import numpy as np
import pandas as pd
import scipy.sparse as sp
from typing import Optional, Literal, List, Union, TYPE_CHECKING
from anndata import AnnData

if TYPE_CHECKING:
    from .de import PseudobulkConfig, BayesianConfig

from .nn.encoder import Encoder

logger = logging.getLogger(__name__)
from .nn.decoder import Decoder
from .nn.training import (
    get_device,
    train_encoder,
    train_decoder,
)
from .nn.inference import get_corrected_expression as _get_corrected_expression
from .teachers.base import BaseTeacher
from .teachers import HarmonyTeacher


class Distiller:
    """Two-phase batch correction via knowledge distillation.

    Phase 1: PCA + Teacher (Harmony) to get batch-corrected Z*
    Phase 2: Train Encoder (X → Z*) and Decoder (Z* → X)

    Parameters
    ----------
    adata : AnnData
        AnnData object with expression data.
    batch_key : str
        Key in adata.obs for batch labels.
    teacher : BaseTeacher, optional
        Teacher method for batch correction.
        Default is HarmonyTeacher().
    n_latent : int, default=50
        Dimension of latent space.
    n_hidden : int, default=128
        Hidden layer size in MLP.
    n_layers : int, default=2
        Number of hidden layers.
    dropout : float, default=0.1
        Dropout rate.
    device : str, default='auto'
        Device for training ('auto', 'cuda', 'mps', 'cpu').
    pca_method : str, default='standard'
        PCA method: 'standard' (log2CPM+1 → PCA) or 'pearson_residuals'
        (Pearson residuals normalization → PCA, better for count data).
    use_batch_conditioning : bool, default=True
        If True, use conditional decoder that takes batch information as input.
        This allows proper batch correction at inference time by specifying
        target_batch. Recommended for most use cases.
    n_batch_embedding : int, default=10
        Dimension of batch embedding in decoder. Uses nn.Embedding for
        scalability when n_batches is large (atlas integration scenarios).

    Examples
    --------
    >>> # Default teacher (HarmonyTeacher)
    >>> model = Distiller(adata, batch_key="batch")
    >>> model.train()

    >>> # Use ComBat
    >>> from scdistill.teachers import CombatTeacher
    >>> model = Distiller(adata, batch_key="batch", teacher=CombatTeacher())

    >>> # Custom Harmony configuration
    >>> from scdistill.teachers import HarmonyTeacher
    >>> teacher = HarmonyTeacher(theta=1.5, covariate_key="condition")
    >>> model = Distiller(adata, batch_key="batch", teacher=teacher)

    >>> # Use Pearson Residuals PCA (better for ZINB/count data)
    >>> model = Distiller(adata, batch_key="batch", pca_method="pearson_residuals")
    """

    def __init__(
        self,
        adata: AnnData,
        batch_key: str,
        teacher: Optional[BaseTeacher] = None,
        n_latent: int = 50,
        n_hidden: int = 128,
        n_layers: int = 2,
        dropout: float = 0.1,
        device: str = 'auto',
        pca_method: Literal["standard", "pearson_residuals"] = "standard",
        use_batch_conditioning: bool = True,
        n_batch_embedding: int = 10,
    ):
        if batch_key not in adata.obs:
            raise ValueError(f"batch_key '{batch_key}' not in adata.obs")

        self.adata = adata
        self.batch_key = batch_key
        self.n_latent = n_latent
        self.n_hidden = n_hidden
        self.n_layers = n_layers
        self.dropout = dropout
        self.pca_method = pca_method
        self.use_batch_conditioning = use_batch_conditioning
        self.n_batch_embedding = n_batch_embedding

        # Initialize teacher (default: HarmonyTeacher)
        self.teacher = teacher if teacher is not None else HarmonyTeacher()

        # Device selection (supports cuda, mps, cpu)
        self.device = get_device(device)

        # Models (student)
        self.encoder: Optional[Encoder] = None
        self.decoder: Optional[Decoder] = None

        # Internal state (can hold sparse matrices)
        self._X_input: Optional[Union[np.ndarray, sp.spmatrix]] = None
        self._X_counts: Optional[Union[np.ndarray, sp.spmatrix]] = None
        self._library_size: Optional[np.ndarray] = None
        self._gene_names: Optional[np.ndarray] = None
        self._batch_labels: Optional[np.ndarray] = None
        self._Z_teacher: Optional[np.ndarray] = None  # Z* from Teacher
        self.is_trained: bool = False
        self.history: dict = {}

    def _prepare_input_data(self, adata: AnnData) -> tuple:
        """Prepare input data: raw counts → log2(CPM + 1).

        Keeps sparse matrices as sparse for memory efficiency.
        Conversion to dense happens on-the-fly during training.
        """
        if 'counts' in adata.layers:
            X_raw = adata.layers['counts']
        else:
            X_raw = adata.X

        # Keep sparse if input is sparse
        is_sparse = sp.issparse(X_raw)

        if is_sparse:
            # Convert to CSR for efficient row operations
            X_counts = sp.csr_matrix(X_raw, dtype=np.float32)
            # Library size calculation works on sparse
            library_size = np.array(X_counts.sum(axis=1)).flatten().astype(np.float32)
            # For sparse: normalize row-wise
            # Use diag matrix multiplication for efficiency
            norm_factors = sp.diags(1e6 / (library_size + 1e-8))
            X_cpm = norm_factors @ X_counts
            # log2(X + 1) = log1p(X) / log(2)
            # Copy data to avoid modifying original, apply log1p in-place
            X_norm = X_cpm.copy()
            X_norm.data = np.log1p(X_norm.data) / np.log(2)
            # Ensure CSR format
            X_norm = sp.csr_matrix(X_norm, dtype=np.float32)
            X_counts = sp.csr_matrix(X_counts, dtype=np.float32)
        else:
            X_counts = np.asarray(X_raw).astype(np.float32)
            library_size = X_counts.sum(axis=1).astype(np.float32)
            X_cpm = X_counts / (library_size[:, np.newaxis] + 1e-8) * 1e6
            X_norm = np.log2(X_cpm + 1).astype(np.float32)

        return X_norm, X_counts, library_size

    def _to_dense(self, X: Union[np.ndarray, sp.spmatrix]) -> np.ndarray:
        """Convert to dense array if sparse."""
        if sp.issparse(X):
            return X.toarray()
        return np.asarray(X)

    def _apply_teacher(
        self,
        Z: np.ndarray,
        batch_labels: np.ndarray,
        verbose: bool = True,
    ) -> np.ndarray:
        """Apply Teacher to get Z*."""
        # Convert batch labels to string for compatibility
        batch_str = np.array([f"batch_{b}" for b in batch_labels])

        # Call teacher's fit_transform
        Z_corrected = self.teacher.fit_transform(
            X_pca=Z,
            batch_labels=batch_str,
            obs=self.adata.obs,
            verbose=verbose,
            adata=self.adata,
            batch_key=self.batch_key,
        )

        return Z_corrected  # (N, n_latent)

    def train(
        self,
        n_epochs_encoder: int = 100,
        n_epochs_decoder: int = 100,
        lr: float = 1e-3,
        batch_size: int = 512,
        early_stopping: bool = True,
        patience: int = 10,
        verbose: bool = True,
    ) -> "Distiller":
        """Train the 3-step model.

        Parameters
        ----------
        n_epochs_encoder : int, default=100
            Epochs for encoder training.
        n_epochs_decoder : int, default=100
            Epochs for decoder training.
        lr : float, default=1e-3
            Learning rate.
        batch_size : int, default=512
            Mini-batch size.
        early_stopping : bool, default=True
            Enable early stopping.
        patience : int, default=10
            Epochs to wait for improvement.
        verbose : bool, default=True
            Print progress.

        Returns
        -------
        self : Distiller
            Returns self to enable method chaining (e.g., model.train().save()).
        """
        # === Prepare data ===
        if verbose:
            logger.info("=" * 60)
            logger.info("Preparing data...")
            logger.info("=" * 60)

        self._gene_names = self.adata.var_names.values
        self._X_input, self._X_counts, self._library_size = self._prepare_input_data(self.adata)
        self._batch_labels = pd.Categorical(self.adata.obs[self.batch_key]).codes

        # Get shape (works for both sparse and dense)
        n_cells = self._X_input.shape[0]
        n_genes = self._X_input.shape[1]

        if verbose:
            is_sparse = sp.issparse(self._X_input)
            logger.info(f"  Cells: {n_cells}")
            logger.info(f"  Genes: {n_genes}")
            logger.info(f"  Batches: {len(np.unique(self._batch_labels))}")
            logger.info(f"  Sparse: {is_sparse}")
            logger.info(f"  Device: {self.device}")

        # === Phase 1: PCA + Teacher (Harmony) ===
        if verbose:
            logger.info("=" * 60)
            logger.info("Phase 1: Computing PCA and applying Teacher")
            logger.info("=" * 60)
            logger.info(f"  PCA method: {self.pca_method}")

        # Compute PCA (standard or Pearson residuals)
        if self.pca_method == "pearson_residuals":
            # Pearson Residuals PCA (better for count data)
            import scanpy as sc
            adata_tmp = self.adata.copy()
            # Ensure we have raw counts
            if 'counts' in adata_tmp.layers:
                adata_tmp.X = adata_tmp.layers['counts']
            # Normalize using Pearson residuals
            sc.experimental.pp.normalize_pearson_residuals(adata_tmp)
            # PCA on Pearson residuals
            sc.pp.pca(adata_tmp, n_comps=self.n_latent, random_state=42)
            Z_pca = adata_tmp.obsm['X_pca']
            var_explained = adata_tmp.uns['pca']['variance_ratio'].sum()
        else:
            # Standard PCA on log2(CPM+1) normalized data
            from sklearn.decomposition import PCA
            X_dense = self._to_dense(self._X_input)
            pca = PCA(n_components=self.n_latent, random_state=42)
            Z_pca = pca.fit_transform(X_dense)
            var_explained = pca.explained_variance_ratio_.sum()

        if verbose:
            logger.info(f"  PCA variance explained: {var_explained:.3f}")
            logger.info(f"  Z_pca shape: {Z_pca.shape}")

        # Apply Teacher (Harmony) to PCA
        if verbose:
            logger.info(f"  Applying {self.teacher.name}...")

        Z_harmony = self._apply_teacher(Z_pca, self._batch_labels, verbose=verbose)

        # Apply per-sample LayerNorm-style normalization to Z_teacher
        # This matches the encoder's LayerNorm output format
        eps = 1e-5
        Z_mean = Z_harmony.mean(axis=1, keepdims=True)
        Z_std = Z_harmony.std(axis=1, keepdims=True)
        self._Z_teacher = (Z_harmony - Z_mean) / (Z_std + eps)

        # Update n_latent to match teacher output dimension
        # Some teachers (Symphony, LIGER) may output different dimensions
        actual_latent_dim = self._Z_teacher.shape[1]
        if actual_latent_dim != self.n_latent:
            if verbose:
                logger.info(f"  Teacher output dim ({actual_latent_dim}) differs from n_latent ({self.n_latent})")
                logger.info(f"  Adjusting n_latent to {actual_latent_dim}")
            self.n_latent = actual_latent_dim

        if verbose:
            logger.info(f"  Z* shape: {self._Z_teacher.shape}")
            logger.info(f"  Z* mean={self._Z_teacher.mean():.4f}, std={self._Z_teacher.std():.4f}")

        # === Phase 2a: Encoder Training ===
        if verbose:
            logger.info("=" * 60)
            logger.info("Phase 2a: Encoder Training (X → Z*)")
            logger.info("=" * 60)

        self.encoder = Encoder(
            n_input=n_genes,
            n_latent=self.n_latent,
            n_hidden=self.n_hidden,
            n_layers=self.n_layers,
            dropout=self.dropout,
        )

        history_encoder = train_encoder(
            encoder=self.encoder,
            X_norm=self._X_input,
            Z_teacher=self._Z_teacher,
            n_epochs=n_epochs_encoder,
            batch_size=batch_size,
            lr=lr,
            device=self.device,
            verbose=verbose,
            early_stopping=early_stopping,
            patience=patience,
        )

        # === Phase 2b: Decoder Training ===
        if verbose:
            logger.info("=" * 60)
            logger.info("Phase 2b: Decoder Training (Z* → X)")
            logger.info("=" * 60)

        # Create NEW decoder with batch conditioning
        n_batches = len(np.unique(self._batch_labels))
        self.decoder = Decoder(
            n_latent=self.n_latent,
            n_output=n_genes,
            n_hidden=self.n_hidden,
            n_layers=self.n_layers,
            dropout=self.dropout,
            n_batches=n_batches,
            n_batch_embedding=self.n_batch_embedding,
            use_batch_conditioning=self.use_batch_conditioning,
        )

        history_decoder = train_decoder(
            encoder=self.encoder,
            decoder=self.decoder,
            X_norm=self._X_input,
            X_counts=self._X_counts,
            library_size=self._library_size,
            batch_labels=self._batch_labels if self.use_batch_conditioning else None,
            n_epochs=n_epochs_decoder,
            batch_size=batch_size,
            lr=lr,
            device=self.device,
            verbose=verbose,
            early_stopping=early_stopping,
            patience=patience,
        )

        # Store history
        self.history = {
            'encoder': history_encoder,
            'decoder': history_decoder,
        }
        self.is_trained = True

        if verbose:
            logger.info("=" * 60)
            logger.info("Training Complete!")
            logger.info("=" * 60)
            logger.info(f"  Encoder final loss: {history_encoder['loss'][-1]:.4f}")
            logger.info(f"  Decoder final loss: {history_decoder['loss'][-1]:.4f}")

        return self

    def get_teacher_representation(self) -> np.ndarray:
        """Get Z* (Teacher-corrected latent from Phase 1)."""
        if self._Z_teacher is None:
            raise ValueError("Model not trained. Run train() first.")
        return self._Z_teacher.copy()

    def get_latent_representation(
        self,
        cell_idx: Optional[np.ndarray] = None,
        batch_size: int = 2048,
    ) -> np.ndarray:
        """Get batch-corrected latent Z from student Encoder.

        Parameters
        ----------
        cell_idx : np.ndarray, optional
            Indices of cells to process. Default is all cells.
        batch_size : int, default=2048
            Batch size for inference (for memory efficiency).

        Returns
        -------
        Z : np.ndarray
            Latent representation (n_cells, n_latent).
        """
        if not self.is_trained:
            raise ValueError("Model not trained. Run train() first.")

        self.encoder.to(self.device)
        self.encoder.eval()

        if cell_idx is None:
            X_input = self._X_input
            n_cells = X_input.shape[0]
        else:
            X_input = self._X_input[cell_idx]
            n_cells = len(cell_idx)

        use_cuda = self.device != "cpu"

        # Pre-allocate result array
        result = np.empty((n_cells, self.n_latent), dtype=np.float32)

        with torch.inference_mode():
            for i in range(0, n_cells, batch_size):
                end_idx = min(i + batch_size, n_cells)
                X_batch = self._to_dense(X_input[i:end_idx])

                # Transfer to device (non_blocking for GPU)
                X_tensor = torch.from_numpy(np.asarray(X_batch, dtype=np.float32))
                X_tensor = X_tensor.to(self.device, non_blocking=use_cuda)

                Z_batch = self.encoder(X_tensor)
                result[i:end_idx] = Z_batch.cpu().numpy()

        return result

    def get_corrected_expression(
        self,
        cell_idx: Optional[np.ndarray] = None,
        genes: Optional[List[str]] = None,
        batch_size: int = 2048,
        target_batch: Optional[int] = None,
    ) -> np.ndarray:
        """Get batch-corrected expression from Decoder.

        Returns expression in log2(CPM+1) scale, matching the input scale.

        Parameters
        ----------
        cell_idx : np.ndarray, optional
            Indices of cells to process.
        genes : list of str, optional
            Gene names to return. Default is all genes.
        batch_size : int, default=2048
            Batch size for inference.
        target_batch : int, optional
            Target batch index for correction. If None, uses first batch (0).
            All cells will be decoded as if they belong to this batch,
            effectively removing batch effects.

        Returns
        -------
        X_corrected : np.ndarray
            Batch-corrected expression in log2(CPM+1) scale.
        """
        if not self.is_trained:
            raise ValueError("Model not trained. Run train() first.")

        # Select cells
        if cell_idx is None:
            X_input = self._X_input
        else:
            X_input = self._X_input[cell_idx]

        # Default target batch
        if target_batch is None:
            target_batch = 0

        # Call inference function
        result = _get_corrected_expression(
            encoder=self.encoder,
            decoder=self.decoder,
            X_input=X_input,
            use_batch_conditioning=self.use_batch_conditioning,
            target_batch=target_batch,
            device=self.device,
            batch_size=batch_size,
        )

        # Filter genes if requested
        if genes is not None:
            gene_idx = [np.where(self._gene_names == g)[0][0] for g in genes if g in self._gene_names]
            result = result[:, gene_idx]

        return result

    def get_gene_names(self) -> np.ndarray:
        """Get gene names."""
        if self._gene_names is None:
            raise ValueError("Model not prepared. Run train() first.")
        return self._gene_names.copy()

    def differential_expression(
        self,
        groupby: str,
        group1: str,
        group2: str,
        sample_key: Optional[str] = None,
        how: Optional[Union["PseudobulkConfig", "BayesianConfig"]] = None,
        fdr_threshold: float = 0.05,
        idx: Optional[np.ndarray] = None,
    ) -> pd.DataFrame:
        """Differential expression analysis.

        Wrapper around scdistill.de.differential_expression that provides
        convenient access to both pseudobulk and Bayesian DE methods.

        LFC Convention: LFC = log2(group1 / group2)
        - Positive LFC: higher in group1
        - Negative LFC: higher in group2

        Parameters
        ----------
        groupby : str
            Column in adata.obs to group cells.
        group1 : str
            Target group (numerator in fold change).
        group2 : str
            Reference group (denominator in fold change).
        sample_key : str, optional
            Column for sample/donor IDs. Required for pseudobulk method.
        how : PseudobulkConfig or BayesianConfig, optional
            DE method configuration. If None, uses PseudobulkConfig().
        fdr_threshold : float, default=0.05
            FDR threshold for DE calling.
        idx : np.ndarray, optional
            Cell indices subset. If None, uses all cells.

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

        Examples
        --------
        >>> # Pseudobulk DE (recommended)
        >>> de_results = model.differential_expression(
        ...     groupby="condition",
        ...     group1="treated",
        ...     group2="control",
        ...     sample_key="patient_id",
        ... )

        >>> # Bayesian MC Dropout DE
        >>> from scdistill.de import BayesianConfig
        >>> de_results = model.differential_expression(
        ...     groupby="condition",
        ...     group1="treated",
        ...     group2="control",
        ...     how=BayesianConfig(n_samples=100),
        ... )
        """
        from scdistill.de import differential_expression as de_func
        from scdistill.de import PseudobulkConfig, BayesianConfig

        if not self.is_trained:
            raise ValueError("Model not trained. Run train() first.")

        # Get cell indices
        if idx is None:
            idx = np.arange(len(self.adata))

        # Get group labels
        group_labels = self.adata.obs[groupby].values[idx]

        # Get indices for each group
        mask1 = group_labels == group1
        mask2 = group_labels == group2
        idx1 = idx[mask1]
        idx2 = idx[mask2]

        # Build config object
        if how is None:
            config = PseudobulkConfig()
        elif isinstance(how, (PseudobulkConfig, BayesianConfig)):
            config = how
        else:
            raise TypeError(
                f"how must be PseudobulkConfig, BayesianConfig, or None. "
                f"Got {type(how).__name__}"
            )

        # Get sample labels for pseudobulk
        if isinstance(config, PseudobulkConfig):
            if sample_key is None:
                raise ValueError("sample_key is required for pseudobulk DE")
            sample_labels = self.adata.obs[sample_key].values
        else:
            sample_labels = None

        # Call de.differential_expression
        return de_func(
            idx1=idx1,
            idx2=idx2,
            gene_names=self._gene_names,
            encoder=self.encoder,
            decoder=self.decoder,
            X_input=self._X_input,
            sample_labels=sample_labels,
            how=config,
            fdr_threshold=fdr_threshold,
            device=self.device,
            use_batch_conditioning=self.use_batch_conditioning,
            target_batch=0,
        )

    def save(self, dir_path: str, overwrite: bool = False):
        """Save the model."""
        import os
        import pickle

        if os.path.exists(dir_path) and not overwrite:
            raise FileExistsError(f"{dir_path} exists. Use overwrite=True.")

        os.makedirs(dir_path, exist_ok=True)

        torch.save(self.encoder.state_dict(), os.path.join(dir_path, 'encoder.pt'))
        torch.save(self.decoder.state_dict(), os.path.join(dir_path, 'decoder.pt'))

        config = {
            'n_latent': self.n_latent,
            'n_hidden': self.n_hidden,
            'n_layers': self.n_layers,
            'dropout': self.dropout,
            'batch_key': self.batch_key,
            'gene_names': self._gene_names.tolist(),
            'pca_method': self.pca_method,
            'use_batch_conditioning': self.use_batch_conditioning,
            'n_batch_embedding': self.n_batch_embedding,
            'n_batches': len(np.unique(self._batch_labels)) if self._batch_labels is not None else 1,
        }
        with open(os.path.join(dir_path, 'config.pkl'), 'wb') as f:
            pickle.dump(config, f)

        # Save Z_teacher for reference
        np.save(os.path.join(dir_path, 'Z_teacher.npy'), self._Z_teacher)

        logger.info(f"Model saved to {dir_path}")

    @classmethod
    def load(cls, dir_path: str, adata: Optional[AnnData] = None) -> "Distiller":
        """Load a saved model."""
        import os
        import pickle

        with open(os.path.join(dir_path, 'config.pkl'), 'rb') as f:
            config = pickle.load(f)

        gene_names = np.array(config['gene_names'])
        n_genes = len(gene_names)

        if adata is None:
            adata = AnnData(np.zeros((1, n_genes)))
            adata.var_names = gene_names
            adata.obs['batch'] = ['A']

        # Get batch conditioning settings (default False for backwards compatibility)
        use_batch_conditioning = config.get('use_batch_conditioning', False)
        n_batch_embedding = config.get('n_batch_embedding', 10)
        n_batches = config.get('n_batches', 1)

        instance = cls(
            adata=adata,
            batch_key=config.get('batch_key', 'batch'),
            n_latent=config['n_latent'],
            n_hidden=config['n_hidden'],
            n_layers=config['n_layers'],
            dropout=config['dropout'],
            pca_method=config.get('pca_method', 'standard'),
            use_batch_conditioning=use_batch_conditioning,
            n_batch_embedding=n_batch_embedding,
        )

        instance._gene_names = gene_names
        instance._X_input, instance._X_counts, instance._library_size = instance._prepare_input_data(adata)
        instance._batch_labels = pd.Categorical(adata.obs[config.get('batch_key', 'batch')]).codes

        instance.encoder = Encoder(
            n_input=n_genes,
            n_latent=config['n_latent'],
            n_hidden=config['n_hidden'],
            n_layers=config['n_layers'],
            dropout=config['dropout'],
        )
        instance.decoder = Decoder(
            n_latent=config['n_latent'],
            n_output=n_genes,
            n_hidden=config['n_hidden'],
            n_layers=config['n_layers'],
            dropout=config['dropout'],
            n_batches=n_batches,
            n_batch_embedding=n_batch_embedding,
            use_batch_conditioning=use_batch_conditioning,
        )

        instance.encoder.load_state_dict(
            torch.load(os.path.join(dir_path, 'encoder.pt'), map_location='cpu', weights_only=True)
        )
        instance.decoder.load_state_dict(
            torch.load(os.path.join(dir_path, 'decoder.pt'), map_location='cpu', weights_only=True)
        )

        instance._Z_teacher = np.load(os.path.join(dir_path, 'Z_teacher.npy'))

        instance.is_trained = True

        logger.info(f"Model loaded from {dir_path}")
        return instance

    def __repr__(self) -> str:
        status = "trained" if self.is_trained else "not trained"
        return f"Distiller(teacher={self.teacher.name}, n_latent={self.n_latent}, {status})"
