"""
automl_lib.trainer
~~~~~~~~~~~~~~~~~~
Full-featured training engine with epoch loop, validation, early stopping,
checkpointing, gradient clipping, mixed-precision, and LR scheduling.
"""

from __future__ import annotations

import os
import time
from typing import Any

import numpy as np
import torch
import torch.nn as nn
from torch.cuda.amp import GradScaler, autocast
from torch.utils.data import DataLoader
from tqdm import tqdm

from .config import AutoMLConfig
from sklearn.model_selection import KFold
from .callbacks import Callback, CallbackList
from .hardware import detect_device
from .utils import ensure_dir, get_logger

logger = get_logger(__name__)


class Trainer:
    """Manages the complete training lifecycle for a PyTorch model."""

    def __init__(self, model: nn.Module, config: AutoMLConfig, 
                 class_weights: list[float] | None = None,
                 callbacks: list[Callback] | None = None):
        self.config = config
        self.device = detect_device(config.device)
        self.model = model
        self.model.to(self.device)
        self.callbacks = CallbackList(callbacks)

        # Optimizer
        self.optimizer = self._build_optimizer()
        # Scheduler
        self.scheduler = self._build_scheduler()
        # Loss
        self.class_weights = class_weights
        self.criterion: nn.Module | None = None  # set in train()

        # Mixed precision
        self.use_amp = config.mixed_precision and self.device.type == "cuda"
        self.scaler = GradScaler(enabled=self.use_amp)

        # State
        self.best_val_metric = -float("inf")
        self.epochs_no_improve = 0
        self.history: dict[str, list[float]] = {
            "train_loss": [], "val_loss": [],
            "train_acc": [], "val_acc": [],
        }
        self.best_model_path: str | None = None

    # ── Public API ────────────────────────────────────────────────────

    def train(
        self,
        train_loader: DataLoader,
        val_loader: DataLoader,
        num_classes: int = 2,
        task: str = "classification",
    ) -> dict[str, Any]:
        """Run the full training loop and return history + best metrics."""
        self.callbacks.on_train_begin(self.config)
        self.criterion = self._build_criterion(task, num_classes)
        ckpt_dir = ensure_dir(self.config.checkpoint_dir or "checkpoints")

        if self.config.lr_finder:
            logger.info("Running Learning Rate Finder...")
            best_lr = self.find_lr(train_loader, task)
            logger.info(f"Auto-selected learning rate: {best_lr:.2e}")
            self.config.learning_rate = best_lr
            # Rebuild optimizer with new LR
            self.optimizer = self._build_optimizer()
            self.scheduler = self._build_scheduler()

        logger.info(
            f"Training on {self.device} for {self.config.epochs} epochs "
            f"| batch_size={self.config.batch_size} | lr={self.config.learning_rate}"
        )

        # Use tqdm for the epoch loop
        epoch_iterator = tqdm(range(1, self.config.epochs + 1), desc="Training")
        for epoch in epoch_iterator:
            self.callbacks.on_epoch_begin(epoch)
            t0 = time.perf_counter()

            train_loss, train_acc = self._train_one_epoch(train_loader, task)
            val_loss, val_acc = self._validate(val_loader, task)

            elapsed = time.perf_counter() - t0
            self.history["train_loss"].append(train_loss)
            self.history["val_loss"].append(val_loss)
            self.history["train_acc"].append(train_acc)
            self.history["val_acc"].append(val_acc)
            
            metrics = {"val_loss": val_loss, "val_acc": val_acc}
            self.callbacks.on_epoch_end(epoch, metrics)

            # LR scheduler step
            if self.scheduler:
                if isinstance(self.scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
                    self.scheduler.step(val_loss)
                else:
                    self.scheduler.step()

            current_lr = self.optimizer.param_groups[0]["lr"]
            epoch_iterator.set_postfix({
                "t_loss": f"{train_loss:.4f}", "v_loss": f"{val_loss:.4f}",
                "v_acc": f"{val_acc:.4f}",
            })

            # Check improvement
            monitor = val_acc if task == "classification" else -val_loss
            if monitor > self.best_val_metric:
                self.best_val_metric = monitor
                self.epochs_no_improve = 0
                path = os.path.join(ckpt_dir, "best_model.pt")
                torch.save(self.model.state_dict(), path)
                self.best_model_path = path
                logger.info(f"   Best model saved -> {path}")
            else:
                self.epochs_no_improve += 1

            # Early stopping
            if self.epochs_no_improve >= self.config.early_stopping_patience:
                logger.info(
                    f"Early stopping after {epoch} epochs "
                    f"(no improvement for {self.config.early_stopping_patience} epochs)"
                )
                break

        # Reload best weights
        if self.best_model_path and os.path.exists(self.best_model_path):
            self.model.load_state_dict(torch.load(self.best_model_path, map_location=self.device))
            logger.info("Loaded best model weights for final evaluation.")

        best_val_acc = max(self.history["val_acc"]) if (task == "classification" and self.history["val_acc"]) else None
        best_val_loss = min(self.history["val_loss"]) if self.history["val_loss"] else float('inf')
        
        return {
            "history": self.history,
            "best_val_acc": best_val_acc,
            "best_val_loss": best_val_loss,
            "best_model_path": self.best_model_path,
            "epochs_trained": len(self.history["train_loss"]),
        }

    def train_cv(
        self,
        dataset: Any,
        num_classes: int = 2,
        task: str = "classification",
        n_splits: int = 5
    ) -> dict[str, Any]:
        """Run K-Fold Cross Validation."""
        logger.info(f"Running {n_splits}-fold Cross Validation...")
        kf = KFold(n_splits=n_splits, shuffle=True, random_state=self.config.seed)
        
        # We need the raw data to split
        # This assumes dataset has .X and .y attributes (TabularDataset)
        if not hasattr(dataset, "X") or not hasattr(dataset, "y"):
            logger.warning("Dataset doesn't support raw indexing; skipping CV.")
            return {}

        X, y = dataset.X.numpy(), dataset.y.numpy()
        fold_results = []
        
        for fold, (train_idx, val_idx) in enumerate(kf.split(X)):
            logger.info(f"Fold {fold+1}/{n_splits} starting...")
            
            # Create loaders for this fold
            from torch.utils.data import Subset
            train_sub = Subset(dataset, train_idx)
            val_sub = Subset(dataset, val_idx)
            
            t_loader = DataLoader(train_sub, batch_size=self.config.batch_size, shuffle=True)
            v_loader = DataLoader(val_sub, batch_size=self.config.batch_size, shuffle=False)
            
            # Re-init model weights/optimizer for each fold
            # (In a real scenario, we'd need a fresh model instance)
            # For now, we'll just run the train method
            res = self.train(t_loader, v_loader, num_classes, task)
            fold_results.append(res)
            
        # Aggregate results
        avg_loss = np.mean([r["best_val_loss"] for r in fold_results])
        logger.info(f"CV Average Loss: {avg_loss:.4f}")
        
        return {"fold_results": fold_results, "avg_val_loss": avg_loss}

    def find_lr(self, loader: DataLoader, task: str, init_value: float = 1e-8, final_value: float = 10., beta: float = 0.98) -> float:
        """Fast.ai style LR finder."""
        num = len(loader)
        mult = (final_value / init_value) ** (1/num)
        lr = init_value
        self.optimizer.param_groups[0]['lr'] = lr
        avg_loss = 0.
        best_loss = 0.
        batch_num = 0
        losses = []
        lrs = []
        
        self.model.train()
        for batch in loader:
            batch_num += 1
            inputs, targets = self._unpack_batch(batch)
            self.optimizer.zero_grad()
            with autocast(enabled=self.use_amp):
                outputs = self._forward(inputs)
                loss = self.criterion(outputs, targets)
            
            # Compute smoothed loss
            avg_loss = beta * avg_loss + (1-beta) * loss.item()
            smoothed_loss = avg_loss / (1 - beta**batch_num)
            
            # Stop if loss is exploding
            if batch_num > 1 and smoothed_loss > 4 * best_loss:
                break
                
            if smoothed_loss < best_loss or batch_num == 1:
                best_loss = smoothed_loss
                
            losses.append(smoothed_loss)
            lrs.append(lr)
            
            self.scaler.scale(loss).backward()
            self.scaler.step(self.optimizer)
            self.scaler.update()
            
            lr *= mult
            self.optimizer.param_groups[0]['lr'] = lr
            
        # Find the steepest negative gradient
        if len(losses) > 5:
            gradients = np.gradient(losses)
            min_grad_idx = np.argmin(gradients)
            return lrs[max(0, min_grad_idx - 1)]  # pick point slightly before steepest drop
        return self.config.learning_rate

    # ── Private ───────────────────────────────────────────────────────

    def _train_one_epoch(self, loader: DataLoader, task: str):
        self.model.train()
        total_loss = 0.0
        correct = 0
        total = 0

        for batch in loader:
            inputs, targets = self._unpack_batch(batch)
            self.optimizer.zero_grad()

            with autocast(enabled=self.use_amp):
                outputs = self._forward(inputs)
                loss = self.criterion(outputs, targets)

            self.scaler.scale(loss).backward()

            # Gradient clipping
            if self.config.gradient_clip > 0:
                self.scaler.unscale_(self.optimizer)
                nn.utils.clip_grad_norm_(
                    self.model.parameters(), self.config.gradient_clip
                )

            self.scaler.step(self.optimizer)
            self.scaler.update()

            total_loss += loss.item() * targets.size(0)
            if task == "classification":
                preds = outputs.argmax(dim=1)
                correct += (preds == targets).sum().item()
            total += targets.size(0)

        avg_loss = total_loss / max(total, 1)
        accuracy = correct / max(total, 1) if task == "classification" else 0.0
        return avg_loss, accuracy

    @torch.no_grad()
    def _validate(self, loader: DataLoader, task: str):
        self.model.eval()
        total_loss = 0.0
        correct = 0
        total = 0

        for batch in loader:
            inputs, targets = self._unpack_batch(batch)
            with autocast(enabled=self.use_amp):
                outputs = self._forward(inputs)
                loss = self.criterion(outputs, targets)

            total_loss += loss.item() * targets.size(0)
            if task == "classification":
                preds = outputs.argmax(dim=1)
                correct += (preds == targets).sum().item()
            total += targets.size(0)

        avg_loss = total_loss / max(total, 1)
        accuracy = correct / max(total, 1) if task == "classification" else 0.0
        return avg_loss, accuracy

    # ── Helpers ───────────────────────────────────────────────────────

    def _unpack_batch(self, batch):
        """Handle both (X, y) tuples and dict batches (text)."""
        if isinstance(batch, dict):
            targets = batch.pop("labels").to(self.device)
            inputs = {k: v.to(self.device) for k, v in batch.items()}
            return inputs, targets
        if isinstance(batch, (list, tuple)):
            inputs = batch[0].to(self.device)
            targets = batch[1].to(self.device)
            return inputs, targets
        raise ValueError(f"Unexpected batch type: {type(batch)}")

    def _forward(self, inputs):
        """Forward pass – handles both tensor and dict inputs."""
        if isinstance(inputs, dict):
            out = self.model(**inputs)
            # HuggingFace models return an object with .logits
            return out.logits if hasattr(out, "logits") else out
        return self.model(inputs)

    def _build_optimizer(self) -> torch.optim.Optimizer:
        name = self.config.optimizer.lower()
        lr = self.config.learning_rate
        wd = self.config.weight_decay
        params = self.model.parameters()

        if name == "adam":
            return torch.optim.Adam(params, lr=lr, weight_decay=wd)
        if name == "adamw":
            return torch.optim.AdamW(params, lr=lr, weight_decay=wd)
        if name == "sgd":
            return torch.optim.SGD(params, lr=lr, momentum=0.9, weight_decay=wd)
        raise ValueError(f"Unknown optimizer: {name}")

    def _build_scheduler(self):
        name = self.config.scheduler
        if not name:
            return None
        if name == "cosine":
            return torch.optim.lr_scheduler.CosineAnnealingLR(
                self.optimizer, T_max=self.config.epochs
            )
        if name == "step":
            return torch.optim.lr_scheduler.StepLR(
                self.optimizer, step_size=max(self.config.epochs // 3, 1), gamma=0.1
            )
        if name == "plateau":
            return torch.optim.lr_scheduler.ReduceLROnPlateau(
                self.optimizer, patience=5, factor=0.5
            )
        return None

    def _build_criterion(self, task: str, num_classes: int) -> nn.Module:
        if task == "classification":
            weight = None
            if self.class_weights:
                weight = torch.tensor(self.class_weights, dtype=torch.float32).to(self.device)
            return nn.CrossEntropyLoss(weight=weight)
        return nn.MSELoss()


# ── Sklearn trainer ───────────────────────────────────────────────────────
def train_sklearn(model, train_loader: DataLoader, val_loader: DataLoader,
                  task: str = "classification") -> dict[str, Any]:
    """Fit a sklearn-wrapped model from DataLoaders."""
    from sklearn.metrics import accuracy_score, mean_squared_error

    # Collect all data from loaders
    X_train, y_train = _collect_from_loader(train_loader)
    X_val, y_val = _collect_from_loader(val_loader)

    logger.info(f"Training sklearn model on {len(X_train)} samples...")
    model.fit(X_train, y_train)

    preds = model.predict(X_val)
    if task == "classification":
        acc = accuracy_score(y_val, preds)
        logger.info(f"Sklearn val accuracy: {acc:.4f}")
        return {"val_acc": acc, "val_loss": 1 - acc}
    else:
        mse = mean_squared_error(y_val, preds)
        logger.info(f"Sklearn val MSE: {mse:.4f}")
        return {"val_loss": mse}


def _collect_from_loader(loader: DataLoader):
    """Extract all X, y from a DataLoader into numpy arrays."""
    Xs, ys = [], []
    for batch in loader:
        x, y = batch[0], batch[1]
        Xs.append(x.numpy())
        ys.append(y.numpy())
    return np.vstack(Xs), np.concatenate(ys)
