"""Stage 4 — evaluate NPC-Reason-SFT on the FROZEN eval, vs the frozen baseline + PREREG.

Same EVAL.lock, same greedy decoding (temp 0, seed 0, max_tokens 12288) as the baseline =
apples-to-apples. Scores with the frozen verifier. Grades the pre-registered bars honestly.
"""

import json
import sys
import time
from pathlib import Path

from transformers import AutoTokenizer
from vllm import LLM, SamplingParams

HERE = Path(__file__).resolve().parent
ROOT = HERE.parent
sys.path.insert(0, str(ROOT))
from verifier.step_verifier import verify_chain  # FROZEN

MERGED = str(ROOT / "sft" / "merged")
MAX_TOKENS = 12288
MAX_MODEL_LEN = 16384

PLAIN_INSTRUCTION = (
    "Solve this math problem. Show your work step by step, then give the final answer "
    "on the last line as \\boxed{{ANSWER}}.\n\nProblem: {problem}"
)
FORMAT_INSTRUCTION = (
    "Solve this math problem. For EVERY load-bearing arithmetic step, write the "
    "computation as an inline checkable assertion in the exact form <<EXPR = RESULT>>, "
    "where EXPR is the arithmetic expression and RESULT is its value (for example "
    "<<3*8 = 24>>). If a quantity is reused, you may name it, e.g. "
    "let total = <<3*8 = 24>>, and reference it later as <<total + 6 = 30>>. Do not "
    "assert any number that drives the answer without wrapping it in <<...>>. End with "
    "the final answer as \\boxed{{ANSWER}}, and make sure it equals the result of your "
    "last <<...>> step.\n\nProblem: {problem}"
)


def load_eval():
    return [json.loads(l) for l in (ROOT / "eval" / "eval_set.jsonl").open()]


def rate(recs, k):
    return round(100 * sum(1 for r in recs if r[k]) / len(recs), 2) if recs else 0.0


def main():
    rows = load_eval()
    tok = AutoTokenizer.from_pretrained(MERGED)
    variants = {"plain": PLAIN_INSTRUCTION, "format": FORMAT_INSTRUCTION}
    jobs = []
    for v, tmpl in variants.items():
        for r in rows:
            p = tok.apply_chat_template([{"role": "user", "content": tmpl.format(problem=r["problem"])}],
                                        tokenize=False, add_generation_prompt=True)
            jobs.append((v, r, p))

    llm = LLM(model=MERGED, dtype="bfloat16", max_model_len=MAX_MODEL_LEN,
              gpu_memory_utilization=0.85, seed=0)
    sp = SamplingParams(temperature=0.0, max_tokens=MAX_TOKENS, seed=0)
    t0 = time.time()
    outs = llm.generate([j[2] for j in jobs], sp)
    gen_s = round(time.time() - t0, 1)

    per = {"plain": [], "format": []}
    raw = HERE / "eval_raw.jsonl"
    with raw.open("w") as f:
        for (v, r, _), o in zip(jobs, outs):
            text = o.outputs[0].text
            rec = verify_chain(text, gold_answer=r["gold"])
            slim = {"verifiable": rec["verifiable"], "correct": bool(rec["correct"]),
                    "verified_and_correct": bool(rec["verified_and_correct"]),
                    "n_assertions": rec["n_assertions"]}
            per[v].append(slim)
            f.write(json.dumps({"id": r["id"], "variant": v, "gold": r["gold"],
                                "output": text, "verifier": slim}, ensure_ascii=False) + "\n")

    summary = {v: {"n": len(rs), "verifiable_rate": rate(rs, "verifiable"),
                   "accuracy": rate(rs, "correct"),
                   "verified_and_correct": rate(rs, "verified_and_correct"),
                   "mean_assertions": round(sum(x["n_assertions"] for x in rs) / len(rs), 2)}
               for v, rs in per.items()}

    base = json.load(open(ROOT / "baseline" / "BASELINE.lock"))["summary"]
    bf = base["format"]
    sf = summary["format"]

    # PREREG bars (b5a49437): +15pp v&c over format baseline; verifiable >=90%; acc regress <=5pp.
    lift = round(sf["verified_and_correct"] - bf["verified_and_correct"], 2)
    acc_delta = round(sf["accuracy"] - bf["accuracy"], 2)
    grade = {
        "primary_lift_vc_pp": lift,
        "primary_SUPPORTED": lift >= 15.0,
        "verifiable_rate": sf["verifiable_rate"],
        "verifiable_SUPPORTED": sf["verifiable_rate"] >= 90.0,
        "accuracy_delta_pp": acc_delta,
        "accuracy_guard_SUPPORTED": acc_delta >= -5.0,
    }
    grade["ALL_THREE"] = all([grade["primary_SUPPORTED"], grade["verifiable_SUPPORTED"],
                              grade["accuracy_guard_SUPPORTED"]])
    rl_gate = grade["ALL_THREE"] or sf["verifiable_rate"] >= 90.0

    result = {"model": MERGED, "decoding": {"greedy": True, "max_tokens": MAX_TOKENS, "seed": 0},
              "gen_seconds": gen_s, "summary": summary, "baseline_format": bf,
              "prereg_grade": grade, "rl_gate_cleared": rl_gate}
    (HERE / "eval_result.json").write_text(json.dumps(result, indent=2))
    print(json.dumps({"summary": summary, "grade": grade, "rl_gate": rl_gate}, indent=2))


if __name__ == "__main__":
    main()