"""Optimizer, learning-rate schedule and EMA of weights."""

from __future__ import annotations

import math
from collections.abc import Collection, Iterator
from contextlib import contextmanager

import torch
from torch import Tensor, nn
from torch.optim.lr_scheduler import LambdaLR

from nida.train.config import TrainConfig

__all__ = ["EMA", "build_optimizer", "build_scheduler", "lr_factor"]


def lr_factor(step: int, *, warmup_steps: int, total_steps: int, min_lr_ratio: float) -> float:
    """Multiplier of the peak learning rate for the optimizer step number ``step`` (0-based).

    Linear warmup ``(step + 1) / warmup_steps`` for ``step < warmup_steps`` (so the very
    first step already has a non-zero rate and step ``warmup_steps - 1`` reaches 1.0), then
    cosine decay from 1.0 at ``step == warmup_steps`` to ``min_lr_ratio`` at
    ``step == total_steps``; constant ``min_lr_ratio`` afterwards.
    """
    if step < 0:
        raise ValueError("step must be >= 0")
    if step < warmup_steps:
        return (step + 1) / warmup_steps
    decay_steps = total_steps - warmup_steps
    progress = 1.0 if decay_steps <= 0 else min(1.0, (step - warmup_steps) / decay_steps)
    cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
    return min_lr_ratio + (1.0 - min_lr_ratio) * cosine


def build_optimizer(model: nn.Module, cfg: TrainConfig) -> torch.optim.AdamW:
    """AdamW with weight decay on weight matrices (``ndim >= 2``) only.

    Biases, norm scales and embedding tables are exempt (decaying a lookup table shrinks
    the rows of rare tokens/classes towards zero between their rare updates).
    """
    embedding_ids = {
        id(module.weight) for module in model.modules() if isinstance(module, nn.Embedding)
    }
    decay: list[nn.Parameter] = []
    no_decay: list[nn.Parameter] = []
    for _, param in model.named_parameters():
        if param.requires_grad:
            exempt = param.ndim < 2 or id(param) in embedding_ids
            (no_decay if exempt else decay).append(param)
    o = cfg.optim
    groups = [
        {"params": decay, "weight_decay": o.weight_decay},
        {"params": no_decay, "weight_decay": 0.0},
    ]
    return torch.optim.AdamW(groups, lr=o.lr, betas=o.betas, eps=o.eps)


def build_scheduler(optimizer: torch.optim.Optimizer, cfg: TrainConfig) -> LambdaLR:
    """``LambdaLR`` over :func:`lr_factor`; call ``scheduler.step()`` once per optimizer step."""
    warmup, total, ratio = cfg.schedule.warmup_steps, cfg.steps, cfg.schedule.min_lr_ratio

    def factor(step: int) -> float:
        return lr_factor(step, warmup_steps=warmup, total_steps=total, min_lr_ratio=ratio)

    return LambdaLR(optimizer, factor)


class EMA:
    """Exponential moving average of a model's weights.

    ``shadow = decay_n * shadow + (1 - decay_n) * param`` after every optimizer step, where
    ``decay_n = decay`` or, with ``warmup``, ``min(decay, (1 + n) / (10 + n))`` for the
    ``n``-th update (0-based) -- a short memory early on so the average is useful before
    ``1 / (1 - decay)`` steps have passed. Buffers are copied, not averaged. The shadow is
    stored in float32 on the model's device.
    """

    def __init__(self, model: nn.Module, *, decay: float, warmup: bool) -> None:
        if not 0.0 <= decay < 1.0:
            raise ValueError("decay must be in [0, 1)")
        self.decay = decay
        self.warmup = warmup
        self.num_updates = 0
        self.shadow: dict[str, Tensor] = {
            name: tensor.detach().clone().float() for name, tensor in _state(model).items()
        }

    def current_decay(self) -> float:
        """Decay that the *next* :meth:`update` will use."""
        if not self.warmup:
            return self.decay
        n = self.num_updates
        return min(self.decay, (1.0 + n) / (10.0 + n))

    @torch.no_grad()
    def update(self, model: nn.Module) -> None:
        decay = self.current_decay()
        params = dict(model.named_parameters())
        for name, tensor in _state(model).items():
            shadow = self.shadow[name]
            if name in params:
                shadow.mul_(decay).add_(tensor.detach().float(), alpha=1.0 - decay)
            else:
                shadow.copy_(tensor.detach().float())
        self.num_updates += 1

    @torch.no_grad()
    def copy_to(self, model: nn.Module) -> None:
        """Overwrite ``model``'s weights with the averaged ones."""
        for name, tensor in _state(model).items():
            tensor.copy_(self.shadow[name].to(tensor.dtype))

    @contextmanager
    def swapped_in(self, model: nn.Module) -> Iterator[None]:
        """Temporarily load the EMA weights into ``model`` (restored on exit, also on error)."""
        backup = {name: t.detach().clone() for name, t in _state(model).items()}
        self.copy_to(model)
        try:
            yield
        finally:
            with torch.no_grad():
                for name, tensor in _state(model).items():
                    tensor.copy_(backup[name])

    def state_dict(self) -> dict[str, object]:
        return {
            "decay": self.decay,
            "warmup": self.warmup,
            "num_updates": self.num_updates,
            "shadow": {k: v.clone() for k, v in self.shadow.items()},
        }

    def load_state_dict(
        self, state: dict[str, object], *, allow_missing: Collection[str] = ()
    ) -> None:
        """Load ``state``; its names must be exactly this EMA's names.

        ``allow_missing`` names tensors that ``state`` must NOT contain (a fine-tune that
        adds parameters, :meth:`nida.train.loop.Trainer.init_from`): they keep their current
        shadow value. The names missing from ``state`` have to be exactly this set.

        Raises:
            ValueError: names, shapes or ``num_updates`` do not match.
        """
        shadow = state["shadow"]
        allowed = set(allow_missing)
        if not allowed <= set(self.shadow):
            raise ValueError(
                f"allow_missing names unknown to the EMA: {sorted(allowed - set(self.shadow))}"
            )
        if not isinstance(shadow, dict) or set(shadow) != set(self.shadow) - allowed:
            raise ValueError("EMA state does not match the model's parameter names")
        for name, tensor in shadow.items():
            if tensor.shape != self.shadow[name].shape:
                raise ValueError(f"EMA tensor {name!r} has shape {tuple(tensor.shape)}")
            self.shadow[name].copy_(tensor)
        num_updates = state["num_updates"]
        if not isinstance(num_updates, int):
            raise ValueError("EMA num_updates must be an int")
        self.num_updates = num_updates


def _state(model: nn.Module) -> dict[str, Tensor]:
    """Parameters and buffers by name, in a deterministic order."""
    out: dict[str, Tensor] = dict(model.named_parameters())
    out.update(model.named_buffers())
    return out