"""Streamed-intra path (C*R > 64) against the upstream reference and small chunks."""
import argparse
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 / "ref"))
import load_local as mamba3  # noqa: E402
from anchor import load_upstream_refs, max_rel_err  # noqa: E402

ap = argparse.ArgumentParser(); ap.add_argument("--mamba-tests", required=True)
up = load_upstream_refs(ap.parse_args().mamba_tests)
dev = "cuda"
worst, fails = 0.0, []

# (S, H, G, P, N, R, C) with C*R > 64 so the streamed path is taken
GRID = [(256, 8, 1, 64, 128, 4, 64), (256, 8, 1, 32, 64, 4, 32),
        (512, 4, 1, 64, 128, 4, 64), (128, 8, 2, 32, 64, 2, 64),
        (256, 8, 1, 64, 128, 1, 128), (192, 4, 1, 64, 128, 8, 32)]

for S, H, G, P, N, R, C in GRID:
    B, Na = 1, N // 2
    torch.manual_seed(0)
    f = lambda *s: torch.randn(*s, device=dev)
    dt = F.softplus(-3.0 + f(B, H, S))
    adt = -F.softplus(f(B, H, S)).clamp(max=-1e-4) * dt
    trap = torch.rand(B, H, S, device=dev) * 0.5
    raw = torch.rand(B, S, H, Na, device=dev)
    q, k, v, z = f(B, S, R, G, N), f(B, S, R, G, N), f(B, S, H, P), f(B, S, H, P)
    qb, kb = f(H, R, N), f(H, R, N)
    mv, mo, mz = (torch.rand(H, R, P, device=dev) / R for _ in range(3))
    D = f(H)

    got = mamba3.forward(q, k, v, qb, kb, mv, mo, raw, adt, dt, trap,
                         z=z, mimo_z=mz, D=D, chunk_size=C)
    ang = mamba3.cumulative_angles(raw, dt)
    dcs, dcsr = mamba3.chunk_decay(adt, C)
    ref, _, _ = up.mamba3_MIMO_chunk_ref(q, k, v, qb, kb, mv, mo, z, mz, ang, dcs, dcsr,
                                         dt, trap, D, chunk_size=C, rotary_dim_divisor=2,
                                         dtype=torch.float32, rotate_pairwise=True)
    e = max_rel_err(got, ref)
    worst = max(worst, e)
    ok = e < 5e-5
    if not ok:
        fails.append(f"S{S}C{C}R{R}")
    print(f"  S={S:4d} H={H:2d} G={G} P={P:3d} N={N:3d} R={R} C={C:3d} CR={C*R:3d}"
          f"   {e:.3e}  {'ok' if ok else 'FAIL'}")

print(f"\nworst {worst:.3e}")
print("PASS" if not fails else "FAIL: " + ", ".join(fails))
sys.exit(0 if not fails else 1)