"""Run the pinned MTPLX final-state exit matrix without loading model weights."""
import importlib.util
import json
import os
from pathlib import Path
import subprocess
import sys

root = Path(sys.argv[1]).resolve()
assert subprocess.check_output(["git", "-C", str(root), "rev-parse", "HEAD"], text=True).strip() == "9882703f3105363ddc37eca9f97aa09a1d387112"
assert os.environ["MTPLX_CONTEXT_COPY"] == "0"
sys.path.insert(0, str(root))
spec = importlib.util.spec_from_file_location("final_state", root / "tests/test_mtp_final_state_covers_committed_tokens.py")
m = importlib.util.module_from_spec(spec)
spec.loader.exec_module(m)
for label, stop_index, depth, omniscient, temperature, maximum in [
    ("bonus primary", 8, 2, False, 0.0, 16),
    ("deferred greedy correction", 8, 3, False, 0.0, 16),
    ("accepted draft stop", 8, 3, True, 0.0, 16),
    ("sampled correction", 8, 3, False, 0.7, 16),
    ("first primary", 6, 3, False, 0.0, 16),
    ("fresh primary depth 1", 7, 1, False, 0.0, 16),
    ("max tokens depth 3", 10000, 3, False, 0.0, 7),
    ("max tokens depth 2", 10000, 2, False, 0.0, 5),
]:
    sampler = m.SamplerConfig(temperature=temperature, top_k=1 if temperature else 0)
    out = m.generate_mtpk(m._runtime(m._ScriptedModel(stop_index, omniscient)), m.PROMPT,
                         speculative_depth=depth, max_tokens=maximum, sampler=sampler, **m.COMMON)
    gap, origin, safe = m._gap(m.PROMPT, out)
    assert gap == 0 and safe
    assert out.final_state.final_committed_mtp_cache[0].offset == len(m.PROMPT) + len(out.tokens) - 1
    if stop_index < 1000:
        assert out.tokens[-1] == m.STOP
    print(json.dumps(dict(case="mtp_exit", label=label, stop_index=stop_index, depth=depth,
                          omniscient=omniscient, temperature=temperature, max_tokens=maximum,
                          tokens=out.tokens, gap=gap, origin=origin, safe=safe)), flush=True)
m.test_ar_final_state_control()
print(json.dumps(dict(case="ar_control", passed=True)), flush=True)