"""Training functions for scDistill.

Two-phase training:
  Phase 1: PCA + Teacher (Harmony) → Z* (batch-corrected target)
  Phase 2a: Encoder Training (X → Z*)
  Phase 2b: Decoder Training (Z* → X, with batch conditioning)
"""

import os
import sys
import logging
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import scipy.sparse as sp
from typing import Dict, Optional, Union
from tqdm import tqdm

from .encoder import Encoder
from .decoder import Decoder

# Force unbuffered output for real-time progress display
os.environ["PYTHONUNBUFFERED"] = "1"

logger = logging.getLogger(__name__)


def get_device(device: str = 'auto') -> str:
    """Get the best available device.

    Parameters
    ----------
    device : str
        Device specification. 'auto' selects best available.
        Options: 'auto', 'cuda', 'mps', 'cpu'

    Returns
    -------
    str
        Device string for PyTorch.
    """
    if device == 'auto':
        if torch.cuda.is_available():
            return 'cuda'
        elif torch.backends.mps.is_available():
            return 'mps'
        else:
            return 'cpu'

    if device == 'cuda' and not torch.cuda.is_available():
        return 'cpu'
    if device == 'mps' and not torch.backends.mps.is_available():
        return 'cpu'

    return device


class SparseAwareDataset(torch.utils.data.Dataset):
    """Dataset that handles both sparse and dense arrays.

    Converts sparse matrices to dense on-the-fly when fetching batches,
    avoiding loading the entire dataset into memory at once.

    Parameters
    ----------
    *arrays : array-like
        Arrays to include in dataset. Can be np.ndarray or scipy.sparse.
    """

    def __init__(self, *arrays):
        self.arrays = arrays
        self.n_samples = self._get_length(arrays[0])

        # Validate all arrays have same length
        for arr in arrays:
            if self._get_length(arr) != self.n_samples:
                raise ValueError("All arrays must have the same length")

    def _get_length(self, arr) -> int:
        """Get length of array (works for both sparse and dense)."""
        if sp.issparse(arr):
            return arr.shape[0]
        return len(arr)

    def _get_row(self, arr, idx: int):
        """Get row from array, converting sparse to dense if needed."""
        if sp.issparse(arr):
            # For sparse matrices, index returns a sparse matrix, convert to dense
            row = arr.getrow(idx).toarray().flatten()
            return row.astype(np.float32)
        row = arr[idx]
        # Handle 1D arrays (e.g., library_size) - return as 1-element array
        if np.isscalar(row) or (isinstance(row, np.ndarray) and row.ndim == 0):
            return np.array([row], dtype=np.float32)
        return np.asarray(row).astype(np.float32)

    def __len__(self) -> int:
        return self.n_samples

    def __getitem__(self, idx: int):
        return tuple(
            torch.from_numpy(self._get_row(arr, idx))
            for arr in self.arrays
        )


def nb_nll(
    x: torch.Tensor,
    mu: torch.Tensor,
    theta: torch.Tensor,
    eps: float = 1e-8,
) -> torch.Tensor:
    """Negative Binomial negative log-likelihood.

    Parameters
    ----------
    x : torch.Tensor
        Observed counts (N, G).
    mu : torch.Tensor
        Expected counts / mean (N, G).
    theta : torch.Tensor
        Dispersion parameter (G,) or (N, G).
    eps : float
        Small constant for numerical stability.

    Returns
    -------
    nll : torch.Tensor
        Mean negative log-likelihood (scalar).
    """
    if theta.dim() == 1:
        theta = theta.unsqueeze(0)

    mu = mu.clamp(min=eps)
    theta = theta.clamp(min=eps)

    ll = (
        torch.lgamma(x + theta)
        - torch.lgamma(theta)
        - torch.lgamma(x + 1)
        + theta * torch.log(theta / (theta + mu))
        + x * torch.log(mu / (theta + mu + eps))
    )

    return -ll.mean()


def train_encoder(
    encoder: Encoder,
    X_norm: Union[np.ndarray, sp.spmatrix],
    Z_teacher: np.ndarray,
    n_epochs: int = 100,
    batch_size: int = 512,
    lr: float = 1e-3,
    device: str = 'auto',
    verbose: bool = True,
    early_stopping: bool = True,
    patience: int = 10,
    num_workers: int = 0,
) -> Dict[str, list]:
    """Train encoder to distill teacher embeddings.

    X → Encoder → Ẑ
    Loss = MSE(Ẑ, Z*)

    Parameters
    ----------
    encoder : Encoder
        Student encoder to train.
    X_norm : array-like
        Normalized expression (N, G). Can be sparse.
    Z_teacher : np.ndarray
        Batch-corrected teacher embeddings Z* (N, D).
    n_epochs : int
        Training epochs.
    batch_size : int
        Mini-batch size.
    lr : float
        Learning rate.
    device : str
        Device ('auto', 'cuda', 'mps', 'cpu').
    verbose : bool
        Print progress.
    early_stopping : bool
        Enable early stopping.
    patience : int
        Epochs to wait.
    num_workers : int
        DataLoader workers.

    Returns
    -------
    history : Dict
        Training history.
    """
    device = get_device(device)

    encoder = encoder.to(device)
    encoder.train()

    dataset = SparseAwareDataset(X_norm, Z_teacher)
    loader = torch.utils.data.DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers,
        pin_memory=(device != 'cpu'),
    )

    optimizer = torch.optim.Adam(encoder.parameters(), lr=lr)

    history = {'epoch': [], 'loss': []}
    best_loss = float('inf')
    patience_counter = 0
    best_state = None

    if verbose:
        logger.info("Encoder Training (X → Z*)")
        logger.info(f"  Epochs: {n_epochs}, LR: {lr}, Device: {device}")

    # Use stderr for tqdm to properly display progress bar
    epoch_iter = tqdm(range(n_epochs), desc="Encoder", disable=not verbose, file=sys.stderr)

    for epoch in epoch_iter:
        total_loss = 0.0
        n_samples = 0

        for x_norm, z_teacher in loader:
            x_norm = x_norm.to(device)
            z_teacher = z_teacher.to(device)

            optimizer.zero_grad()

            z_pred = encoder(x_norm)
            loss = F.mse_loss(z_pred, z_teacher)

            loss.backward()
            optimizer.step()

            total_loss += loss.item() * len(x_norm)
            n_samples += len(x_norm)

        avg_loss = total_loss / n_samples
        history['epoch'].append(epoch)
        history['loss'].append(avg_loss)
        epoch_iter.set_postfix(loss=f"{avg_loss:.4f}")

        if early_stopping:
            if avg_loss < best_loss - 1e-4:
                best_loss = avg_loss
                patience_counter = 0
                best_state = {k: v.cpu().clone() for k, v in encoder.state_dict().items()}
            else:
                patience_counter += 1
                if patience_counter >= patience:
                    if verbose:
                        logger.info(f"  Early stopping at epoch {epoch + 1}")
                    if best_state:
                        encoder.load_state_dict(best_state)
                    break

    encoder.eval()

    if verbose:
        logger.info(f"  Final loss: {history['loss'][-1]:.4f}")

    return history


def train_decoder(
    encoder: Encoder,
    decoder: Decoder,
    X_norm: Union[np.ndarray, sp.spmatrix],
    X_counts: Union[np.ndarray, sp.spmatrix],
    library_size: np.ndarray,
    batch_labels: Optional[np.ndarray] = None,
    n_epochs: int = 100,
    batch_size: int = 512,
    lr: float = 1e-3,
    device: str = 'auto',
    verbose: bool = True,
    early_stopping: bool = True,
    patience: int = 10,
    num_workers: int = 0,
) -> Dict[str, list]:
    """Train decoder with frozen encoder.

    Ẑ → Decoder → X̂  (Encoder frozen)
    Loss = NB_NLL(X, X̂)

    The Decoder learns to map from batch-corrected Z* space to X.
    Encoder is frozen. Decoder is initialized from scratch.

    Parameters
    ----------
    encoder : Encoder
        Trained student encoder (frozen).
    decoder : Decoder
        New decoder to train from scratch.
    X_norm : array-like
        Normalized expression (N, G). Can be sparse.
    X_counts : array-like
        Raw counts (N, G). Can be sparse.
    library_size : np.ndarray
        Library size per cell (N,).
    batch_labels : np.ndarray, optional
        Batch indices per cell (N,). Required if decoder.use_batch_conditioning=True.
    n_epochs : int
        Training epochs.
    batch_size : int
        Mini-batch size.
    lr : float
        Learning rate.
    device : str
        Device ('auto', 'cuda', 'mps', 'cpu').
    verbose : bool
        Print progress.
    early_stopping : bool
        Enable early stopping.
    patience : int
        Epochs to wait.
    num_workers : int
        DataLoader workers.

    Returns
    -------
    history : Dict
        Training history.
    """
    device = get_device(device)

    encoder = encoder.to(device)
    decoder = decoder.to(device)

    # Freeze encoder
    encoder.eval()
    for param in encoder.parameters():
        param.requires_grad = False

    decoder.train()

    # Check if batch conditioning is enabled
    use_batch_conditioning = getattr(decoder, 'use_batch_conditioning', False)
    if use_batch_conditioning and batch_labels is None:
        raise ValueError("batch_labels required when decoder.use_batch_conditioning=True")

    # Create dataset with or without batch labels
    if batch_labels is not None:
        dataset = SparseAwareDataset(X_norm, X_counts, library_size, batch_labels)
    else:
        dataset = SparseAwareDataset(X_norm, X_counts, library_size)

    loader = torch.utils.data.DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers,
        pin_memory=(device != 'cpu'),
    )

    # Only optimize decoder
    optimizer = torch.optim.Adam(decoder.parameters(), lr=lr)

    history = {'epoch': [], 'loss': []}
    best_loss = float('inf')
    patience_counter = 0
    best_state = None

    if verbose:
        logger.info("Decoder Training (Encoder frozen)")
        logger.info(f"  Epochs: {n_epochs}, LR: {lr}, Device: {device}")

    # Use stderr for tqdm to properly display progress bar
    epoch_iter = tqdm(range(n_epochs), desc="Decoder", disable=not verbose, file=sys.stderr)

    for epoch in epoch_iter:
        total_loss = 0.0
        n_samples = 0

        for batch_data in loader:
            # Unpack batch data (with or without batch_labels)
            if batch_labels is not None:
                x_norm, x_counts, lib_size, batch_ids = batch_data
                batch_ids = batch_ids.to(device).squeeze(-1).long()
            else:
                x_norm, x_counts, lib_size = batch_data
                batch_ids = None

            x_norm = x_norm.to(device)
            x_counts = x_counts.to(device)
            lib_size = lib_size.to(device).squeeze(-1)  # (batch, 1) -> (batch,)

            optimizer.zero_grad()

            # Get batch-corrected Z from frozen encoder
            with torch.no_grad():
                z = encoder(x_norm)

            # Decode with batch conditioning
            mu, theta = decoder.forward_nb(z, lib_size, batch=batch_ids)
            loss = nb_nll(x_counts, mu, theta)

            loss.backward()
            optimizer.step()

            total_loss += loss.item() * len(x_norm)
            n_samples += len(x_norm)

        avg_loss = total_loss / n_samples
        history['epoch'].append(epoch)
        history['loss'].append(avg_loss)
        epoch_iter.set_postfix(loss=f"{avg_loss:.4f}")

        if early_stopping:
            if avg_loss < best_loss - 1e-4:
                best_loss = avg_loss
                patience_counter = 0
                best_state = {k: v.cpu().clone() for k, v in decoder.state_dict().items()}
            else:
                patience_counter += 1
                if patience_counter >= patience:
                    if verbose:
                        logger.info(f"  Early stopping at epoch {epoch + 1}")
                    if best_state:
                        decoder.load_state_dict(best_state)
                    break

    decoder.eval()

    if verbose:
        logger.info(f"  Final loss: {history['loss'][-1]:.4f}")

    return history
