#!/usr/bin/env python
"""Read the model's answers yourself: how each one ended, and the text where it looped or stopped.

    . /nvme/wsp-lab/code/scripts/env.sh && cd /nvme/wsp-lab/code
    export CUDA_VISIBLE_DEVICES=""                                  # CPU only, a few seconds
    python labs/peek.py quick/base                                  # one line per answer
    python labs/peek.py quick/base --id gsm8k_train_5628            # the end of answer 0 of that problem
    python labs/peek.py dev/p50c --id gsm8k_train_5628 --answer 1 --chars 3000
    python labs/peek.py build/base --id math_train_algebra_1625 --answer 1 --tokens 2893 2929   # one window

Run it on the VM that made the run: the answer texts stay there (scripts/pull_runs.sh copies only the per-answer
numbers to the laptop). For an answer that ends in a verbatim loop, it also finds the loop in tokens: where the
repetition starts, how many tokens one cycle has, and the text of one cycle. --tokens prints any stretch of the
answer by token position, counted from the start of the answer as in Lab 8's windows (results/lab8/capstone.json).
"""
import argparse
import glob
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  # noqa: E402

ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument("run", help="run name under runs/, e.g. quick/base or dev/p50c")
ap.add_argument("--id", help="problem id to print (from the list); without it, list every answer")
ap.add_argument("--answer", type=int, default=0, help="which of the problem's answers (0 = the first)")
ap.add_argument("--chars", type=int, default=1500, help="how many characters to print from the end")
ap.add_argument("--tokens", type=int, nargs=2, metavar=("START", "END"),
                help="also print answer tokens START to END (e.g. a window from results/lab8/capstone.json)")
a = ap.parse_args()

run_dir = os.path.join(hf.LAB, "runs", a.run)
samples = list(loops.iter_samples(run_dir))
if not samples:
    sys.exit(f"no answer texts in {run_dir}; is this the VM that ran {a.run}?")


def ending(s):
    if s["cut"]:
        return "cut at the cap " + ("after thinking" if s["closed"] else "while thinking")
    return "finished" if s["closed"] else "stopped while thinking"


if not a.id:
    print(f"{'problem':36s} {'answer':>6s} {'tokens':>7s}  {'how it ended':29s} {'loops at the end':16s} "
          f"{'correct':7s}  model said / truth")
    for rec, j, s, text in samples:
        loop = "yes" if s["verbatim_loop_end"] else ""
        print(f"{rec['id']:36s} {j:6d} {s['tokens']:7,d}  {ending(s):29s} {loop:16s} {'yes' if s['correct'] else '':7s}"
              f"  {str(s['pred'])[:20]} / {str(rec.get('answer', '-'))[:20]}")
    n, cut = len(samples), sum(s["cut"] for _, _, s, _ in samples)
    print(f"\n{n} answers: {sum(s['correct'] for _, _, s, _ in samples)} correct, {cut} cut at the cap, "
          f"{sum(s['cut'] and s['verbatim_loop_end'] for _, _, s, _ in samples)} of the cut ones end in a verbatim loop")
    print(f"read one: python labs/peek.py {a.run} --id <problem> [--answer N]")
    sys.exit()

hit = [(rec, j, s, text) for rec, j, s, text in samples if rec["id"] == a.id and j == a.answer]
if not hit:
    sys.exit(f"no answer {a.answer} for problem {a.id} in {a.run}; list them with: python labs/peek.py {a.run}")
rec, j, s, text = hit[0]
done = json.load(open(sorted(glob.glob(os.path.join(run_dir, "shard*.done.json")))[0]))
problem = loops.read_rows(hf.LAB, os.path.basename(done["data"]).removesuffix(".jsonl")).get(rec["id"], {})
print(f"problem {rec['id']}: {problem.get('problem', problem.get('prompt', '?'))}")
print(f"truth: {rec.get('answer', '-')}   model said: {s['pred']}   correct: {'yes' if s['correct'] else 'no'}")
print(f"answer {j}: {s['tokens']:,} tokens, {ending(s)}")

if s["verbatim_loop_end"] or a.tokens:
    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained(hf.model_dir("base"))
    ids = tok(text, add_special_tokens=False).input_ids
found = loops.find_loop(ids) if s["verbatim_loop_end"] else None
if found:
    onset, period = found
    print(f"it loops: from token {onset:,} of {len(ids):,}, one block of {period} tokens repeats about "
          f"{(len(ids) - 1 - onset) / period:.0f} times. One cycle:")
    print("    " + repr(tok.decode(ids[onset:onset + period])))
i = text.find("</think>")
if i >= 0:
    print(f"\n--- where it closed its thinking (character {i:,} of {len(text):,}) ---")
    print(text[max(0, i - 400):i + 200])
print(f"\n--- the last {a.chars:,} characters ---")
print(text[-a.chars:])
if a.tokens:
    t0, t1 = max(0, a.tokens[0]), min(len(ids), a.tokens[1])
    print(f"\n--- answer tokens {t0:,} to {t1:,} (of {len(ids):,}) ---")
    print(tok.decode(ids[t0:t1]))