#!/usr/bin/env python3
"""Focal-SAM (ICML 2025) — sharpness-aware minimization with focal class weighting.

SAM: w' = w + ρ · g/||g||; compute loss @ w'; step from w.
Focal: instead of uniform ρ, scale perturbation per-class by focal weight
(1 - p_c)^γ × n_c^(-α), penalizing tail classes more strongly.

Reference: Li et al. "Focal-SAM: Focal Sharpness-Aware Minimization for
Long-Tailed Classification", ICML 2025.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

class FocalSAM:
    """Two-pass SAM wrapper.

    Usage:
        sam = FocalSAM(optimizer, rho=0.05, class_counts=[...], gamma=2.0, alpha=0.25)
        # First pass:
        loss = criterion(model(x), y); loss.backward()
        sam.first_step(zero_grad=True)
        # Recompute forward+loss @ perturbed weights:
        loss2 = criterion(model(x), y); loss2.backward()
        sam.second_step(zero_grad=True)
    """
    def __init__(self, optimizer, rho=0.05, class_counts=None,
                 gamma=2.0, alpha=0.25, device='cuda'):
        self.opt = optimizer
        self.rho = rho
        self.gamma = gamma
        self.alpha = alpha
        if class_counts is not None:
            cc = np.asarray(class_counts, dtype=np.float64).clip(min=1)
            # Focal-like weight: tail classes get larger ρ
            # w_c = (cc_head / cc_c)^alpha, normalized to mean=1
            w = (cc.max() / cc) ** alpha
            w = w / w.mean()
            self.register_weights = torch.tensor(w, dtype=torch.float32, device=device)
        else:
            self.register_weights = None
        # Buffers for perturbation
        self._e_w = []

    def _grad_norm(self):
        norms = []
        for group in self.opt.param_groups:
            for p in group['params']:
                if p.grad is None: continue
                norms.append(p.grad.norm(p=2))
        return torch.norm(torch.stack(norms), p=2) if norms else torch.tensor(0.0)

    @torch.no_grad()
    def first_step(self, zero_grad=False, focal_scale=1.0):
        """Perturb weights in direction of gradient."""
        grad_norm = self._grad_norm()
        scale = self.rho * focal_scale / (grad_norm + 1e-12)
        self._e_w = []
        for group in self.opt.param_groups:
            for p in group['params']:
                if p.grad is None:
                    self._e_w.append(None); continue
                e_w = p.grad * scale
                p.add_(e_w)
                self._e_w.append(e_w)
        if zero_grad: self.opt.zero_grad(set_to_none=True)

    @torch.no_grad()
    def second_step(self, zero_grad=False):
        """Restore weights and step with new gradient (computed at perturbed)."""
        i = 0
        for group in self.opt.param_groups:
            for p in group['params']:
                if self._e_w[i] is not None:
                    p.sub_(self._e_w[i])
                i += 1
        self.opt.step()
        if zero_grad: self.opt.zero_grad(set_to_none=True)

    def compute_focal_scale(self, targets):
        """Given a batch of targets, compute mean focal weight for rho scaling.

        Rho scales batch-wise by average of per-sample focal weights.
        """
        if self.register_weights is None: return 1.0
        w_per_sample = self.register_weights[targets]
        return float(w_per_sample.mean().item())

if __name__ == '__main__':
    # Minimal smoke test
    torch.manual_seed(0)
    m = nn.Sequential(nn.Linear(10, 20), nn.ReLU(), nn.Linear(20, 5))
    opt = torch.optim.AdamW(m.parameters(), lr=1e-3)
    cc = [1000, 500, 100, 30, 5]  # long-tail
    sam = FocalSAM(opt, rho=0.05, class_counts=cc, device='cpu')
    x = torch.randn(8, 10); y = torch.tensor([0,0,1,1,2,3,4,4])
    fs = sam.compute_focal_scale(y); print(f"focal scale: {fs:.3f}")
    loss = F.cross_entropy(m(x), y); loss.backward()
    sam.first_step(zero_grad=True, focal_scale=fs)
    loss2 = F.cross_entropy(m(x), y); loss2.backward()
    sam.second_step(zero_grad=True)
    print("SAM step OK")