"""
utils.metrics — generic, reusable metric functions.

Shared by every experiment in experiments/. Anything specific to a
particular dataset or training recipe stays out of this module.
"""

from __future__ import annotations

import numpy as np
import torch
import torch.nn.functional as F


@torch.no_grad()
def accuracy(model, loader, device):
    model.eval()
    correct = total = 0
    for imgs, labels in loader:
        imgs   = imgs.to(device)
        labels = labels.squeeze().long().to(device)
        preds  = model(imgs).argmax(1)
        correct += (preds == labels).sum().item()
        total   += labels.size(0)
    return correct / total


@torch.no_grad()
def weight_norm(model):
    return sum(p.data.norm(2).item() ** 2 for p in model.parameters()) ** 0.5


@torch.no_grad()
def feature_rank(model, loader, device, n=200, hook_module_attr="avgpool"):
    """
    Effective rank of penultimate-layer features = exp(entropy of normalised
    singular values). Fan et al. 2024 — the most reliable progress measure
    for grokking. The rank collapses at the transition.

    `hook_module_attr` is the attribute name on the model whose forward output
    we treat as the penultimate representation. Defaults to ResNet's avgpool.
    """
    model.eval()
    feats = []
    target = getattr(model, hook_module_attr)
    hook = target.register_forward_hook(
        lambda m, i, o: feats.append(o.view(o.size(0), -1).cpu())
    )
    count = 0
    for imgs, _ in loader:
        model(imgs.to(device))
        count += imgs.size(0)
        if count >= n:
            break
    hook.remove()
    F_mat = torch.cat(feats)[:n]
    try:
        _, s, _ = torch.svd(F_mat)
        s = s / (s.sum() + 1e-10)
        return torch.exp(-(s * torch.log(s + 1e-10)).sum()).item()
    except Exception:
        return float("nan")


def irm_penalty(model, envs, device):
    """
    IRMv1 penalty (Arjovsky et al. 2019). For each environment, the squared
    gradient of the loss w.r.t. a dummy scalar w=1 in the logits.

    LOW value  = invariant predictor = causal features found.
    HIGH value = environment-specific (spurious) features still in use.

    `envs` is a list of dicts {"x": tensor, "y": tensor} already on `device`.
    Returns (mean, var) over environments.
    """
    model.eval()
    penalties = []
    for env in envs:
        w      = torch.tensor(1.0, requires_grad=True, device=device)
        logits = model(env["x"]) * w
        loss   = F.cross_entropy(logits, env["y"])
        grad   = torch.autograd.grad(loss, w, create_graph=False)[0]
        penalties.append(grad.item() ** 2)
    t = torch.tensor(penalties)
    return t.mean().item(), t.var().item()


@torch.no_grad()
def shortcut_ratio(model, loader, device):
    """
    Border-confidence / center-confidence proxy for artifact reliance.
    Ratio > 1 means the model trusts the borders (where scanner markers,
    laterality letters, and other artefacts live) more than the center
    (where actual anatomy is). On CheXpert, replace with GradCAM pointed at
    known artifact locations vs. anatomical regions.

    Returns (center_conf_mean, border_conf_mean).
    """
    model.eval()
    cc, bc = [], []
    for imgs, _ in loader:
        imgs = imgs.to(device)
        B, C, H, W = imgs.shape
        hs, he = H // 4, 3 * H // 4
        ws, we = W // 4, 3 * W // 4
        center = F.interpolate(imgs[:, :, hs:he, ws:we], size=(H, W),
                               mode='bilinear', align_corners=False)
        border = imgs.clone(); border[:, :, hs:he, ws:we] = 0.
        cc.append(F.softmax(model(center), 1).max(1).values.mean().item())
        bc.append(F.softmax(model(border), 1).max(1).values.mean().item())
    return float(np.mean(cc)), float(np.mean(bc))