"""

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 trainiqConfig

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: trainiqConfig, 

                 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)

