from __future__ import annotations

"""Gradient monitoring for training diagnostics.

Computes per-group gradient norms after backward pass to detect
vanishing/exploding gradients in different components:
  - MONA adapters (drone + satellite)
  - LoRA (DGTRS-CLIP text encoder)
  - TextFusionMLP
  - GatedFusion gates (alpha_q, alpha_g)
  - Learnable temperature (logit_scale)
"""

import logging

import torch
import torch.nn as nn

LOGGER = logging.getLogger("caption_test.grad_monitor")

# Parameter group patterns for classification.
_GROUP_PATTERNS = {
    "mona_drone": "drone_encoder",
    "mona_sat": "sat_encoder",
    "lora_text": "text_encoder",
    "text_fusion_mlp": "text_fusion",
    "gate_q": "fusion_query",
    "gate_g": "fusion_gallery",
}


def compute_gradient_norms(
    model: nn.Module,
    loss_fn: nn.Module | None = None,
) -> dict[str, float]:
    """Compute L2 gradient norms per parameter group.

    Args:
        model: The model after backward().
        loss_fn: Loss module (for logit_scale gradient).

    Returns:
        Dict mapping group name to gradient L2 norm.
    """
    group_grads: dict[str, list[torch.Tensor]] = {k: [] for k in _GROUP_PATTERNS}
    group_grads["other"] = []

    for name, param in model.named_parameters():
        if not param.requires_grad or param.grad is None:
            continue

        matched = False
        for group_name, pattern in _GROUP_PATTERNS.items():
            if pattern in name:
                group_grads[group_name].append(param.grad.detach().flatten())
                matched = True
                break
        if not matched:
            group_grads["other"].append(param.grad.detach().flatten())

    norms: dict[str, float] = {}
    for group_name, grads in group_grads.items():
        if grads:
            all_grads = torch.cat(grads)
            norms[f"grad_norm/{group_name}"] = float(all_grads.norm(2).item())
            norms[f"grad_max/{group_name}"] = float(all_grads.abs().max().item())

    # Logit scale gradient (from loss module).
    if loss_fn is not None and hasattr(loss_fn, "logit_scale"):
        ls = loss_fn.logit_scale
        if ls is not None and ls.grad is not None:
            norms["grad_norm/logit_scale"] = float(ls.grad.detach().abs().item())

    return norms


def log_gradient_summary(norms: dict[str, float]) -> None:
    """Log gradient norms summary to console."""
    norm_parts = []
    for k, v in sorted(norms.items()):
        if k.startswith("grad_norm/"):
            name = k.replace("grad_norm/", "")
            norm_parts.append(f"{name}={v:.4f}")
    if norm_parts:
        LOGGER.info("grad norms: %s", "  ".join(norm_parts))
