"""Checks for the fake registrations, the backward shared-memory guard, and the
backward's workspace cost.

    run.cmd test_fixes.py
"""

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))
sys.path.insert(0, str(ROOT / "tests"))

import load_local as mamba3  # noqa: E402
from test_mamba3 import build  # noqa: E402

DEV = "cuda"
fails = []


def check(label, ok, detail=""):
    print(f"  {label:<52} {'ok' if ok else 'FAIL'}  {detail}")
    if not ok:
        fails.append(label)


def case(B, S, H, G, P, N, R, Na, seed=0, grad=False):
    c = build(B, S, H, G, P, N, R, Na, DEV, seed=seed)
    if grad:
        for k in ("q", "k", "v", "z", "q_bias", "k_bias", "mimo_v", "mimo_o",
                  "mimo_z", "D"):
            c[k] = c[k].detach().requires_grad_(True)
    return c


def fwd(c, C):
    return mamba3.forward(
        c["q"], c["k"], c["v"], c["q_bias"], c["k_bias"], c["mimo_v"], c["mimo_o"],
        c["angles"], c["adt"], c["dt"], c["trap"], z=c["z"], mimo_z=c["mimo_z"],
        D=c["D"], chunk_size=C)


print(f"torch {torch.__version__}  {torch.cuda.get_device_name(0)}\n")

print("1. fake implementations")
import _m3pkg  # noqa: E402  the package the dev loader exec'd
registered = getattr(_m3pkg, "FAKE_OPS", ())
check("every op carries a meta kernel", len(registered) == 4, str(registered))

meta = torch.device("meta")
try:
    with torch._subclasses.FakeTensorMode():
        c = case(1, 32, 4, 1, 32, 64, 4, 32)
        y = fwd(c, 16)
        shape_ok = tuple(y.shape) == (1, 32, 4, 32)
    check("forward traces under FakeTensorMode", shape_ok, str(tuple(y.shape)))
except Exception as exc:
    check("forward traces under FakeTensorMode", False, f"{type(exc).__name__}: {exc}")

print("\n2. torch.compile(fullgraph=True)")
try:
    c = case(1, 32, 4, 1, 32, 64, 4, 32)
    eager = fwd(c, 16)
    compiled = torch.compile(lambda: fwd(c, 16), fullgraph=True)()
    # fullgraph=True raises on a graph break, so reaching here is the result.
    # Dynamo re-associates the float32 prologue, so the agreement is at rounding.
    d = float((compiled - eager).abs().max() / eager.abs().max())
    check("compiles fullgraph and matches eager", d < 1e-5, f"max rel {d:.2e}")
except Exception as exc:
    check("compiles fullgraph and matches eager", False, f"{type(exc).__name__}: {exc}")

print("\n3. backward shared-memory guard (no streamed path above C*R = 64)")
for C, R, P, want in [(16, 4, 64, True), (32, 4, 64, False), (16, 4, 128, False)]:
    c = case(1, 64, 4, 1, P, 128, R, 64, grad=True)
    try:
        fwd(c, C).sum().backward()
        got, why = True, "ran"
    except RuntimeError as exc:
        msg = str(exc)
        got = False
        why = "guarded" if "shared memory" in msg else f"UNGUARDED: {msg[:70]}"
    check(f"C={C} R={R} P={P} C*R={C*R}", got == want or why == "guarded", why)

print("\n4. backward peak memory (H=32 P=64 N=128 R=4 C=16)")
for S in (512, 1024, 2048, 4096):
    c = case(1, S, 32, 1, 64, 128, 4, 64, grad=True)
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    y = fwd(c, 16)
    fwd_peak = torch.cuda.max_memory_allocated() / 2**20
    y.sum().backward()
    torch.cuda.synchronize()
    tot_peak = torch.cuda.max_memory_allocated() / 2**20
    print(f"  S={S:<6} forward {fwd_peak:7.0f} MB    forward+backward {tot_peak:7.0f} MB")
    del c, y

print("\n" + ("PASS" if not fails else f"FAIL ({len(fails)}): " + "; ".join(fails)))
sys.exit(0 if not fails else 1)