#!/usr/bin/env python
"""Prove the engine's cached decode path equals a full recompute.

This is the spike_infer check. verify.py is the same idea against generate.py
/ model_v2. Both ship; they test different stacks.

Greedy is deterministic, so the two paths must produce the same token ids,
not similar text. This is the test that caught MemoryCacheBranch rebuilding
from a one-token window: max |dlogit| was 7.24e-01 as shipped, vs ~1e-05
float32 floor with the branch off.

Three checks:

  1. token identity  -- cached greedy ids == full-recompute greedy ids
  2. logit agreement -- max |dlogit| at each step, expect the f32 floor (~1e-5)
  3. slot count      -- KV cache has len(loop_plan) entries, one per
                        (loop pass, layer), plus the trailing Engram slot

CPU by default so it does not steal a GPU from a training job:

    python verify_cache.py                 # cpu, float32, package.json ckpt
    python verify_cache.py --device cuda
    python verify_cache.py --ckpt checkpoints/sft_7100.pt
"""
import argparse
import os
import sys

import torch

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)

from spike_infer import SpikeEngine   # noqa: E402


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--device", default="cpu")
    ap.add_argument("--steps", type=int, default=16)
    ap.add_argument("--prompt", default="The capital of France is")
    ap.add_argument("--tol", type=float, default=1e-3)
    ap.add_argument("--model-dir", default=None,
                    help="alternate exported weights dir (safetensors)")
    ap.add_argument("--ckpt", default=None,
                    help="released stage: checkpoints/base_62k.pt, "
                         "checkpoints/sft_7100.pt, "
                         "checkpoints/dpo_3200.pt (default via package.json)")
    args = ap.parse_args()

    eng = SpikeEngine(HERE, device=args.device, dtype=torch.float32,
                      fast_cache=True, verbose=True, model_dir=args.model_dir,
                      ckpt_file=args.ckpt)
    model, tok = eng.model, eng.tok
    dev = eng.device
    inner = getattr(model, "model", model)

    n_slots = int(inner._n_cache_slots)
    plan_ok = n_slots == len(inner.loop_plan)
    print(f"\n[slots] loop_count={inner.loop_count} layers={eng.cfg.num_hidden_layers} "
          f"-> {n_slots} KV slots  {'OK' if plan_ok else 'MISMATCH'}")

    ids = tok.encode(args.prompt)
    ids = ids.tolist() if hasattr(ids, "tolist") else list(ids)
    bos = getattr(tok, "bos_token_id", None)
    if bos is not None and (not ids or ids[0] != bos):
        ids = [int(bos)] + ids

    # ---- path A: cached, incremental (the engine's own decode path) --------
    eng._reset_runtime()
    eng._engram_kw_ctx = {}
    cached_ids, cached_logits = [], []
    logits, _ = eng._forward(ids)
    for _ in range(args.steps):
        nxt = int(logits[0].argmax())
        cached_ids.append(nxt)
        cached_logits.append(logits[0].detach().float().clone())
        logits, _ = eng._forward([nxt])

    # ---- path B: no cache, full recompute of the whole prefix every step ---
    eng.mods["model_v2"].reset_memory_cache(model)
    full_ids, full_logits = [], []
    seq = list(ids)
    with torch.no_grad():
        for _ in range(args.steps):
            t = torch.tensor([seq], device=dev)
            pos = torch.arange(len(seq), device=dev).unsqueeze(0)
            out = model(t, position_ids=pos, past_key_values=None, use_cache=False)
            lg = out.logits[0, -1]
            nxt = int(lg.argmax())
            full_ids.append(nxt)
            full_logits.append(lg.detach().float().clone())
            seq.append(nxt)

    # ---- compare -----------------------------------------------------------
    same = cached_ids == full_ids
    n_match = sum(a == b for a, b in zip(cached_ids, full_ids))
    diffs = [float((a - b).abs().max()) for a, b in zip(cached_logits, full_logits)]
    worst = max(diffs) if diffs else 0.0

    print(f"\n[tokens] cached : {cached_ids}")
    print(f"[tokens] full   : {full_ids}")
    print(f"[tokens] identical {n_match}/{args.steps}")
    print(f"[logits] max |delta| = {worst:.3e}  (first step {diffs[0]:.3e})")
    print(f"\ncached : {eng._visible(cached_ids)!r}")
    print(f"full   : {eng._visible(full_ids)!r}")

    ok = same and worst < args.tol and plan_ok
    print(f"\n{'PASS' if ok else 'FAIL'} -- cached decoding "
          f"{'matches' if same else 'DIVERGES FROM'} full recompute")
    return 0 if ok else 1


if __name__ == "__main__":
    raise SystemExit(main())