#!/usr/bin/env python
"""Lab 4 (also used in Labs 2 and 5): how did each answer end? Reads the full answer texts of finished runs.

    . /nvme/wsp-lab/code/scripts/env.sh && cd /nvme/wsp-lab/code
    export CUDA_VISIBLE_DEVICES=""                                # CPU only, a few seconds per run
    # sampler VM
    python labs/lab4_endings.py                                   # Lab 4, the alpha sweep -> results/lab4/endings.json
    python labs/lab4_endings.py dev/base dev/p50c dev/l15only dev/l19only dev/all1 dev/all2 dev/randdir dev/randcols \
        --out lab4/endings-controls                               # Lab 4, which parts matter
    python labs/lab4_endings.py dev/base dev/base+temp1.0 dev/base+pp0 dev/base+pp2 dev/base+rep1.1 dev/base+max4k \
        dev/base+nothink dev/p50c+nothink --out lab5/endings      # Lab 5 (the two thinking-off runs are skipped)
    python labs/lab4_endings.py dev/base dev/p50c dev/rebuilt --out lab3/endings                     # Lab 3
    python labs/lab4_endings.py dev/base dev/p50c dev/capA-a6 dev/capB-a6 --out lab8/endings          # Lab 8
    # trainer VM
    python labs/lab4_endings.py bench/base/math500 bench/p50c/math500 bench/base/gsm8k bench/p50c/gsm8k \
        bench/base/ifeval bench/p50c/ifeval --out lab2/endings    # Lab 2
    python labs/lab4_endings.py dev/rebuilt-pairs2 dev/rebuilt-windows36-36 dev/rebuilt-layers15-19-21 \
        --out lab3/endings-tries                                  # Lab 3's experiments
    # either VM
    python labs/lab4_endings.py dev/a-4 --show 2                  # also print where 2 answers closed their thinking

Without --out, any run list other than the sweep is written to results/lab4/try-<names>/endings.json, so the main
results stay as they are.

It runs on the VM that made the runs, on the CPU, because the answer texts stay there (scripts/pull_runs.sh copies
only the per-answer numbers to the laptop). Every answer lands in exactly one of four groups:
  cut while thinking       it hit the 32,768-token cap before it ever wrote </think>
  stopped while thinking   the model emitted its end token before it ever wrote </think>, so there is no final answer
  closed out of a loop     it wrote </think> right after a stretch of verbatim repetition
  closed normally          it wrote </think> after ordinary reasoning
"Closed out of a loop" is a rough text check: the 60 characters that end 100 characters before </think> appear
at least 3 times in the 6,000 characters before </think>. An answer that closed its thinking can still be cut
later, while it writes the final answer; it stays in its "closed" group.
"""
import argparse
import glob
import hashlib
import json
import os
import sys

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
from wsp import hf, loops, protocol  # noqa: E402

SWEEP = ["dev/base", "dev/a-4", "dev/a1", "dev/a2", "dev/a4", "dev/p50c", "dev/a8", "dev/a12"]
GROUPS = ["cut while thinking", "stopped while thinking", "closed out of a loop", "closed normally"]
SNIP, GAP, LOOK, MIN_COUNT = 60, 100, 6000, 3

ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument("runs", nargs="*", default=SWEEP, help="run names under runs/ (default: the alpha sweep)")
ap.add_argument("--show", type=int, default=0, metavar="N",
                help="print the text around </think> for N correct answers that closed out of a loop, per run")
ap.add_argument("--out", metavar="PATH",
                help="write results/PATH.json instead (e.g. lab5/endings)")
a = ap.parse_args()


def closed_out_of_loop(text, i):
    """i = where </think> starts. True if the text just before it is a verbatim loop (rough check)."""
    if i < LOOK:
        return False
    snip = text[i - GAP - SNIP:i - GAP]
    return text[i - LOOK:i].count(snip) >= MIN_COUNT


def group(sample, text):
    i = text.find("</think>")
    if i < 0:
        return GROUPS[0] if sample["cut"] else GROUPS[1]
    return GROUPS[2] if closed_out_of_loop(text, i) else GROUPS[3]


res = {}
for run in a.runs:
    run_dir = os.path.join(hf.LAB, "runs", run)
    done = sorted(glob.glob(os.path.join(run_dir, "shard*.done.json")))
    if not done:
        sys.exit(f"no finished shards in {run_dir}; is this the VM that ran {run}?")
    if not protocol.PRESETS[json.load(open(done[0]))["preset"]]["thinking"]:
        print(f"skipping {run}: thinking is off, so its answers have no </think> to look for")
        continue
    counts = {g: {"answers": 0, "correct": 0} for g in GROUPS}
    shown = 0
    for rec, j, s, text in loops.iter_samples(run_dir):
        g = group(s, text)
        counts[g]["answers"] += 1
        counts[g]["correct"] += bool(s["correct"])
        if g == GROUPS[2] and s["correct"] and shown < a.show:
            shown += 1
            i = text.find("</think>")
            print(f"--- {run}: {rec['id']} answer {j}, {s['tokens']:,} tokens, closed out of a loop and correct")
            print(text[i - 600:i + 300].rstrip() + "\n")
    n = sum(c["answers"] for c in counts.values())
    if n == 0:
        sys.exit(f"no answer texts found in {run_dir}; is this the VM that ran {run}?")
    res[run] = {"answers": n, "correct": sum(c["correct"] for c in counts.values()), "groups": counts}

w = max(len(r) for r in res)
print(f"{'run':{w}s} answers correct | " + " | ".join(GROUPS) + "   (each: answers, of which correct)")
for run, r in res.items():
    cells = [f"{r['groups'][g]['answers']:>{len(g) - 6}d} ({r['groups'][g]['correct']:3d})" for g in GROUPS]
    print(f"{run:{w}s} {r['answers']:7d} {r['correct']:7d} | " + " | ".join(cells))

if not res:
    sys.exit("nothing to report")
if a.out:
    path = os.path.join(ROOT, "results", a.out.removesuffix(".json") + ".json")
elif a.runs != SWEEP:
    name = "-".join(r.replace("/", "_") for r in a.runs)
    if len(name) > 60:
        name = name[:50] + "-" + hashlib.md5(name.encode()).hexdigest()[:6]
    path = os.path.join(ROOT, "results", "lab4", "try-" + name, "endings.json")
else:
    path = os.path.join(ROOT, "results", "lab4", "endings.json")
os.makedirs(os.path.dirname(path), exist_ok=True)
json.dump({"check": {"snippet": SNIP, "gap": GAP, "look_back": LOOK, "min_count": MIN_COUNT}, "runs": res},
          open(path, "w"), indent=1)
print(f"\nwrote {os.path.relpath(path, ROOT)}")