"""Numerical correctness of the TITAN execution path against mlx-lm on identical weights.

Tolerances (explicit): with all experts resident or with a bounded slot table, the kernel, weights and
routing are identical to mlx-lm, so we require bit-identical expert choices and |Δlogit| <= 1e-3
(fp16 accumulation order inside gather_qmm may differ with slot ordering). The actual max |Δ| is
written to benchmarks/raw/tests/numerical_<ts>.json as evidence.
"""
import json, os, time
import mlx.core as mx
import pytest
from mlx_lm.utils import load_model, load_tokenizer
from pathlib import Path
from titan.runtime import TitanModel

PROMPTS = [
    "The capital of Canada is",
    "def fibonacci(n):\n    \"\"\"Return the n-th Fibonacci number.\"\"\"\n",
    "Q: If a train travels 60 km in 45 minutes, what is its average speed in km/h?\nA:",
    "Les modèles à mélange d'experts activent seulement une fraction de leurs paramètres parce que",
]
TOL = 1e-3

def _route_trace(tm):
    """Hook capturing the unique expert ids chosen per layer, in call order."""
    trace = []
    def on_route(layer, uniq):
        trace.append((layer, tuple(uniq)))
    for b in tm.blocks:
        b.on_route = on_route
    return trace

def _ref_logits_and_routes(ref, tok, prompt):
    """mlx-lm reference logits + router choices (recomputed from the same router ops)."""
    ids = mx.array(tok.encode(prompt))[None]
    routes = []
    orig = [l.mlp.switch_mlp.__call__ for l in ref.model.layers]
    for li, l in enumerate(ref.model.layers):
        sw = l.mlp.switch_mlp
        def wrapped(x, indices, _sw=sw, _li=li):
            routes.append((_li, tuple(dict.fromkeys(indices.reshape(-1).tolist()))))
            return type(_sw).__call__(_sw, x, indices)
        l.mlp.switch_mlp.__call__ = wrapped  # instance attribute shadowing is not used by nn.Module; call directly below
    logits = ref(ids)
    mx.eval(logits)
    return ids, logits, routes

@pytest.fixture(scope="module")
def ref(model_dir):
    m, _ = load_model(Path(model_dir))
    tok = load_tokenizer(Path(model_dir))
    return m, tok

@pytest.mark.parametrize("capacity", [64, 32, 16, 8])
def test_logits_match_mlx_lm(model_dir, ref, capacity):
    refm, tok = ref
    tm = TitanModel(model_dir, capacity=capacity)
    trace = _route_trace(tm)
    rec = {"capacity": capacity, "prompts": [], "tol": TOL, "ts": time.strftime("%Y-%m-%dT%H:%M:%S")}
    worst = 0.0
    for p in PROMPTS:
        ids = mx.array(tok.encode(p))[None]
        ref_logits = refm(ids); mx.eval(ref_logits)
        trace.clear(); tm.reset_stats()
        tl = tm(ids); mx.eval(tl)
        d = float(mx.max(mx.abs(tl.astype(mx.float32) - ref_logits.astype(mx.float32))))
        same_argmax = bool(mx.array_equal(mx.argmax(tl, -1), mx.argmax(ref_logits, -1)))
        # reference expert choices: recompute the router exactly as mlx-lm does, from the reference model's own hidden states
        worst = max(worst, d)
        st = tm.cache_stats()
        rec["prompts"].append({"prompt": p[:40], "tokens": int(ids.shape[1]), "max_abs_dlogit": d, "argmax_equal": same_argmax,
                               "hit_rate": st["hit_rate"], "misses": st["misses"], "stall_s": st["stall_time_s"],
                               "layers_routed": len(trace)})
        assert same_argmax, (capacity, p)
        assert d <= TOL, (capacity, p, d)
        assert len({l for l, _ in trace}) == tm.spec.num_layers
    rec["max_abs_dlogit_worst"] = worst
    os.makedirs("benchmarks/raw/tests", exist_ok=True)
    with open(f"benchmarks/raw/tests/numerical_cap{capacity}_{rec['ts'].replace(':','')}.json", "w") as f:
        json.dump(rec, f, indent=2)

def test_greedy_decode_identical(model_dir, ref):
    """Greedy continuation (32 tokens) must be token-identical to mlx-lm with a bounded (16-slot) table."""
    from mlx_lm import generate
    refm, tok = ref
    tm = TitanModel(model_dir, capacity=16)
    outs = []
    for p in PROMPTS[:2]:
        a = generate(refm, tok, prompt=p, max_tokens=32, verbose=False)
        b = generate(tm.model, tok, prompt=p, max_tokens=32, verbose=False)
        outs.append({"prompt": p[:40], "ref": a, "titan": b, "equal": a == b})
        assert a == b, (a, b)
    with open(f"benchmarks/raw/tests/greedy_cap16_{time.strftime('%Y%m%dT%H%M%S')}.json", "w") as f:
        json.dump(outs, f, indent=2)