"""Contrast readout floor: a standalone check for equilibrium-propagation-style code.

WHAT THIS DETECTS
-----------------
Learning rules in the equilibrium-propagation / coupled-learning family read the weight
update from a *contrast*: a small perturbation d applied on top of a much larger state o.
A very common implementation forms the perturbed state and then recovers the perturbation
by subtracting the state back out:

    z = o + d                      # perturbed state, stored in float32
    contrast = 0.5 * (z - o) ** 2  # read the contrast, differentiate w.r.t. parameters

Algebraically z - o == d. In floating point it does not. Forming z rounds away every
component of d below eps * |o| (absorption). The subtraction z - o is then exact, so it
faithfully returns the *already damaged* copy of d and raises no error. The damage grows as
the displacement-to-activation ratio falls, which happens as models get wider or deeper, so
it imitates an algorithmic scaling wall.

The fix costs nothing: read the contrast from the stored perturbation itself. The linear
form -<d, o> has the same derivative with respect to the parameters and never touches the
rounded copy.

HOW TO USE
----------
1. Zero-setup demo of the mechanism and the fix:

       python contrast_floor_check.py

2. Check your own tensors. Pass the state and the perturbation you actually use:

       from contrast_floor_check import check_pair
       check_pair(o, d, name="block3")

3. Scan a whole model. Collect (state, perturbation) pairs per layer and call:

       from contrast_floor_check import report
       report({"block0": (o0, d0), "block1": (o1, d1), ...})

INTERPRETATION
--------------
The headline quantity is the displacement ratio RMS(d) / RMS(o). Treat it as a screening
number rather than a verdict, for two reasons.

Two thresholds, worth separating. The recovered perturbation first picks up error that is
close to orthogonal (`recovered_rel_err` rises while `recovered_projection` stays near one).
That part behaves like noise and averages over steps. Only further down does the surviving
perturbation shrink along its own direction (`recovered_projection` falls), and that part is
a consistent shortfall which accumulates across training. The projection is the number to
watch, and it turns later than the error does.

Synthetic pairs underestimate. Running this on random tensors at a given RMS ratio gives a
lower bound on the damage a real implementation suffers at the same ratio, because a real
system re-forms the perturbed state several times per step and transmits the perturbation
down a chain that has already been rounded, so the loss compounds, and because what matters
is the share of the *update-carrying* mass sitting under the grid rather than the share of
elements. In our transformer language models the per-layer onset sat near 5e-7 in float32,
about four times epsilon, and layers an order of magnitude below that were missing most of
their contribution to the parameter update, far more than a matched synthetic pair loses.
Measure with your own tensors, and if you can, measure the parameter update itself against a
double precision or backprop reference.

This file is standalone: only PyTorch is required, and it is released for anyone to run
against their own implementation.
"""

import torch


def _rms(t):
    return float(t.detach().float().pow(2).mean().sqrt())


def check_pair(o, d, name="", dtype=None, verbose=True):
    """Measure how much of the perturbation d survives being added to the state o.

    o     state tensor the perturbation is applied on top of
    d     perturbation tensor (the quantity the contrast is supposed to carry)
    dtype storage dtype to test; defaults to o's dtype

    Returns a dict with the displacement ratio, the fraction of elements below epsilon,
    and the relative error of the recovered perturbation against the exact one.
    """
    o = o.detach()
    d = d.detach()
    st = dtype or o.dtype
    eps = torch.finfo(st).eps

    o64, d64 = o.double(), d.double()
    ratio = _rms(d64) / max(_rms(o64), 1e-300)

    # elementwise exposure: how many components sit under the storage resolution
    rel = (d64.abs() / o64.abs().clamp_min(1e-300)).flatten()
    frac_below = float((rel < eps).float().mean())
    q10 = float(rel.kthvalue(max(1, int(0.10 * rel.numel()))).values)

    # what the subtract-back-out readout actually returns
    z = (o.to(st) + d.to(st))
    d_recovered = (z - o.to(st)).double()
    err = float((d_recovered - d64).norm() / max(float(d64.norm()), 1e-300))

    # the damage that matters is the signed part along the true perturbation
    proj = float((d_recovered * d64).sum() / d64.pow(2).sum().clamp_min(1e-300))

    out = dict(name=name, dtype=str(st), eps=eps, disp_ratio=ratio,
               frac_below_eps=frac_below, ratio_q10=q10,
               recovered_rel_err=err, recovered_projection=proj)
    if verbose:
        flag = "OK " if err < 1e-3 else ("WARN" if err < 0.1 else "DAMAGED")
        print(f"[{flag}] {name or 'pair':<16} ratio {ratio:.2e}  "
              f"below-eps {100*frac_below:5.1f}%  q10 {q10:.2e}  "
              f"recovered err {err:.3f}  projection {proj:.3f}")
    return out


def report(pairs, dtype=None):
    """Run check_pair over a dict of {name: (state, perturbation)} and summarize."""
    eps = torch.finfo(dtype or torch.float32).eps
    print(f"contrast readout floor check   storage eps = {eps:.2e}")
    print(f"{'':18}{'ratio':>10}{'<eps':>10}{'err':>9}{'proj':>8}")
    rows = [check_pair(o, d, name=k, dtype=dtype, verbose=False) for k, (o, d) in pairs.items()]
    for r in rows:
        flag = "ok" if r["recovered_rel_err"] < 1e-3 else ("warn" if r["recovered_rel_err"] < 0.1 else "DAMAGED")
        print(f"{r['name']:<18}{r['disp_ratio']:>10.2e}{100*r['frac_below_eps']:>9.1f}%"
              f"{r['recovered_rel_err']:>9.3f}{r['recovered_projection']:>8.3f}  {flag}")
    worst = max(rows, key=lambda r: r["recovered_rel_err"])
    if worst["recovered_rel_err"] >= 0.1:
        print(f"\nThe readout loses a large part of the perturbation, worst at "
              f"'{worst['name']}' with {100*(1-worst['recovered_projection']):.0f}% of it missing. "
              f"Read the contrast from the stored perturbation instead of differencing the states.")
    elif worst["recovered_rel_err"] >= 1e-3:
        print(f"\nSome loss at '{worst['name']}'. Watch this as models get wider, since the "
              f"displacement ratio falls with width.")
    else:
        print("\nNo material loss at these magnitudes.")
    return rows


def _demo():
    torch.manual_seed(0)
    print(__doc__.split("HOW TO USE")[0].strip()[:0] or "", end="")
    print("Demonstration: one state, perturbations spanning eight decades.\n")
    o = torch.randn(4096, 512)
    pairs = {}
    for k in range(4, 12):
        scale = 10.0 ** (-k)
        pairs[f"ratio 1e-{k}"] = (o, torch.randn_like(o) * scale)
    report(pairs, dtype=torch.float32)

    print("\nSame perturbations, but the contrast is read from the stored tensor.")
    print("This is the fix: the derivative is identical and nothing is rounded away.\n")
    o_ = o.float()
    for k in (4, 8, 11):
        d = torch.randn_like(o) * (10.0 ** (-k))
        # damaged path: differentiate 0.5*||z-o||^2  ->  cotangent is (z-o)
        cot_bad = ((o_ + d.float()) - o_).double()
        # fixed path: differentiate -<d,o>            ->  cotangent is d itself
        cot_good = d.double()
        ref = d.double()
        e_bad = float((cot_bad - ref).norm() / ref.norm())
        e_good = float((cot_good - ref).norm() / ref.norm())
        print(f"  ratio 1e-{k:<3} subtract-back-out error {e_bad:8.3f}   stored-perturbation error {e_good:.3e}")

    print("\nIf your code differences two large states to obtain a small contrast, the left "
          "column is what your update carries.")


if __name__ == "__main__":
    _demo()