"""Standard PyTorch training-loop trainer.

``LoopTrainer`` extends :class:`~matrice_models.training.trainer.Trainer`
with a built-in epoch loop — the pattern used by classification
(``ml_pytorch_vision_classification``) and video-action
(``r2plus1d``) models.  Subclasses only implement the per-epoch work;
all boilerplate (optimizer / scheduler setup, metric formatting,
checkpointing, early-stop checks, ActionTracker logging) lives here.

Example::

    class ResNetTrainer(LoopTrainer):
        def setup_data(self) -> None: ...
        def setup_model(self) -> None: ...
        def get_label_map(self) -> dict[str, str]: ...
        def configure_criterion(self) -> nn.Module:
            return nn.CrossEntropyLoss()

        def configure_optimizer(self) -> optim.Optimizer:
            return optim.AdamW(self.model.parameters(), lr=self.config.learning_rate)

        def train_one_epoch(self, epoch: int) -> dict[str, float]: ...
        def validate_one_epoch(self, epoch: int) -> dict[str, float]: ...
        def run_evaluation(self) -> list[dict[str, Any]]: ...


    trainer = ResNetTrainer(action_id, model_info, config)
    result = trainer.fit()
"""

from __future__ import annotations

import logging
import math
from abc import abstractmethod
from typing import TYPE_CHECKING, Any

from matrice_models.training.trainer import Trainer, TrainResult


if TYPE_CHECKING:
    from pathlib import Path

    import torch.nn as nn
    import torch.optim as optim
    from torch.optim.lr_scheduler import LRScheduler

    from matrice_models.config.base import ModelInfo, TrainConfig
    from matrice_models.training.callbacks import TrainerCallback


logger = logging.getLogger(__name__)


class LoopTrainer(Trainer):
    """Trainer with a standard PyTorch epoch loop.

    Subclasses implement four extra hooks beyond the ones inherited from
    :class:`Trainer`:

    ==================  ======================================================
    Hook                Responsibility
    ==================  ======================================================
    configure_criterion Return the loss function (e.g. ``nn.CrossEntropyLoss``)
    configure_optimizer Return the optimiser for ``self.model``
    train_one_epoch     Run one training epoch, return a flat metrics dict
    validate_one_epoch  Run one validation epoch, return a flat metrics dict
    ==================  ======================================================

    Optionally override :meth:`configure_scheduler` to add a learning-rate
    schedule.

    The built-in :meth:`train_loop` then orchestrates:

    1. Criterion / optimiser / scheduler creation.
    2. ``on_train_begin`` callback.
    3. Per-epoch: ``on_epoch_begin`` → train → validate → scheduler step
       → format metrics → ``on_epoch_end`` → checkpoint → early-stop check.
    4. ``on_train_end`` callback and :class:`TrainResult` population.
    """

    # ------------------------------------------------------------------ #
    #  Construction                                                        #
    # ------------------------------------------------------------------ #

    def __init__(
        self,
        action_id: str | None,
        model_info: ModelInfo,
        config: TrainConfig,
        *,
        checkpoint_dir: str | Path = "checkpoints",
        callbacks: list[TrainerCallback] | None = None,
        monitor: str = "acc@1",
        monitor_mode: str = "max",
    ) -> None:
        """Initialise the loop trainer.

        Args:
            action_id: Platform action ID (``None`` for local runs).
            model_info: Model identity and capabilities.
            config: Training hyperparameters.
            checkpoint_dir: Directory for checkpoint files.
            callbacks: Optional list of callbacks.  Defaults to
                ``[ActionTrackerCallback()]``.
            monitor: Validation metric to track for *best model*
                selection.  Must match a key returned by
                :meth:`validate_one_epoch`.
            monitor_mode: ``"max"`` when higher is better (accuracy)
                or ``"min"`` when lower is better (loss).
        """
        if monitor_mode not in ("min", "max"):
            msg = f"monitor_mode must be 'min' or 'max', got '{monitor_mode}'"
            raise ValueError(msg)

        super().__init__(
            action_id,
            model_info,
            config,
            checkpoint_dir=checkpoint_dir,
            callbacks=callbacks,
        )

        self.monitor = monitor
        self.monitor_mode = monitor_mode

        self.criterion: nn.Module | None = None
        self.optimizer: optim.Optimizer | None = None
        self.scheduler: LRScheduler | None = None

    # ------------------------------------------------------------------ #
    #  Abstract hooks — subclasses MUST implement                          #
    # ------------------------------------------------------------------ #

    @abstractmethod
    def configure_criterion(self) -> nn.Module:
        """Return the loss function.

        Called once before the epoch loop begins.  The returned module
        is stored as ``self.criterion`` so :meth:`train_one_epoch` and
        :meth:`validate_one_epoch` can access it.

        Example::

            def configure_criterion(self) -> nn.Module:
                return nn.CrossEntropyLoss().to(self.device)
        """

    @abstractmethod
    def configure_optimizer(self) -> optim.Optimizer:
        """Return the optimiser for ``self.model``.

        Called once after ``setup_model`` has populated ``self.model``
        and the device has been resolved.  The returned optimiser is
        stored as ``self.optimizer``.

        Example::

            def configure_optimizer(self) -> optim.Optimizer:
                return torch.optim.AdamW(
                    self.model.parameters(),
                    lr=self.config.learning_rate,
                )
        """

    @abstractmethod
    def train_one_epoch(self, epoch: int) -> dict[str, float]:
        """Run one training epoch and return metrics.

        Typical implementations iterate ``self.train_loader``, compute
        the forward/backward pass, and accumulate metrics.

        The trainer guarantees that ``self.model``, ``self.criterion``,
        ``self.optimizer``, and ``self.device`` are already set.

        Args:
            epoch: Zero-based epoch index.

        Returns:
            Flat dict of metric values, e.g.
            ``{"loss": 0.45, "acc@1": 0.82, "acc@5": 0.96}``.
        """

    @abstractmethod
    def validate_one_epoch(self, epoch: int) -> dict[str, float]:
        """Run one validation epoch and return metrics.

        Similar to :meth:`train_one_epoch` but in ``torch.no_grad()``
        mode and without gradient updates.

        Args:
            epoch: Zero-based epoch index.

        Returns:
            Flat dict of metric values, e.g.
            ``{"loss": 0.52, "acc@1": 0.78, "acc@5": 0.94}``.
            Return an **empty dict** if no validation set exists.
        """

    # ------------------------------------------------------------------ #
    #  Overridable hooks — subclasses MAY override                         #
    # ------------------------------------------------------------------ #

    def configure_scheduler(self) -> LRScheduler | None:
        """Return an optional learning-rate scheduler.

        Called once after :meth:`configure_optimizer`.  The default
        returns ``None`` (constant LR).  Override to add e.g.
        ``StepLR``, ``CosineAnnealingLR``, etc.

        Returns:
            A scheduler instance, or ``None`` for constant LR.

        Example::

            def configure_scheduler(self) -> LRScheduler:
                return torch.optim.lr_scheduler.StepLR(
                    self.optimizer,
                    step_size=30,
                    gamma=0.1,
                )
        """
        return None

    # ------------------------------------------------------------------ #
    #  train_loop implementation                                           #
    # ------------------------------------------------------------------ #

    def train_loop(self) -> None:
        """Execute the standard PyTorch training loop.

        Lifecycle::

            configure_criterion → configure_optimizer → configure_scheduler
            on_train_begin
            for epoch in range(epochs):
                on_epoch_begin
                train_one_epoch  →  validate_one_epoch  →  scheduler.step()
                format_metrics → on_epoch_end → log_epoch → checkpoint → early_stop?
            on_train_end(result)

        Populates ``self._result`` before returning.
        """
        self.criterion = self.configure_criterion()
        self.optimizer = self.configure_optimizer()
        self.scheduler = self.configure_scheduler()

        best_value = -math.inf if self.monitor_mode == "max" else math.inf
        best_epoch = 0
        epochs_completed = 0

        self.notify("on_train_begin", self)

        for epoch in range(self.config.epochs):
            self.notify("on_epoch_begin", self, epoch)

            train_metrics = self.train_one_epoch(epoch)
            val_metrics = self.validate_one_epoch(epoch)

            if self.scheduler is not None:
                self.scheduler.step()

            epoch_metrics = _format_epoch_metrics(train_metrics, val_metrics)
            self.notify("on_epoch_end", self, epoch, epoch_metrics)
            self.log_epoch(epoch, epoch_metrics)

            current = val_metrics.get(self.monitor)
            is_best = current is not None and _is_improvement(current, best_value, self.monitor_mode)
            if is_best:
                best_value = current  # type: ignore[assignment]
                best_epoch = epoch

            self._save_loop_checkpoint(epoch, is_best=is_best)

            epochs_completed = epoch + 1

            if self._should_stop():
                logger.info("Early stopping triggered at epoch %d", epoch)
                break

        self._result = TrainResult(
            best_metric_name=self.monitor,
            best_metric_value=best_value if math.isfinite(best_value) else 0.0,
            best_epoch=best_epoch,
            epochs_completed=epochs_completed,
        )

        self.notify("on_train_end", self, self._result)

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

    def _save_loop_checkpoint(self, epoch: int, *, is_best: bool) -> None:
        """Persist a checkpoint using :func:`checkpoint.save_checkpoint`."""
        from matrice_models.training.checkpoint import save_checkpoint

        state: dict[str, Any] = {
            "epoch": epoch,
            "architecture": self.model_info.architecture,
        }

        if self.model is not None and hasattr(self.model, "state_dict"):
            state["state_dict"] = self.model.state_dict()
        if self.optimizer is not None:
            state["optimizer"] = self.optimizer.state_dict()
        if self.scheduler is not None:
            state["scheduler"] = self.scheduler.state_dict()

        written = save_checkpoint(
            state,
            self.model,
            is_best=is_best,
            checkpoint_dir=self.checkpoint_dir,
        )

        if is_best and self._result is not None:
            self._result.checkpoint_paths.extend(written)

        if is_best:
            best_pth = self.checkpoint_dir / "model_best.pth.tar"
            if best_pth.exists():
                self.action_tracker.upload_checkpoint(str(best_pth))

    def _should_stop(self) -> bool:
        """Check whether any callback has requested early stopping."""
        return any(getattr(cb, "should_stop", False) for cb in self.callbacks)


# ------------------------------------------------------------------ #
#  Module-level helpers                                                #
# ------------------------------------------------------------------ #


def _format_epoch_metrics(
    train_metrics: dict[str, float],
    val_metrics: dict[str, float],
) -> list[dict[str, Any]]:
    """Convert flat metric dicts into ActionTracker epoch-metric format.

    Args:
        train_metrics: Metrics from :meth:`train_one_epoch`.
        val_metrics: Metrics from :meth:`validate_one_epoch`.

    Returns:
        List of ``{"splitType", "metricName", "metricValue"}`` dicts.
    """
    formatted: list[dict[str, Any]] = []

    for name, value in train_metrics.items():
        formatted.append({"splitType": "train", "metricName": name, "metricValue": value})
    for name, value in val_metrics.items():
        formatted.append({"splitType": "val", "metricName": name, "metricValue": value})

    return formatted


def _is_improvement(
    current: float,
    best: float,
    mode: str,
) -> bool:
    """Return ``True`` when *current* is strictly better than *best*."""
    if mode == "max":
        return current > best
    return current < best
