"""Branch-gain profile rho_k and the NSR constant c(scale) = L * mean(rho) (THEORY.md T2/corollary).
Validation-only: uses the simulation autograd for exact J_F^T v."""
import torch
import torch.nn as nn
from .zbp import zbp_blocks, ZBPConfig


@torch.enable_grad()
def rho_profile(model, x, y):
    blocks = zbp_blocks(model)
    cfgs = [b.cfg for b in blocks]
    for b in blocks:
        b.cfg = b.cfg.replace(mode="exact")
        b.capture = True
    model.zero_grad(set_to_none=True)
    nn.functional.cross_entropy(model(x).flatten(0, 1), y.flatten()).backward()
    out = {}
    for b, c in zip(blocks, cfgs):
        b.cfg = c
        b.capture = False
        b.captured = None
    # rho via a recorded pass: |J_F^T v| / |v| per block from the Recorder
    from .zbp import Recorder
    for b in blocks:
        b.cfg = b.cfg.replace(mode="zero")
    with Recorder() as rec:
        model.zero_grad(set_to_none=True)
        nn.functional.cross_entropy(model(x).flatten(0, 1), y.flatten()).backward()
    for b, c in zip(blocks, cfgs):
        b.cfg = c
    model.zero_grad(set_to_none=True)
    rows = rec.summary()
    for name, r in rows.items():
        out[name] = (r["gnorm"] / max(r["vnorm"], 1e-30)) ** 2 if "vnorm" in r else None
    return out