"""02_forward_parity.py -- verify NKI kernel forward matches an eager reference.

Runs a Mamba-3 forward on the eager path and on the NKI-accelerated path with
identical weights, then compares the outputs. Both paths should agree at
cos_sim > 0.999999 (numerical noise from bf16-adjacent MLIR lowering).

Runs both SISO and MIMO modes.

Run:
  python 02_forward_parity.py
"""

import torch
from _loader import load_kernel


def compare(name, y_nki, y_ref, tol_cos=0.999):
    y_nki_f = y_nki.cpu().float().reshape(-1)
    y_ref_f = y_ref.cpu().float().reshape(-1)
    max_abs = (y_nki_f - y_ref_f).abs().max().item()
    ref_norm = y_ref_f.norm().item()
    nki_norm = y_nki_f.norm().item()
    if ref_norm < 1e-12 and nki_norm < 1e-12:
        cos = 1.0
    else:
        cos = torch.dot(y_nki_f, y_ref_f).item() / (ref_norm * nki_norm + 1e-12)
    status = "PASS" if cos > tol_cos else "FAIL"
    print(f"  [{status}] {name}: cos_sim={cos:.6f}  max_abs={max_abs:.3e}  ||ref||={ref_norm:.3e}")
    return cos > tol_cos


def parity_one(mamba3, mode: str, seqlen: int):
    print(f"\n--- {mode.upper()} parity, seqlen={seqlen} ---")

    if mode == "siso":
        kwargs_ref = dict(d_model=1024, d_state=128, headdim=64, chunk_size=64, mimo_rank=1,
                          use_nki_ssd=False)
        kwargs_nki = dict(d_model=1024, d_state=128, headdim=64, chunk_size=64, mimo_rank=1,
                          use_nki_ssd=True)
    else:  # mimo
        kwargs_ref = dict(d_model=1024, d_state=128, headdim=64, chunk_size=16, mimo_rank=4,
                          use_nki_ssd=False)
        kwargs_nki = dict(d_model=1024, d_state=128, headdim=64, chunk_size=16, mimo_rank=4,
                          use_nki_ssd=True)

    torch.manual_seed(0)
    mixer_ref = mamba3.NeuronMamba3Mixer(**kwargs_ref).to("neuron")
    torch.manual_seed(0)
    mixer_nki = mamba3.NeuronMamba3Mixer(**kwargs_nki).to("neuron")

    # Copy weights from ref to nki (both were seeded identically, but this makes it explicit)
    mixer_nki.load_state_dict(mixer_ref.state_dict())

    torch.manual_seed(100)
    u = torch.randn(1, seqlen, 1024, device="neuron")

    y_ref, cache_ref = mixer_ref(u)
    y_nki, cache_nki = mixer_nki(u)
    try:
        torch.neuron.synchronize()
    except Exception:
        pass

    ok_y = compare("y_out", y_nki, y_ref, tol_cos=0.999)
    ok_state = compare("ssm_state", cache_nki.ssm_state, cache_ref.ssm_state, tol_cos=0.999)
    return ok_y and ok_state


def main():
    print("=" * 60)
    print("Forward parity: NKI kernel vs eager reference")
    print("=" * 60)
    mamba3 = load_kernel()
    all_pass = True
    for mode in ["siso", "mimo"]:
        for seqlen in [64, 256]:
            all_pass = parity_one(mamba3, mode, seqlen) and all_pass
    print()
    print("=" * 60)
    print("ALL PARITY TESTS PASS" if all_pass else "SOME TESTS FAILED")
    print("=" * 60)


if __name__ == "__main__":
    main()