"""Check the independent reference against the upstream Mamba-3 reference."""

import argparse
import sys

import torch

sys.path.insert(0, str(__import__("pathlib").Path(__file__).parent))
from anchor import build_case, load_upstream_refs, max_rel_err  # noqa: E402
from mamba3_ref import mimo_step_ref  # noqa: E402

GRID = [
    #  S    H  G   P    N   R
    (128,  8, 1, 32,  64, 4),
    (128,  8, 1, 32,  64, 1),
    (256,  4, 1, 64, 128, 4),
    (256,  8, 2, 32,  64, 2),
    (64,   8, 1, 64, 128, 8),
]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--mamba-tests", required=True)
    ap.add_argument("--device", default="cpu")
    args = ap.parse_args()
    up = load_upstream_refs(args.mamba_tests)

    print(f"torch {torch.__version__}  device={args.device}")
    print("mimo_step_ref (ours) vs mamba3_MIMO_step_ref (upstream):")
    worst = 0.0
    for S, H, G, P, N, R in GRID:
        B, C = 1, 16
        c = build_case(B, S, H, G, P, N, R, C, args.device, torch.float32, seed=0)

        ours, _ = mimo_step_ref(
            c["q"], c["k"], c["v"], c["ADT"], c["dt"], c["trap"],
            c["q_bias"], c["k_bias"], c["angles"], c["mimo_v"], c["mimo_o"],
            D=c["D"], z=c["z"], mimo_z=c["mimo_z"],
        )
        ref, _ = up.mamba3_MIMO_step_ref(
            c["q"], c["k"], c["v"], c["ADT"], c["dt"], c["trap"],
            c["q_bias"], c["k_bias"], c["angles"], c["mimo_v"], c["mimo_o"],
            D=c["D"], Z=c["z"], MIMO_Z=c["mimo_z"],
        )
        rel = max_rel_err(ours, ref)
        worst = max(worst, rel)
        print(f"  S={S:4d} H={H:2d} G={G} P={P:3d} N={N:3d} R={R}  max_rel={rel:.3e}")

    # fused-norm variant
    c = build_case(1, 128, 8, 1, 32, 64, 4, 16, args.device, torch.float32, seed=3)
    w = torch.randn((8, 32), device=args.device, dtype=torch.float32)
    ours, _ = mimo_step_ref(
        c["q"], c["k"], c["v"], c["ADT"], c["dt"], c["trap"], c["q_bias"], c["k_bias"],
        c["angles"], c["mimo_v"], c["mimo_o"], D=c["D"], z=c["z"], mimo_z=c["mimo_z"],
        fused_norm=True, outproj_norm_weight=w,
    )
    ref, _ = up.mamba3_MIMO_step_ref(
        c["q"], c["k"], c["v"], c["ADT"], c["dt"], c["trap"], c["q_bias"], c["k_bias"],
        c["angles"], c["mimo_v"], c["mimo_o"], D=c["D"], Z=c["z"], MIMO_Z=c["mimo_z"],
        fused_norm=True, outproj_norm_weight=w,
    )
    rel = max_rel_err(ours, ref)
    worst = max(worst, rel)
    print(f"  fused_norm                              max_rel={rel:.3e}")

    print(f"\nworst: {worst:.3e}")
    ok = worst < 1e-5
    print("PASS" if ok else "FAIL")
    return 0 if ok else 1


if __name__ == "__main__":
    sys.exit(main())