"""Exactly paired test of the forced code-check module on recorded runs.

For every recorded answer whose thinking closed, the same thinking is replayed (prompt now declares the python
tool), the model is made to write a check program right after </think>, the program runs, and the model continues
with the output in a new thinking turn and gives its final answer. Truncated records are copied unchanged, so the
difference to the source run is due to the check alone.

usage: posthoc_check.py --src results/m180_base/mini.jsonl --out results/m180_chk/mini.jsonl [--max-tokens 6144]
"""
import argparse
import json
import os
import sys
import threading
import time
import urllib.request
import zlib
from concurrent.futures import ThreadPoolExecutor

from transformers import AutoTokenizer

sys.argv, _argv = sys.argv[:1] + ["--inp", "x", "--out", "x"], sys.argv  # reuse tir.py's helpers without running it
_src = open(os.path.join(os.path.dirname(os.path.abspath(__file__)), "tir.py"), encoding="utf-8").read()
exec(_src[:_src.index("def solve(")])  # noqa: S102  (constants, TOOL, SYSTEMS, run_code, post, CALL_RE)
sys.argv = _argv

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from grade import last_boxed  # noqa: E402

ap2 = argparse.ArgumentParser()
ap2.add_argument("--src", required=True)
ap2.add_argument("--out", required=True)
ap2.add_argument("--items", nargs="+", default=["data/coc/math180.jsonl"])
ap2.add_argument("--max-tokens", type=int, default=6144)
ap2.add_argument("--system", default="verify")
ap2.add_argument("--check-line", default="# Independent check of the answer: compute it directly / by brute force\n")
ap2.add_argument("--workers", type=int, default=32)
ap2.add_argument("--mode", default="check", choices=["check", "switch"],
                 help="check: code check after the thinking closes; switch: if the thinking runs longer than "
                      "--switch-at tokens, stop it there and make the model compute the answer with a program")
ap2.add_argument("--switch-at", type=int, default=3072)
ap2.add_argument("--more-calls", type=int, default=0, help="further tool calls the model may make after the forced one")
ap2.add_argument("--switch-text", default="\n\nThis is taking long. Let me compute the answer directly with a program.\n")
a2 = ap2.parse_args()
ENDS = {i for i in range(len(tok)) if (t := tok.convert_ids_to_tokens(i)) is not None and (t.endswith(".") or t.endswith("Ċ"))}
items = {}
for p in a2.items:
    for l in open(p, encoding="utf-8"):
        r = json.loads(l)
        items[r["id"]] = r
SAMP = {"temperature": 0.6, "top_p": 0.95, "top_k": 50, "min_p": 0.0, "samplers": ["top_k", "top_p", "temperature"]}


def check(rec):
    try:
        return _check(rec)
    except Exception as e:
        return {**rec, "check_error": repr(e)[:200]}


def _check(rec):
    c = rec["completion"]
    it = items[rec["id"]]
    think = c[len("<think>"):].split("</think>", 1)[0]
    if a2.mode == "check" and (rec["truncated"] or "</think>" not in c):
        return rec
    th = tok(think, add_special_tokens=False)["input_ids"]
    if a2.mode == "switch":
        if len(th) <= a2.switch_at:
            return rec
        cut = next((j for j in range(a2.switch_at, a2.switch_at // 2, -1) if th[j - 1] in ENDS), a2.switch_at)
        think = tok.decode(th[:cut]) + a2.switch_text
        th = tok(think, add_special_tokens=False)["input_ids"]
    th = th + [THINK_END]
    msgs = ([{"role": "system", "content": SYSTEMS[a2.system]}] if SYSTEMS[a2.system] else []) + \
        [{"role": "user", "content": it["prompt"]}]
    ids = tok.apply_chat_template(msgs, tools=[TOOL], tokenize=True, add_generation_prompt=True, return_dict=False)
    used = len(th)
    seed = zlib.crc32(rec["id"].encode()) % 1000003 * 7 + rec["sample"] + 7
    pre = tok(FORCE_PREFIX + a2.check_line, add_special_tokens=False)["input_ids"]
    r2 = post({"prompt": ids + th + pre, "n_predict": max(64, min(1536, a2.max_tokens - used)), "seed": seed,
               "cache_prompt": False, "stop": ["</parameter>"], **SAMP})
    used += r2["tokens_predicted"]
    code = a2.check_line + r2["content"]
    out = run_code(code)
    ids = ids + th + pre + tok(r2["content"].rstrip("\n") + "\n</parameter>\n</function>\n</tool_call><|im_end|>\n"
                               "<|im_start|>tool\n<tool_response>\n" + out + "\n</tool_response><|im_end|>\n"
                               "<|im_start|>assistant\n<think>", add_special_tokens=False)["input_ids"]
    log = ""
    for extra in range(a2.more_calls + 1):  # the model may run more code (e.g. after an error) before answering
        r3 = post({"prompt": ids, "n_predict": max(256, a2.max_tokens - used), "seed": seed + 1 + extra,
                   "cache_prompt": False, "return_tokens": True, **SAMP})
        used += r3["tokens_predicted"]
        after = r3["content"]
        m = CALL_RE.search(after.split("</think>", 1)[1]) if "</think>" in after else None
        if not m or extra == a2.more_calls or used >= a2.max_tokens:
            break
        out2 = run_code(m.group(1))
        turn = after.split("</tool_call>", 1)[0] + "</tool_call>"
        log += turn.replace("</think>", "\n", 1) + "\n<tool_response>\n" + out2 + "\n</tool_response>\n"
        ids = ids + [t for t in r3["tokens"] if t != EOS] + tok(
            "<|im_end|>\n<|im_start|>tool\n<tool_response>\n" + out2 + "\n</tool_response><|im_end|>\n"
            "<|im_start|>assistant\n<think>", add_special_tokens=False)["input_ids"]
    after = log + after
    if "</think>" not in after:  # the post-check turn ran out: keep the original answer (switch: no answer)
        if a2.mode == "switch":
            return {**rec, "completion": "<think>" + think + after, "n_tokens": used, "truncated": True,
                    "check": {"code": code, "out": out, "kept": "none"}}
        return {**rec, "check": {"code": code, "out": out, "kept": "original"}}
    t2, final = after.split("</think>", 1)
    if "\\boxed" in it["prompt"] and last_boxed(final) is None:
        # after a tool turn the model often answers in prose: read the answer out in the requested form
        last = [t for t in r3["tokens"] if t != EOS]
        k = last.index(THINK_END) + 1 if THINK_END in last else len(last)
        pr = post({"prompt": ids + last[:k] + tok("The final answer is $\\boxed{", add_special_tokens=False)["input_ids"],
                   "n_predict": 48, "temperature": 0.0, "cache_prompt": False, "stop": ["}$", "}\n"]})
        final = final.rstrip() + "\n\nThe final answer is $\\boxed{" + pr["content"] + "}$"
    return {**rec, "completion": "<think>" + think + "\n" + FORCE_PREFIX + code + "\n</parameter>\n</function>\n"
            "</tool_call>\n<tool_response>\n" + out + "\n</tool_response>\n" + t2 + "</think>" + final,
            "n_tokens": used, "truncated": False, "check": {"code": code, "out": out, "kept": "new"}}


recs = [json.loads(l) for l in open(a2.src, encoding="utf-8")]
os.makedirs(os.path.dirname(a2.out), exist_ok=True)
t0 = time.time()
with ThreadPoolExecutor(a2.workers) as ex:
    res = list(ex.map(check, recs))
with open(a2.out, "w", encoding="utf-8") as f:
    for r in res:
        f.write(json.dumps(r, ensure_ascii=False) + "\n")
print(f"{len(res)} records, {sum('check' in r for r in res)} checked, {(time.time() - t0) / 60:.1f} min")