"""H4 cheap fusion: one gather_qmm for up+gate must be bit-identical to two (row-independent dot products)."""
import json, os, time, mlx.core as mx, pytest
from pathlib import Path
from mlx_lm.utils import load_model, load_tokenizer
from titan.runtime import TitanModel

def test_fused_upgate_bit_identical(model_dir):
    refm, _ = load_model(Path(model_dir)); tok = load_tokenizer(Path(model_dir))
    from titan.registry import load_spec
    E = load_spec(model_dir).num_experts
    rec = {"model": model_dir, "ts": time.strftime("%Y-%m-%dT%H:%M:%S"), "cases": []}
    for cap in (E, max(8, E // 4)):
        tm = TitanModel(model_dir, capacity=cap, fused_upgate=True)
        if cap == E: tm.warm_all()
        for p in ["The capital of Canada is", "def fibonacci(n):\n    \"\"\"Return the n-th Fibonacci number.\"\"\"\n"]:
            ids = mx.array(tok.encode(p))[None]
            ref = refm(ids); out = tm(ids); mx.eval(ref, out)
            d = float(mx.max(mx.abs(out.astype(mx.float32) - ref.astype(mx.float32))))
            rec["cases"].append({"capacity": cap, "prompt": p[:30], "max_abs_dlogit": d})
            assert d <= 1e-3, (cap, p, d)
    os.makedirs("benchmarks/raw/tests", exist_ok=True)
    json.dump(rec, open(f"benchmarks/raw/tests/fused_upgate_{tm.spec.model_type}_{rec['ts'].replace(':','')}.json", "w"), indent=2)