# Gradient-norm clipping callback.
# Fires in the on_before_optimizer_step hook so clipping runs after all
# grad accumulation and GradScaler.unscale_ but before optimizer.step.

from __future__ import annotations

import torch

from imaginaire.utils import log
from imaginaire.utils.callback import Callback


class GradClipCallback(Callback):
    """Clip global gradient L2 norm to `max_norm`. Logs the pre-clip norm to wandb.

    :param max_norm: clip threshold. 1.0 is a reasonable default for diffusion.
    :param log_every_n: iterations between wandb scalar log emissions (gradient norm).
    """

    def __init__(self, max_norm: float = 1.0, log_every_n: int = 10) -> None:
        super().__init__()
        self.max_norm = float(max_norm)
        self.log_every_n = int(log_every_n)

    def on_before_optimizer_step(self, model_ddp, optimizer, scheduler,
                                 grad_scaler, iteration: int = 0) -> None:
        # Unscale fp16/bf16 grads so clip_grad_norm_ sees true magnitudes.
        # Cosmos training does not use a true GradScaler for bf16 so this is a no-op,
        # but keep the call for safety if float16 is ever enabled upstream.
        try:
            grad_scaler.unscale_(optimizer)
        except (AttributeError, RuntimeError):
            pass

        params = [p for g in optimizer.param_groups for p in g['params']
                  if p.grad is not None]
        if not params:
            return

        total_norm = torch.nn.utils.clip_grad_norm_(params, max_norm=self.max_norm)

        if iteration % self.log_every_n == 0:
            total_norm_f = float(total_norm) if torch.isfinite(total_norm) else float('nan')
            log.debug(f'[GradClip] iter {iteration} grad_norm={total_norm_f:.4f} '
                      f'(clip@{self.max_norm})')
            # Best-effort wandb log; imaginaire wires metrics through trainer hooks,
            # so we only push to the logger if one is attached.
            try:
                from imaginaire.utils import logger as _logger
                _logger.log_metrics({'grad_norm': total_norm_f, 'samples': iteration})
            except Exception:
                pass