# ------------------------------------------------------------------------
# RF-DETR
# Copyright (c) 2025 Roboflow. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
# ------------------------------------------------------------------------

"""Exponential Moving Average callback compatible with ``ModelEma``."""

from __future__ import annotations

import math
import warnings
from copy import deepcopy
from typing import Any, Optional

import torch
from pytorch_lightning import Callback, LightningModule, Trainer
from torch.optim.swa_utils import AveragedModel


class RFDETREMACallback(Callback):
    """Exponential Moving Average with optional tau-based warm-up.

    Drop-in replacement for ``rfdetr.util.utils.ModelEma`` implemented as a
    plain Lightning callback around :class:`torch.optim.swa_utils.AveragedModel`.
    The ``_avg_fn`` reproduces the exact same formula as ``ModelEma``
    (1-indexed ``updates`` counter, optional ``tau`` warm-up).

    Args:
        decay: Base EMA decay factor. Corresponds to ``TrainConfig.ema_decay``.
        tau: Warm-up time constant (in optimizer steps). When > 0 the
            effective decay ramps from 0 towards *decay* following
            ``decay * (1 - exp(-updates / tau))``. Corresponds to
            ``TrainConfig.ema_tau``.
        use_buffers: Whether buffers are averaged in addition to parameters.
        update_interval_steps: Update EMA every N optimizer steps.
    """

    def __init__(
        self,
        decay: float = 0.993,
        tau: int = 100,
        use_buffers: bool = True,
        update_interval_steps: int = 1,
    ) -> None:
        super().__init__()
        self._decay = decay
        self._tau = tau
        self._use_buffers = use_buffers
        self._update_interval_steps = max(1, int(update_interval_steps))

        self._average_model: Optional[AveragedModel] = None
        self._latest_update_step = 0
        self._latest_update_epoch = -1
        self._swapped_state_dict: Optional[dict[str, torch.Tensor]] = None
        self._pending_average_state_dict: Optional[dict[str, Any]] = None

    def _avg_fn(
        self,
        averaged_param: torch.Tensor,
        model_param: torch.Tensor,
        num_averaged: int,
    ) -> torch.Tensor:
        """Compute the EMA update for a single parameter tensor.

        Matches the ``ModelEma`` formula where ``updates`` is 1-indexed:
        PTL's ``num_averaged`` starts at 0 (incremented *after* calling
        ``avg_fn``), so ``updates = num_averaged + 1`` reproduces the
        same sequence of effective decay values.

        Args:
            averaged_param: Current EMA parameter value.
            model_param: Corresponding live model parameter value.
            num_averaged: Number of models averaged so far (0-indexed).

        Returns:
            Updated EMA parameter tensor.
        """
        updates = num_averaged + 1  # match ModelEma 1-indexed counter
        if self._tau > 0:
            effective_decay = self._decay * (1 - math.exp(-updates / self._tau))
        else:
            effective_decay = self._decay
        return averaged_param * effective_decay + model_param * (1.0 - effective_decay)

    def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str) -> None:
        """Initialise the averaged model at fit start.

        Args:
            trainer: The Lightning Trainer instance.
            pl_module: The ``RFDETRModelModule`` being trained.
            stage: Current trainer stage (``"fit"``, ``"validate"``, ...).
        """
        if stage != "fit":
            return

        self._average_model = AveragedModel(
            model=pl_module,
            device=pl_module.device,
            use_buffers=self._use_buffers,
            avg_fn=self._avg_fn,
        )
        # The averaged model is inference-only; PTL never calls .eval() on it
        # because it is not registered as a Lightning module.  Without this,
        # dropout layers stay in training mode and produce ~random outputs.
        self._average_model.eval()

        if self._pending_average_state_dict is not None:
            self._average_model.load_state_dict(self._pending_average_state_dict)
            self._pending_average_state_dict = None
        elif hasattr(pl_module, "_pending_legacy_ema_state"):
            legacy_ema_state = getattr(pl_module, "_pending_legacy_ema_state")
            if isinstance(legacy_ema_state, dict):
                incompatible = self._average_model.module.model.load_state_dict(legacy_ema_state, strict=False)
                if incompatible.missing_keys or incompatible.unexpected_keys:
                    warnings.warn(
                        "Legacy EMA checkpoint loaded with non-exact key match; "
                        f"missing={len(incompatible.missing_keys)} "
                        f"unexpected={len(incompatible.unexpected_keys)}.",
                        UserWarning,
                        stacklevel=2,
                    )
            delattr(pl_module, "_pending_legacy_ema_state")

    def should_update(
        self,
        step_idx: Optional[int] = None,
        epoch_idx: Optional[int] = None,
    ) -> bool:
        """Return ``True`` after every optimizer step and every epoch end.

        The base ``WeightAveraging`` only updates on steps. This override
        also triggers an update at epoch boundaries, matching RF-DETR's
        existing EMA behaviour.

        Args:
            step_idx: Index of the last optimizer step, or ``None``.
            epoch_idx: Index of the last epoch, or ``None``.

        Returns:
            Whether the averaged model should be updated.
        """
        return step_idx is not None or epoch_idx is not None

    def _swap_models(self, pl_module: LightningModule) -> None:
        """Swap live model weights with averaged EMA weights."""
        if self._average_model is None:
            return
        if self._swapped_state_dict is None:
            self._swapped_state_dict = deepcopy(pl_module.state_dict())
            pl_module.load_state_dict(self._average_model.module.state_dict(), strict=True)
            return
        pl_module.load_state_dict(self._swapped_state_dict, strict=True)
        self._swapped_state_dict = None

    def on_train_batch_end(
        self,
        trainer: Trainer,
        pl_module: LightningModule,
        outputs: Any,
        batch: Any,
        batch_idx: int,
    ) -> None:
        """Update EMA after optimizer steps."""
        if self._average_model is None:
            return
        step_idx = trainer.global_step - 1
        if trainer.global_step <= self._latest_update_step:
            return

        self._latest_update_step = trainer.global_step
        should_update_step = trainer.global_step % self._update_interval_steps == 0
        if should_update_step and self.should_update(step_idx=step_idx):
            self._average_model.update_parameters(pl_module)

    def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
        """Optionally update EMA at epoch boundaries."""
        if self._average_model is None:
            return
        if trainer.current_epoch > self._latest_update_epoch and self.should_update(epoch_idx=trainer.current_epoch):
            self._average_model.update_parameters(pl_module)
            self._latest_update_epoch = trainer.current_epoch

    def on_test_epoch_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
        """Evaluate tests using averaged EMA weights."""
        self._swap_models(pl_module)

    def on_test_epoch_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
        """Restore live weights after test evaluation."""
        self._swap_models(pl_module)

    def on_train_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
        """Leave the module in EMA state after training finishes."""
        if self._average_model is not None:
            pl_module.load_state_dict(self._average_model.module.state_dict(), strict=True)
        self._swapped_state_dict = None

    def state_dict(self) -> dict[str, Any]:
        """Return callback state for checkpointing."""
        state = {
            "latest_update_step": self._latest_update_step,
            "latest_update_epoch": self._latest_update_epoch,
        }
        if self._average_model is not None:
            state["average_model_state_dict"] = self._average_model.state_dict()
        return state

    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
        """Restore callback state from checkpoints."""
        self._latest_update_step = state_dict.get("latest_update_step", 0)
        self._latest_update_epoch = state_dict.get("latest_update_epoch", -1)
        self._pending_average_state_dict = state_dict.get("average_model_state_dict")

    def get_ema_model_state_dict(self) -> Optional[dict[str, torch.Tensor]]:
        """Expose EMA model weights for external checkpoint callbacks."""
        if self._average_model is None or not hasattr(self._average_model.module, "model"):
            return None
        return {k: v.detach().clone() for k, v in self._average_model.module.model.state_dict().items()}
