"""Run a fixed set of forensic probes through a TinyLiquid checkpoint.

Usage:
  .venv/bin/python research/probe.py --ckpt ckpt/distill
"""

import argparse
from pathlib import Path

import torch

from model.config import TinyLiquidConfig, CONFIGS
from model.tiny_liquid import TinyLiquid
from model.utils import latest_ckpt
from data.tokenizer import load_tokenizer

PROBES = [
    ("analyst", "Find discrepancies between: Account A: The meeting started at 9am and ended at 11am. Account B: The meeting started at 9am and ran until noon."),
    ("analyst", "Two accounts describe the same event. Account A: 'No officials were present.' Account B: 'An official arrived later.' What can you conclude?"),
    ("analyst", "Evaluate this claim: 'Crime in the city doubled last year because of the new policy.' The report shows incidents rose from 1,000 to 2,000 while reporting methods changed."),
    ("analyst", "What are the weak links in a theory claiming one actor caused three unrelated disasters?"),
    ("skeptic", "Attack this conclusion: 'The stock dropped after the announcement, so investors rejected the announcement.'"),
]


def parse_args():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", default="ckpt/distill")
    ap.add_argument("--tok", default="data/tokenizer.json")
    ap.add_argument("--max-new", type=int, default=140)
    ap.add_argument("--threads", type=int, default=8)
    return ap.parse_args()


def main():
    args = parse_args()
    torch.set_num_threads(args.threads)
    tok = load_tokenizer(args.tok)
    ckpt = latest_ckpt(args.ckpt)
    assert ckpt, f"no checkpoints in {args.ckpt}"
    sd = torch.load(ckpt, map_location="cpu")
    cfg_dict = dict(sd.get("config", CONFIGS["tiny10m"]))
    cfg = TinyLiquidConfig(vocab_size=tok.get_vocab_size(),
                           **{k: v for k, v in cfg_dict.items() if k != "vocab_size"})
    model = TinyLiquid(cfg)
    model.load_state_dict(sd["model"])
    model.eval()
    print(f"== {ckpt} (step {sd.get('step','?')}) ==\n", flush=True)

    P_TOKEN = {"analyst": "<|analyst|>", "skeptic": "<|skeptic|>"}
    P_ID = {"analyst": 1, "skeptic": 2}
    for persona, prompt in PROBES:
        p = P_TOKEN[persona] + "<|user|>" + prompt + "<|assistant|>"
        ids = tok.encode(p).ids
        out = model.generate(tok, ids, persona_id=P_ID[persona], max_new=args.max_new,
                             temperature=0.65, top_k=40, repetition_penalty=1.4,
                             no_repeat_ngram_size=4)
        print(f"--- [{persona}] {prompt}\n{tok.decode(out[len(ids):])}\n", flush=True)


if __name__ == "__main__":
    main()