"""Minimal loader for the fitted reference operators (A_l, b_l).

The reference predicts each layer's next state from its current one:

    h_{l+1} ~= A_l @ h_l + b_l

and the *innovation* is the part it does not predict:

    eps_l = h_{l+1} - (A_l @ h_l + b_l)

Extract your own hidden states (`output_hidden_states=True`), pick the matching
`condition` and `variant`, and subtract. No dependency on the fitting pipeline.

    ops = load_operators("operators_gemma.pt")
    eps = innovation(ops, h_l, h_next, layer=12)          # h_*: [..., d]

Use `variant="raw"` for hidden_states as returned; `variant="normed"` if you feed
`input_layernorm_l(h_l)` (what block l actually reads — matters most for Gemma's
(1+gamma) RMSNorm). `condition` selects which fit: PT-COMP (base weights),
IT-COMP (instruct weights, no template), IT-CHAT (instruct + chat template).
"""
import torch


def load_operators(path):
    """Load an operators_<family>.pt bundle (dict: meta + per-condition A/b)."""
    return torch.load(path, map_location="cpu", weights_only=False)


def _block(ops, condition, variant):
    if condition not in ops:
        raise KeyError(f"condition {condition!r} not in {[k for k in ops if k!='meta']}")
    return ops[condition][variant]


def predict_next(ops, h, layer, condition="IT-COMP", variant="raw"):
    """A_l @ h + b_l for h of shape [..., d]. Returns the predicted h_{l+1}."""
    blk = _block(ops, condition, variant)
    A = blk["A"][layer].to(h.dtype).to(h.device)
    b = blk["b"][layer].to(h.dtype).to(h.device)
    return h @ A.T + b


def innovation(ops, h_l, h_next, layer, condition="IT-COMP", variant="raw"):
    """eps_l = h_{l+1} - (A_l @ h_l + b_l). h_l, h_next: [..., d]."""
    return h_next - predict_next(ops, h_l, layer, condition, variant)


if __name__ == "__main__":
    import sys
    ops = load_operators(sys.argv[1] if len(sys.argv) > 1 else "operators_gemma.pt")
    m = ops["meta"]
    print(f"family={m['family']} d={m['d']} n_layers={m['n_layers']}")
    print(f"conditions={list(m['conditions'])} variants={list(m['variants'])}")
    d, nL = m["d"], m["n_layers"]
    h = torch.randn(4, d)                       # 4 fake token states
    eps = innovation(ops, h, torch.randn(4, d), layer=nL // 2)
    print(f"demo innovation shape: {tuple(eps.shape)}")