"""#2 — leading-abscissa control for the ept non-conservative operator (ports the aep-dynamics
'control the LEADING spectral signal, surgically' finding to C512).

Why not jacreg: jacreg penalizes ||J_nc||_F^2 (Hutchinson) — the WHOLE Jacobian norm. That is blunt:
it over-constrains directions that never threaten stability, and when the controller ramps it high it
HIJACKS the task gradient (the known jr-hijack failure). The aep leading-vs-lagging result says the
right knob is the leading SPECTRAL ABSCISSA, not the norm.

What we control: the numerical abscissa  omega(J_nc) = lambda_max( (J_nc + J_nc^T)/2 )  = the 2-norm
log-norm mu_2(J_nc) = one-sided Lipschitz constant. It upper-bounds the spectral abscissa, governs the
transient growth ||e^{Jt}||, and its crossing past (1+c) IS the free-phase Hopf (J_F = J_nc - (1+c)I).
Power iteration on the SYMMETRIC PART -> matvec-only (jvp+vjp of nc_force, the same primitives jacreg
and the AEP correction already call), so it scales and is analog-compatible (no eigendecomposition).
One-sided ReLU penalty = a LEADING signal: acts only as the abscissa nears the margin, so unlike jacreg
it does not over-contract / hijack when the operator is already safe.
"""
import torch
from torch.autograd.functional import jvp, vjp


def num_abscissa(blk, zs, cache, iters=3):
    """Power-iterate Sym(J_nc)=(J_nc+J_nc^T)/2 at zs for the leading eigenpair. Returns (v_detached, lambda_float).
    lambda = v^T Sym(J_nc) v = v^T J_nc v (Rayleigh quotient at the leading eigenvector) = numerical abscissa."""
    z = zs.detach()
    v = cache.get('v')
    if v is None or v.shape != z.shape or v.dtype != z.dtype or v.device != z.device:
        v = torch.randn_like(z)
    v = v / (v.norm() + 1e-12)
    with torch.no_grad():
        for _ in range(iters):
            Sv = 0.5 * (jvp(blk.nc_force, z, v)[1] + vjp(blk.nc_force, z, v)[1])   # Sym(J_nc) v
            v = Sv / (Sv.norm() + 1e-12)
        lam = float((v * jvp(blk.nc_force, z, v)[1]).sum() / (v * v).sum())        # v^T J_nc v
    cache['v'] = v
    return v, lam


def eig_penalty(blk, zs, eigreg, margin, cache, iters=3):
    """Grads of the one-sided leading-abscissa penalty  R = eigreg * relu(omega(J_nc) - margin)^2.
    Returns ({id(p): grad}, omega) — omega logged as the controller signal. Empty grads when below margin."""
    v, lam0 = num_abscissa(blk, zs, cache, iters)
    if lam0 <= margin:                                     # below the stability margin: leading signal off
        return {}, lam0
    z = zs.detach()
    with torch.enable_grad():
        Jv = jvp(blk.nc_force, z, v, create_graph=True)[1]           # differentiable in theta (nc_force params)
        lam = (v * Jv).sum() / (v * v).sum()                        # numerical abscissa, v fixed
        R = eigreg * torch.relu(lam - margin) ** 2
        gr = torch.autograd.grad(R, blk.block, allow_unused=True)
    return {id(p): g for p, g in zip(blk.block, gr) if g is not None}, lam0