#!/usr/bin/env python3
# SPDX-License-Identifier: AGPL-3.0-or-later
# Copyright (C) 2026 David (Corey) Herrera
"""
gen_wkv7_fixture.py -- Emit a C header of fixed-seed WKV7 fixtures for the
gradient-check test (test_wkv7_bwd.c).

It imports a *verified* WKV7 reference oracle (wkv7_reference.py, whose manual
reverse-scan backward matches autograd to rel<1e-9 in fp64; it is not shipped
here) and dumps, for a few fixed-seed shapes: the inputs (r,w,k,v,a,b), the
upstream grad dy, the forward y, and the reference grads (dr,dw,dk,dv,da,db).

Run:  python3 rwkvx/tests/gen_wkv7_fixture.py /path/to/wkv7_reference.py \
          > rwkvx/tests/wkv7_fixture.h
      (or set WKV7_REFERENCE=/path/to/wkv7_reference.py instead of the argument)
(torch CPU is fine).  The generated header is committed so the C test runs
without torch; regenerate only if the oracle changes.
"""
import os
import sys
import importlib.util

ORACLE = sys.argv[1] if len(sys.argv) > 1 else os.environ.get("WKV7_REFERENCE")
if not ORACLE or not os.path.isfile(ORACLE):
    sys.exit("usage: gen_wkv7_fixture.py <path/to/wkv7_reference.py> "
             "(or set WKV7_REFERENCE)")

spec = importlib.util.spec_from_file_location("wkv7_reference", ORACLE)
oracle = importlib.util.module_from_spec(spec)
spec.loader.exec_module(oracle)

import torch  # noqa: E402


def emit_array(name, t):
    """Emit a flat C float array from a tensor (row-major [B,T,H,N])."""
    flat = t.detach().to(torch.float64).reshape(-1).tolist()
    print(f"static const float {name}[] = {{")
    line = "    "
    for i, x in enumerate(flat):
        tok = f"{x:.9e}f,"
        if len(line) + len(tok) > 96:
            print(line)
            line = "    "
        line += tok + " "
    if line.strip():
        print(line)
    print("};")


def emit_case(idx, B, T, H, N, seed, dy_mode):
    ins = oracle._rand_inputs(B, T, H, N, torch.float64, seed=seed,
                              requires_grad=True)
    r, w, k, v, a, b = ins
    y = oracle.wkv7_forward(r, w, k, v, a, b)
    if dy_mode == "ones":
        dy = torch.ones_like(y)
    else:
        g = torch.Generator().manual_seed(seed + 1000)
        dy = torch.randn(*y.shape, generator=g, dtype=torch.float64)
    y.backward(dy)
    grads = {"dr": r.grad, "dw": w.grad, "dk": k.grad,
             "dv": v.grad, "da": a.grad, "db": b.grad}

    p = f"c{idx}_"
    print(f"\n/* ---- case {idx}: B={B} T={T} H={H} N={N} seed={seed} "
          f"dy={dy_mode} ---- */")
    print(f"#define {p.upper()}B {B}")
    print(f"#define {p.upper()}T {T}")
    print(f"#define {p.upper()}H {H}")
    print(f"#define {p.upper()}N {N}")
    for nm, t in [("r", r), ("w", w), ("k", k), ("v", v), ("a", a), ("b", b),
                  ("dy", dy), ("y", y)]:
        emit_array(p + nm, t)
    for nm, t in grads.items():
        emit_array(p + nm, t)


def main():
    print("/* AUTO-GENERATED by rwkvx/tests/gen_wkv7_fixture.py -- DO NOT EDIT.")
    print(" * Source oracle: wkv7_reference.py (fp64 PyTorch WKV7 reference)")
    print(" * (verified: manual backward == autograd, rel<1e-9 fp64).")
    print(" * Reference grads are fp64, stored as float. */")
    print("#ifndef WKV7_FIXTURE_H")
    print("#define WKV7_FIXTURE_H")
    # canonical fixture from the doc's gradient-check plan
    emit_case(0, 1, 4, 1, 4, 42, "ones")
    # larger, multi-head, random dy -> exercises the reverse scan harder
    emit_case(1, 1, 16, 2, 16, 7, "randn")
    # a third with B>1 and different N
    emit_case(2, 2, 8, 3, 8, 123, "randn")
    print("\n#endif /* WKV7_FIXTURE_H */")


if __name__ == "__main__":
    main()