"""Trainer — fit, evaluate, predict, save, load in one line each."""

from __future__ import annotations

from pathlib import Path

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset, random_split
from tqdm import tqdm


def _best_device() -> torch.device:
    if torch.cuda.is_available():
        return torch.device("cuda")
    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


class Trainer:
    """Train any ``nn.Module`` with minimal boilerplate.

    Examples::

        trainer = Trainer(model)
        trainer.fit(X_train, y_train, epochs=10)
        metrics = trainer.evaluate(X_test, y_test)
        trainer.save("model.pt")

        # Dynamic LR + layer growth
        trainer = Trainer(model, scheduler="plateau")
        trainer.fit(X, y, epochs=50, grow_epochs=10)
    """

    def __init__(
        self,
        model: nn.Module,
        *,
        lr: float = 1e-3,
        optimizer: str = "adam",
        loss: nn.Module | None = None,
        device: torch.device | str | None = None,
        scheduler: str | None = None,
        name: str | None = None,
    ):
        """
        Args:
            model: Any ``nn.Module`` (use ``Model.*`` factories or bring your own).
            lr: Learning rate.
            optimizer: One of "adam", "sgd", "adamw".
            loss: Custom loss function (auto-detected if None).
            device: Force a device ("cpu", "cuda", "mps") or auto-detect.
            scheduler: LR scheduler — "plateau", "cosine", "warmup_cosine", "step", or None.
            name: Friendly name for this training run.
        """
        self.device = torch.device(device) if device else _best_device()
        self.model = model.to(self.device)
        self.name = name or getattr(model, "name", "neurologic")
        self._task = getattr(model, "_task", "classification")
        self._lr = lr
        self._optimizer_name = optimizer
        self._scheduler_type = scheduler

        optim_cls = {"adam": torch.optim.Adam, "sgd": torch.optim.SGD, "adamw": torch.optim.AdamW}
        self._optim_cls = optim_cls[optimizer]
        self.optimizer = self._optim_cls(self.model.parameters(), lr=lr)
        self._scheduler: torch.optim.lr_scheduler.LRScheduler | None = None

        if loss is not None:
            self.loss_fn = loss
        elif self._task == "regression":
            self.loss_fn = nn.MSELoss()
        else:
            self.loss_fn = nn.CrossEntropyLoss()

    @property
    def lr(self) -> float:
        """Current learning rate."""
        return self.optimizer.param_groups[0]["lr"]

    def summary(self) -> str:
        """Print a readable summary of the model."""
        total = sum(p.numel() for p in self.model.parameters())
        trainable = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
        lines = [
            f"  Name:       {self.name}",
            f"  Model:      {self.model.__class__.__name__}",
            f"  Task:       {self._task}",
            f"  Device:     {self.device}",
            f"  Parameters: {trainable:,} trainable / {total:,} total",
            f"  Optimizer:  {self.optimizer.__class__.__name__} (lr={self._lr})",
            f"  Loss:       {self.loss_fn.__class__.__name__}",
            f"  Scheduler:  {self._scheduler_type or 'none'}",
        ]
        text = "\n".join(lines)
        print(text)
        return text

    # ------------------------------------------------------------------
    # Training
    # ------------------------------------------------------------------

    def fit(
        self,
        X: torch.Tensor,
        y: torch.Tensor,
        *,
        epochs: int = 10,
        batch_size: int = 32,
        validation_split: float = 0.1,
        early_stopping: int = 0,
        grow_epochs: int = 0,
        grow_size: int | None = None,
        grow_max: int = 5,
    ) -> dict[str, list[float]]:
        """Train the model.

        Args:
            X: Input tensor.
            y: Target tensor.
            epochs: Number of training epochs.
            batch_size: Mini-batch size.
            validation_split: Fraction held out for validation (0 to disable).
            early_stopping: Stop if val_loss doesn't improve for N epochs (0 = off).
            grow_epochs: Add a layer every N epochs if loss plateaus (0 = off).
                         Only works with ``Model.classifier`` models.
            grow_size: Width of grown layers (None = match previous layer).
            grow_max: Maximum number of layers to grow.

        Returns:
            Dict with ``"train_loss"``, ``"val_loss"``, and ``"lr"`` lists per epoch.
        """
        X = self._to_tensor(X)
        y = self._to_tensor(y)

        dataset = TensorDataset(X, y)

        if validation_split > 0:
            val_size = max(1, int(len(dataset) * validation_split))
            train_size = len(dataset) - val_size
            train_ds, val_ds = random_split(dataset, [train_size, val_size])
        else:
            train_ds = dataset
            val_ds = None

        train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True)
        val_loader = DataLoader(val_ds, batch_size=batch_size) if val_ds else None

        # Set up LR scheduler
        self._scheduler = self._make_scheduler(epochs)

        history: dict[str, list[float]] = {"train_loss": [], "val_loss": [], "lr": []}
        best_val_loss = float("inf")
        patience_counter = 0
        grow_count = 0
        last_grow_loss = float("inf")

        print(f"Training '{self.name}' on {self.device} — {len(train_ds)} train samples")

        for epoch in range(1, epochs + 1):
            # Train
            self.model.train()
            running_loss = 0.0
            pbar = tqdm(train_loader, desc=f"Epoch {epoch}/{epochs}", leave=False)
            for xb, yb in pbar:
                xb, yb = xb.to(self.device), yb.to(self.device)
                self.optimizer.zero_grad()
                out = self.model(xb)
                loss = self.loss_fn(out, yb)
                loss.backward()
                self.optimizer.step()
                running_loss += loss.item() * xb.size(0)
                pbar.set_postfix(loss=f"{loss.item():.4f}")

            train_loss = running_loss / len(train_ds)
            history["train_loss"].append(train_loss)
            history["lr"].append(self.lr)

            # Validate
            if val_loader:
                val_loss = self._compute_loss(val_loader)
                history["val_loss"].append(val_loss)

                lr_str = f"  lr={self.lr:.6f}" if self._scheduler_type else ""
                print(f"  Epoch {epoch}/{epochs}  train_loss={train_loss:.4f}  val_loss={val_loss:.4f}{lr_str}")

                # Step scheduler
                if isinstance(self._scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
                    self._scheduler.step(val_loss)
                elif self._scheduler:
                    self._scheduler.step()

                # Early stopping
                if early_stopping > 0:
                    if val_loss < best_val_loss:
                        best_val_loss = val_loss
                        patience_counter = 0
                    else:
                        patience_counter += 1
                        if patience_counter >= early_stopping:
                            print(f"  Early stopping at epoch {epoch} (no improvement for {early_stopping} epochs)")
                            break

                # Layer growth
                metric_for_grow = val_loss
            else:
                lr_str = f"  lr={self.lr:.6f}" if self._scheduler_type else ""
                print(f"  Epoch {epoch}/{epochs}  train_loss={train_loss:.4f}{lr_str}")

                if self._scheduler and not isinstance(self._scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
                    self._scheduler.step()

                metric_for_grow = train_loss

            # Layer growth check
            if grow_epochs > 0 and epoch % grow_epochs == 0 and grow_count < grow_max:
                if hasattr(self.model, "grow") and metric_for_grow >= last_grow_loss * 0.95:
                    size = self.model.grow(grow_size)
                    self.model = self.model.to(self.device)
                    # Rebuild optimizer to pick up new parameters
                    self.optimizer = self._optim_cls(self.model.parameters(), lr=self.lr)
                    self._scheduler = self._make_scheduler(epochs - epoch)
                    grow_count += 1
                    print(f"  >> Grew layer #{grow_count} (size={size}), model now: {self.model}")
                last_grow_loss = metric_for_grow

        print(f"Done. Final train_loss={history['train_loss'][-1]:.4f}")
        return history

    # ------------------------------------------------------------------
    # Evaluation
    # ------------------------------------------------------------------

    def evaluate(self, X: torch.Tensor, y: torch.Tensor, *, batch_size: int = 32) -> dict[str, float]:
        """Evaluate the model and return metrics.

        Returns dict with ``"loss"`` and, for classification, ``"accuracy"``.
        """
        X = self._to_tensor(X)
        y = self._to_tensor(y)
        loader = DataLoader(TensorDataset(X, y), batch_size=batch_size)

        loss = self._compute_loss(loader)
        metrics: dict[str, float] = {"loss": loss}

        if self._task == "classification":
            metrics["accuracy"] = self._compute_accuracy(loader)

        return metrics

    # ------------------------------------------------------------------
    # Prediction
    # ------------------------------------------------------------------

    def predict(self, X: torch.Tensor, *, batch_size: int = 32) -> torch.Tensor:
        """Run inference and return predictions."""
        X = self._to_tensor(X)
        self.model.eval()
        preds = []
        loader = DataLoader(TensorDataset(X), batch_size=batch_size)
        with torch.no_grad():
            for (xb,) in loader:
                out = self.model(xb.to(self.device))
                if self._task == "classification":
                    out = out.argmax(dim=1)
                preds.append(out.cpu())
        return torch.cat(preds)

    # ------------------------------------------------------------------
    # Save / Load
    # ------------------------------------------------------------------

    def save(self, path: str | Path) -> None:
        """Save model + optimizer state."""
        torch.save(
            {
                "model_state": self.model.state_dict(),
                "optimizer_state": self.optimizer.state_dict(),
                "model": self.model,
                "name": self.name,
            },
            path,
        )
        print(f"Saved '{self.name}' to {path}")

    @classmethod
    def load(cls, path: str | Path, *, device: str | None = None) -> "Trainer":
        """Load a saved trainer."""
        checkpoint = torch.load(path, map_location=device or _best_device(), weights_only=False)
        model = checkpoint["model"]
        trainer = cls(model, device=device, name=checkpoint.get("name"))
        trainer.model.load_state_dict(checkpoint["model_state"])
        trainer.optimizer.load_state_dict(checkpoint["optimizer_state"])
        return trainer

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _to_tensor(self, data: torch.Tensor) -> torch.Tensor:
        if not isinstance(data, torch.Tensor):
            return torch.tensor(data, dtype=torch.float32)
        return data

    def _make_scheduler(self, epochs: int) -> torch.optim.lr_scheduler.LRScheduler | None:
        if self._scheduler_type == "plateau":
            return torch.optim.lr_scheduler.ReduceLROnPlateau(
                self.optimizer, mode="min", factor=0.5, patience=3,
            )
        if self._scheduler_type == "step":
            return torch.optim.lr_scheduler.StepLR(self.optimizer, step_size=max(1, epochs // 3), gamma=0.1)
        if self._scheduler_type == "cosine":
            return torch.optim.lr_scheduler.CosineAnnealingLR(self.optimizer, T_max=max(1, epochs))
        if self._scheduler_type == "warmup_cosine":
            warmup_epochs = max(1, epochs // 10)

            def lr_lambda(epoch: int) -> float:
                if epoch < warmup_epochs:
                    return (epoch + 1) / warmup_epochs
                progress = (epoch - warmup_epochs) / max(1, epochs - warmup_epochs)
                return 0.5 * (1.0 + __import__("math").cos(__import__("math").pi * progress))

            return torch.optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda)
        return None

    def _compute_loss(self, loader: DataLoader) -> float:
        self.model.eval()
        total_loss = 0.0
        total = 0
        with torch.no_grad():
            for xb, yb in loader:
                xb, yb = xb.to(self.device), yb.to(self.device)
                out = self.model(xb)
                total_loss += self.loss_fn(out, yb).item() * xb.size(0)
                total += xb.size(0)
        return total_loss / total

    def _compute_accuracy(self, loader: DataLoader) -> float:
        self.model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for xb, yb in loader:
                xb, yb = xb.to(self.device), yb.to(self.device)
                preds = self.model(xb).argmax(dim=1)
                correct += (preds == yb).sum().item()
                total += yb.size(0)
        return correct / total
