"""
WaveDL - Cross-Validation Utilities
====================================

Internal module for K-fold cross-validation. Called by train.py when --cv flag is used.

This module provides:
    - SimpleDataset: In-memory dataset for CV
    - train_fold: Single fold training function
    - run_cross_validation: Main CV orchestration

Author: Ductho Le (ductho.le@outlook.com)
Version: 1.0.0
"""

import json
import logging
import os
import pickle
from datetime import datetime
from typing import Any

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from sklearn.metrics import mean_absolute_error, r2_score
from sklearn.model_selection import KFold, StratifiedKFold
from sklearn.preprocessing import StandardScaler
from torch.utils.data import DataLoader


# ==============================================================================
# SIMPLE DATASET
# ==============================================================================
class CVDataset(torch.utils.data.Dataset):
    """Simple in-memory dataset for cross-validation."""

    def __init__(self, X: np.ndarray, y: np.ndarray, expected_spatial_ndim: int = None):
        """
        Initialize CV dataset with explicit channel dimension handling.

        Args:
            X: Input data with shape (N, *spatial_dims) or (N, C, *spatial_dims)
            y: Target data (N, T)
            expected_spatial_ndim: Expected number of spatial dimensions (1, 2, or 3).
                If provided, uses explicit logic instead of heuristics.
                If None, falls back to ndim-based inference (legacy behavior).

        Channel Dimension Logic:
            - If X.ndim == expected_spatial_ndim + 1: Add channel dim (N, *spatial) -> (N, 1, *spatial)
            - If X.ndim == expected_spatial_ndim + 2: Already has channel (N, C, *spatial)
            - If expected_spatial_ndim is None: Use legacy ndim-based inference
        """
        if expected_spatial_ndim is not None:
            # Explicit mode: use expected_spatial_ndim to determine if channel exists
            if X.ndim == expected_spatial_ndim + 1:
                # Shape is (N, *spatial) - needs channel dimension
                X = np.expand_dims(X, axis=1)
            elif X.ndim == expected_spatial_ndim + 2:
                # Shape is (N, C, *spatial) - already has channel
                pass
            else:
                raise ValueError(
                    f"Input shape {X.shape} incompatible with expected_spatial_ndim={expected_spatial_ndim}. "
                    f"Expected ndim={expected_spatial_ndim + 1} or {expected_spatial_ndim + 2}, got {X.ndim}."
                )
        else:
            # Legacy mode: infer from ndim (for backwards compatibility)
            # Assumes single-channel data without explicit channel dimension
            if X.ndim == 2:  # 1D signals: (N, L) -> (N, 1, L)
                X = X[:, np.newaxis, :]
            elif X.ndim == 3:  # 2D images: (N, H, W) -> (N, 1, H, W)
                X = X[:, np.newaxis, :, :]
            elif X.ndim == 4:  # 3D volumes: (N, D, H, W) -> (N, 1, D, H, W)
                X = X[:, np.newaxis, :, :, :]
            # ndim >= 5 assumed to already have channel dimension

        self.X = torch.tensor(X, dtype=torch.float32)
        self.y = torch.tensor(y, dtype=torch.float32)

    def __len__(self) -> int:
        return len(self.X)

    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
        return self.X[idx], self.y[idx]


# ==============================================================================
# SINGLE FOLD TRAINING
# ==============================================================================
def train_fold(
    fold: int,
    model: nn.Module,
    train_loader: DataLoader,
    val_loader: DataLoader,
    criterion: nn.Module,
    optimizer: torch.optim.Optimizer,
    scheduler,
    device: torch.device,
    epochs: int,
    patience: int,
    scaler: StandardScaler,
    logger: logging.Logger,
) -> dict[str, Any]:
    """
    Train and evaluate a single CV fold.

    Args:
        fold: Fold index (0-based)
        model: PyTorch model
        train_loader: Training data loader
        val_loader: Validation data loader
        criterion: Loss function
        optimizer: Optimizer
        scheduler: LR scheduler
        device: Torch device
        epochs: Max epochs
        patience: Early stopping patience
        scaler: Target scaler (for physical units)
        logger: Logger instance

    Returns:
        Dictionary with fold results and metrics
    """
    best_val_loss = float("inf")
    patience_ctr = 0
    best_state = None
    history = []

    # Determine if scheduler steps per batch (OneCycleLR) or per epoch
    # Use isinstance check since class name 'OneCycleLR' != 'onecycle' string in is_epoch_based
    from torch.optim.lr_scheduler import OneCycleLR

    step_per_batch = isinstance(scheduler, OneCycleLR)

    for epoch in range(epochs):
        # Training
        model.train()
        train_loss = 0.0
        train_samples = 0

        for x, y in train_loader:
            x, y = x.to(device), y.to(device)

            optimizer.zero_grad()
            pred = model(x)
            loss = criterion(pred, y)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()

            # Per-batch LR scheduling (OneCycleLR)
            if step_per_batch:
                scheduler.step()

            train_loss += loss.item() * x.size(0)
            train_samples += x.size(0)

        avg_train_loss = train_loss / train_samples

        # Validation
        model.eval()
        val_loss = 0.0
        val_samples = 0
        all_preds = []
        all_targets = []

        with torch.inference_mode():
            for x, y in val_loader:
                x, y = x.to(device), y.to(device)
                pred = model(x)
                loss = criterion(pred, y)

                val_loss += loss.item() * x.size(0)
                val_samples += x.size(0)

                all_preds.append(pred.cpu())
                all_targets.append(y.cpu())

        avg_val_loss = val_loss / val_samples

        # Compute metrics (guard for tiny datasets)
        y_pred = torch.cat(all_preds).numpy()
        y_true = torch.cat(all_targets).numpy()
        r2 = r2_score(y_true, y_pred) if len(y_true) >= 2 else float("nan")
        mae = np.abs((y_pred - y_true) * scaler.scale_).mean()

        history.append(
            {
                "epoch": epoch + 1,
                "train_loss": avg_train_loss,
                "val_loss": avg_val_loss,
                "r2": r2,
                "mae": mae,
            }
        )

        # LR scheduling (epoch-based only, not for per-batch schedulers)
        if not step_per_batch and hasattr(scheduler, "step"):
            if isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
                scheduler.step(avg_val_loss)
            else:
                scheduler.step()

        # Early stopping
        if avg_val_loss < best_val_loss:
            best_val_loss = avg_val_loss
            patience_ctr = 0
            best_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
        else:
            patience_ctr += 1

        if patience_ctr >= patience:
            logger.info(f"    Fold {fold + 1}: Early stopping at epoch {epoch + 1}")
            break

    # Restore best model and compute final metrics
    if best_state:
        model.load_state_dict(best_state)

    model.eval()
    all_preds = []
    all_targets = []

    with torch.inference_mode():
        for x, y in val_loader:
            x, y = x.to(device), y.to(device)
            pred = model(x)
            all_preds.append(pred.cpu())
            all_targets.append(y.cpu())

    y_pred = torch.cat(all_preds).numpy()
    y_true = torch.cat(all_targets).numpy()

    # Inverse transform for physical units
    y_pred_phys = scaler.inverse_transform(y_pred)
    y_true_phys = scaler.inverse_transform(y_true)

    results = {
        "fold": fold + 1,
        "best_val_loss": best_val_loss,
        "r2": r2_score(y_true, y_pred) if len(y_true) >= 2 else float("nan"),
        "mae_normalized": mean_absolute_error(y_true, y_pred),
        "mae_physical": mean_absolute_error(y_true_phys, y_pred_phys),
        "epochs_trained": len(history),
        "history": history,
    }

    # Per-target metrics (guard for tiny folds)
    for i in range(y_true.shape[1]):
        if len(y_true) >= 2:
            results[f"r2_target_{i}"] = r2_score(y_true[:, i], y_pred[:, i])
        else:
            results[f"r2_target_{i}"] = float("nan")
        results[f"mae_target_{i}"] = mean_absolute_error(
            y_true_phys[:, i], y_pred_phys[:, i]
        )

    return results


# ==============================================================================
# MAIN CV ORCHESTRATION
# ==============================================================================
def run_cross_validation(
    # Data
    X: np.ndarray,
    y: np.ndarray,
    # Model
    model_name: str,
    in_shape: tuple[int, ...],
    out_size: int,
    # CV settings
    folds: int = 5,
    stratify: bool = False,
    stratify_bins: int = 10,
    # Training settings
    batch_size: int = 128,
    lr: float = 1e-3,
    epochs: int = 100,
    patience: int = 20,
    weight_decay: float = 1e-4,
    # Components
    loss_name: str = "mse",
    optimizer_name: str = "adamw",
    scheduler_name: str = "plateau",
    # Output
    output_dir: str = "./cv_results",
    workers: int = 4,
    seed: int = 2025,
    logger: logging.Logger | None = None,
) -> dict[str, Any]:
    """
    Run K-fold cross-validation.

    Args:
        X: Input data
        y: Target data
        model_name: Model architecture name
        in_shape: Input shape (excluding batch and channel)
        out_size: Number of output targets
        folds: Number of CV folds
        stratify: Use stratified splitting
        stratify_bins: Number of bins for stratification
        batch_size: Batch size
        lr: Learning rate
        epochs: Max epochs per fold
        patience: Early stopping patience
        weight_decay: Weight decay
        loss_name: Loss function name
        optimizer_name: Optimizer name
        scheduler_name: Scheduler name
        output_dir: Output directory
        workers: DataLoader workers
        seed: Random seed
        logger: Logger instance

    Returns:
        Summary dictionary with aggregated results
    """
    # Setup
    os.makedirs(output_dir, exist_ok=True)

    if logger is None:
        logging.basicConfig(
            level=logging.INFO,
            format="%(asctime)s | %(levelname)s | %(message)s",
            datefmt="%H:%M:%S",
        )
        logger = logging.getLogger("CV-Trainer")

    # Set seeds
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # Auto-detect optimal DataLoader workers if not specified (matches train.py behavior)
    if workers < 0:
        cpu_count = os.cpu_count() or 4
        num_gpus = torch.cuda.device_count() if torch.cuda.is_available() else 1
        # Heuristic: 4-16 workers per GPU, bounded by available CPU cores
        workers = min(16, max(2, (cpu_count - 2) // max(1, num_gpus)))
        logger.info(
            f"⚙️  Auto-detected workers: {workers} (CPUs: {cpu_count}, GPUs: {num_gpus})"
        )

    logger.info(f"🚀 K-Fold Cross-Validation ({folds} folds)")
    logger.info(f"   Model: {model_name} | Device: {device}")
    logger.info(
        f"   Loss: {loss_name} | Optimizer: {optimizer_name} | Scheduler: {scheduler_name}"
    )
    logger.info(f"   Data shape: X={X.shape}, y={y.shape}")

    # Setup cross-validation
    if stratify:
        # Bin targets for stratification (regression)
        y_binned = np.digitize(
            y[:, 0], np.percentile(y[:, 0], np.linspace(0, 100, stratify_bins + 1))
        )
        kfold = StratifiedKFold(n_splits=folds, shuffle=True, random_state=seed)
        splits = list(kfold.split(X, y_binned))
    else:
        kfold = KFold(n_splits=folds, shuffle=True, random_state=seed)
        splits = list(kfold.split(X))

    # Import factories
    from wavedl.models import build_model
    from wavedl.utils import get_loss, get_optimizer, get_scheduler

    # Run folds
    fold_results = []

    for fold, (train_idx, val_idx) in enumerate(splits):
        logger.info(f"\n{'=' * 60}")
        logger.info(f"📊 Fold {fold + 1}/{folds}")
        logger.info(f"   Train: {len(train_idx)} samples, Val: {len(val_idx)} samples")

        # Split data
        X_train, X_val = X[train_idx], X[val_idx]
        y_train, y_val = y[train_idx], y[val_idx]

        # Fit scaler on training data only
        scaler = StandardScaler()
        y_train_scaled = scaler.fit_transform(y_train)
        y_val_scaled = scaler.transform(y_val)

        # Create datasets and loaders with explicit spatial dimensionality
        spatial_ndim = len(in_shape)
        train_ds = CVDataset(
            X_train.astype(np.float32),
            y_train_scaled.astype(np.float32),
            expected_spatial_ndim=spatial_ndim,
        )
        val_ds = CVDataset(
            X_val.astype(np.float32),
            y_val_scaled.astype(np.float32),
            expected_spatial_ndim=spatial_ndim,
        )

        train_loader = DataLoader(
            train_ds,
            batch_size=batch_size,
            shuffle=True,
            num_workers=workers,
            pin_memory=True,
        )
        val_loader = DataLoader(
            val_ds,
            batch_size=batch_size,
            shuffle=False,
            num_workers=workers,
            pin_memory=True,
        )

        # Build model
        model = build_model(model_name, in_shape=in_shape, out_size=out_size)
        model = model.to(device)

        # Setup training components
        criterion = get_loss(loss_name)
        optimizer = get_optimizer(
            optimizer_name, model.parameters(), lr=lr, weight_decay=weight_decay
        )
        scheduler = get_scheduler(
            scheduler_name,
            optimizer,
            epochs=epochs,
            steps_per_epoch=len(train_loader) if scheduler_name == "onecycle" else None,
        )

        # Train fold
        results = train_fold(
            fold=fold,
            model=model,
            train_loader=train_loader,
            val_loader=val_loader,
            criterion=criterion,
            optimizer=optimizer,
            scheduler=scheduler,
            device=device,
            epochs=epochs,
            patience=patience,
            scaler=scaler,
            logger=logger,
        )

        fold_results.append(results)

        logger.info(
            f"    Fold {fold + 1} Results: R²={results['r2']:.4f}, MAE={results['mae_physical']:.4f}"
        )

        # Save fold model
        fold_dir = os.path.join(output_dir, f"fold_{fold + 1}")
        os.makedirs(fold_dir, exist_ok=True)
        torch.save(model.state_dict(), os.path.join(fold_dir, "model.pth"))
        with open(os.path.join(fold_dir, "scaler.pkl"), "wb") as f:
            pickle.dump(scaler, f)

    # ==============================================================================
    # AGGREGATE RESULTS
    # ==============================================================================
    logger.info(f"\n{'=' * 60}")
    logger.info("📈 Cross-Validation Summary")
    logger.info("=" * 60)

    r2_scores = [r["r2"] for r in fold_results]
    mae_scores = [r["mae_physical"] for r in fold_results]
    val_losses = [r["best_val_loss"] for r in fold_results]

    summary = {
        "config": {
            "model": model_name,
            "folds": folds,
            "stratify": stratify,
            "stratify_bins": stratify_bins,
            "batch_size": batch_size,
            "lr": lr,
            "epochs": epochs,
            "patience": patience,
            "loss": loss_name,
            "optimizer": optimizer_name,
            "scheduler": scheduler_name,
        },
        "timestamp": datetime.now().isoformat(),
        "folds": folds,
        "r2_mean": float(np.mean(r2_scores)),
        "r2_std": float(np.std(r2_scores)),
        "mae_mean": float(np.mean(mae_scores)),
        "mae_std": float(np.std(mae_scores)),
        "val_loss_mean": float(np.mean(val_losses)),
        "val_loss_std": float(np.std(val_losses)),
        "fold_results": fold_results,
    }

    logger.info(f"   R² Score:    {summary['r2_mean']:.4f} ± {summary['r2_std']:.4f}")
    logger.info(f"   MAE (phys):  {summary['mae_mean']:.4f} ± {summary['mae_std']:.4f}")
    logger.info(
        f"   Val Loss:    {summary['val_loss_mean']:.6f} ± {summary['val_loss_std']:.6f}"
    )

    # Per-target summary
    for i in range(out_size):
        r2_target = [r.get(f"r2_target_{i}", np.nan) for r in fold_results]
        mae_target = [r.get(f"mae_target_{i}", np.nan) for r in fold_results]
        logger.info(
            f"   Target {i}: R²={np.mean(r2_target):.4f}±{np.std(r2_target):.4f}, "
            f"MAE={np.mean(mae_target):.4f}±{np.std(mae_target):.4f}"
        )

    # Save summary
    with open(os.path.join(output_dir, "cv_summary.json"), "w") as f:
        summary_save = summary.copy()
        for r in summary_save["fold_results"]:
            r["history"] = None  # Too large
        json.dump(summary_save, f, indent=2)

    # Save detailed results as CSV
    results_df = pd.DataFrame(
        [{k: v for k, v in r.items() if k != "history"} for r in fold_results]
    )
    results_df.to_csv(os.path.join(output_dir, "cv_results.csv"), index=False)

    logger.info(f"\n✅ Results saved to: {output_dir}")

    return summary
