"""Backward reproducibility: repeated runs must be bitwise identical."""
import sys
from pathlib import Path

import torch
import torch.nn.functional as F

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))
import load_local as mamba3  # noqa: E402

dev = "cuda"
NAMES = ["q", "k", "v", "z", "q_bias", "k_bias", "mimo_v", "mimo_o", "mimo_z", "D"]
fails = []

for S, H, G, P, N, R, C in [(512, 16, 1, 64, 128, 4, 16), (256, 8, 2, 32, 64, 2, 32)]:
    B, Na = 1, N // 2
    torch.manual_seed(0)
    f = lambda *s: torch.randn(*s, device=dev)
    shp = {"q": (B, S, R, G, N), "k": (B, S, R, G, N), "v": (B, S, H, P),
           "z": (B, S, H, P), "q_bias": (H, R, N), "k_bias": (H, R, N),
           "mimo_v": (H, R, P), "mimo_o": (H, R, P), "mimo_z": (H, R, P), "D": (H,)}
    base = {n: f(*shp[n]) for n in NAMES}
    dtv = F.softplus(-3.0 + f(B, H, S))
    adtv = -F.softplus(f(B, H, S)).clamp(max=-1e-4) * dtv
    trv = torch.rand(B, H, S, device=dev) * 0.5
    thv = torch.rand(B, S, H, Na, device=dev) * 0.4
    dy = torch.randn(B, S, H, P, device=dev)

    def once():
        t = {n: base[n].detach().clone().requires_grad_(True) for n in NAMES}
        sch = {n: x.clone().requires_grad_(True) for n, x in
               {"dt": dtv, "adt": adtv, "trap": trv, "th": thv}.items()}
        y = mamba3.forward(t["q"], t["k"], t["v"], t["q_bias"], t["k_bias"], t["mimo_v"],
                           t["mimo_o"], sch["th"], sch["adt"], sch["dt"], sch["trap"],
                           z=t["z"], mimo_z=t["mimo_z"], D=t["D"], chunk_size=C)
        y.backward(dy)
        g = {n: t[n].grad.clone() for n in NAMES}
        g.update({n: sch[n].grad.clone() for n in sch})
        return y.detach().clone(), g

    y1, g1 = once()
    print(f"S={S} H={H} G={G} P={P} N={N} R={R} C={C}")
    for trial in range(3):
        y2, g2 = once()
        if not torch.equal(y1, y2):
            fails.append(f"forward@S{S}")
        for n in g1:
            if not torch.equal(g1[n], g2[n]):
                fails.append(f"d{n}@S{S}")
    bad = sorted(set(x for x in fails if f"@S{S}" in x))
    print(f"   forward + 14 gradients over 3 repeats: "
          f"{'all bitwise identical' if not bad else 'DIFFER: ' + ', '.join(bad)}")

print("\nPASS" if not fails else "\nFAIL: " + ", ".join(sorted(set(fails))))
sys.exit(0 if not fails else 1)