"""Moment-Sharp spectral control for Linear layers.

This module is a training intervention, not a replacement for the mandated bench
harness.  K=2 uses the exact nonnegative moment extremizer and is differentiable
only through the post-step rescaling operation (the controller itself is kept
outside autograd).
"""
import math
import torch
from torch import nn


def u2_from_moments(m1, m2, d, eps=1e-12):
    """Maximum possible eigenvalue given sum and sum of squares."""
    d = int(d)
    if d <= 1:
        return max(float(m1), eps)
    disc = max(0.0, d * float(m2) - float(m1) ** 2)
    return max(float(m1) / d + math.sqrt((d - 1) * disc) / d, eps)


def exact_or_hutchinson_moments(weight, probes=8, generator=None):
    """Estimate tr(A), tr(A^2), A=W^T W, with exact small-matrix fallback.

    Hutchinson is used for larger matrices to keep the mechanism aligned with
    the proposed implementation; returned values are detached Python floats.
    """
    w = weight.detach()
    d = w.shape[1]
    if d <= 256:
        a = w.T @ w
        return float(torch.trace(a)), float(torch.trace(a @ a)), d
    gen = generator or torch.Generator(device=w.device)
    if generator is None:
        gen.manual_seed(12345)
    z = torch.randint(0, 2, (probes, d), device=w.device, generator=gen,
                      dtype=torch.int64).to(w.dtype).mul_(2).sub_(1)
    az = (z @ w.T) @ w
    a2z = (az @ w.T) @ w
    return float((z * az).sum(1).mean()), float((z * a2z).sum(1).mean()), d


@torch.no_grad()
def moment_sharp_rescale(model, target_sigma=2.0, probes=8, ema=None,
                         generator=None):
    """Rescale each Linear weight when its K=2 certified bound exceeds target."""
    observed = []
    for layer in model.modules():
        if not isinstance(layer, nn.Linear):
            continue
        m1, m2, d = exact_or_hutchinson_moments(layer.weight, probes, generator)
        key = id(layer)
        if ema is not None:
            old = ema.get(key, (m1, m2))
            m1, m2 = .9 * old[0] + .1 * m1, .9 * old[1] + .1 * m2
            ema[key] = (m1, m2)
        bound_sq = u2_from_moments(m1, m2, d)
        before = float(torch.linalg.matrix_norm(layer.weight, 2))
        if bound_sq > target_sigma ** 2:
            layer.weight.mul_(target_sigma / math.sqrt(bound_sq + 1e-12))
        after = float(torch.linalg.matrix_norm(layer.weight, 2))
        observed.append({"bound_sigma": math.sqrt(bound_sq),
                         "true_sigma_before": before,
                         "true_sigma_after": after})
    return observed


def mechanism_signature(model, target_sigma=2.0):
    """Measured NN-scale signature: bound, observed sigma, and certified status."""
    rows = []
    for layer in model.modules():
        if isinstance(layer, nn.Linear):
            m1, m2, d = exact_or_hutchinson_moments(layer.weight, probes=32)
            b = math.sqrt(u2_from_moments(m1, m2, d))
            s = float(torch.linalg.matrix_norm(layer.weight.detach(), 2))
            rows.append((b, s))
    return {"predicted_bound_sigma": [x[0] for x in rows],
            "observed_true_sigma": [x[1] for x in rows],
            "max_bound_minus_observed": max((b-s for b,s in rows), default=0.),
            "target_sigma": target_sigma,
            "confirmed": bool(all(b + 1e-5 >= s for b,s in rows))}