"""Benchmark TinyLiquid: val perplexity, forensic probes, samples, speed.

Writes bench/metrics.json for the model card and prints a summary.

Usage:
  .venv/bin/python eval/bench.py --ckpt ckpt/dpo --out bench/metrics.json
"""

import argparse
import json
import time
from pathlib import Path

import numpy as np
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 = [
    ("Discrepancy", "Two accounts describe the same event. Account A: 'The meeting started at 9am and ended at 11am.' Account B: 'The meeting started at 9am and ran until noon.' Find discrepancies.", "end time"),
    ("Contradiction", "Account A: 'No officials were present.' Account B: 'An official arrived later.' What can you conclude?", "not necessarily a contradiction"),
    ("Overclaim", "Evaluate: 'Crime doubled last year because of the new policy.' Incidents rose 1,000 to 2,000 while reporting methods changed.", "overclaim"),
    ("Pattern", "What are the weak links in a theory claiming one actor caused three unrelated disasters?", "no direct evidence"),
    ("Verification", "Verify: 'The building was built in 2015 and never renovated.' Assessor record: built 2015, renovation permit 2019.", "false"),
    ("Timeline", "Witness saw a van at 3:30pm; outage reported 3:40pm; inspection 6pm; no report before 5pm. Reconstruct.", "gap"),
    ("History", "2019: 'officials rule out sabotage.' 2021 filing: a second investigation opened the same week. What changed?", "parallel investigation hidden"),
    ("Politics", "Statement: 'We balanced the budget without raising taxes.' Budget includes a reassessment raising collections 9%.", "misleading"),
    ("Source chain", "A claim rests on: company blog, a wire story repeating it, an analyst note quoting the wire. Rate the evidence.", "single chain"),
]


def load(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 = TinyLiquidConfig(vocab_size=tok.get_vocab_size(),
                           **{k: v for k, v in sd["config"].items() if k != "vocab_size"})
    model = TinyLiquid(cfg)
    model.load_state_dict(sd["model"])
    model.eval()
    return tok, model, ckpt, cfg


@torch.no_grad()
def val_ppl(model, tok, val_bin, batches, batch, seq):
    arr = np.fromfile(val_bin, dtype=np.uint16).astype(np.int64)
    t = torch.from_numpy(arr)
    n = (len(t) - 1) // seq
    rng = np.random.RandomState(0)
    total, count = 0.0, 0
    for _ in range(batches):
        s = int(rng.randint(0, n - batch))
        idx = torch.arange(s * seq, (s + batch) * seq, dtype=torch.long)
        buf = torch.stack([t[int(i): int(i) + seq] for i in idx])
        x, y = buf[:, :-1], buf[:, 1:]
        logits = model(x)
        loss = torch.nn.functional.cross_entropy(
            logits.view(-1, logits.size(-1)), y.reshape(-1))
        total += loss.item() * y.numel()
        count += y.numel()
    return float(np.exp(total / count))


@torch.no_grad()
def probe_hits(model, tok):
    hits = 0
    for name, q, want in PROBES:
        prompt = "<|analyst|><|user|>" + q + "<|assistant|>"
        ids = tok.encode(prompt).ids
        out = tok.decode(model.generate(tok, ids, persona_id=1, max_new=60,
                                        temperature=0.4, top_k=40,
                                        repetition_penalty=1.5,
                                        no_repeat_ngram_size=4)[len(ids):]).lower()
        words = want.split()
        hit = any(w in out for w in words)
        hits += int(hit)
    return hits, len(PROBES)


@torch.no_grad()
def speed(model, tok, n_tokens=40):
    ids = tok.encode("<|analyst|><|user|>Evaluate this claim: 'X caused Y.'<|assistant|>").ids
    t0 = time.time()
    model.generate(tok, ids, persona_id=1, max_new=n_tokens, temperature=0.6,
                   top_k=40, repetition_penalty=1.4, no_repeat_ngram_size=4)
    dt = time.time() - t0
    return n_tokens / dt


def samples(model, tok):
    prompts = [
        "<|analyst|><|user|>Verify: 'The bridge was painted in 2019 and never repainted.' Records show a 2022 repaint permit.<|assistant|>",
        "<|analyst|><|user|>What's the most common mistake you see in research?<|assistant|>",
        "<|skeptic|><|user|>Attack this conclusion: 'Three failures in one week with vans nearby is deliberate.'<|assistant|>",
    ]
    out = []
    for p in prompts:
        persona = 2 if p.startswith("<|skeptic|>") else 1
        ids = tok.encode(p).ids
        gen = tok.decode(model.generate(tok, ids, persona_id=persona, max_new=90,
                                        temperature=0.6, top_k=40,
                                        repetition_penalty=1.4,
                                        no_repeat_ngram_size=4)[len(ids):]).strip()
        out.append({"prompt": p.split("<|user|>")[1].split("<|assistant|>")[0],
                    "persona": "skeptic" if persona == 2 else "analyst",
                    "output": gen})
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", default="ckpt/dpo")
    ap.add_argument("--tok", default="data/tokenizer.json")
    ap.add_argument("--val", default="data/valid.bin")
    ap.add_argument("--out", default="bench/metrics.json")
    ap.add_argument("--val-batches", type=int, default=10)
    ap.add_argument("--batch", type=int, default=16)
    ap.add_argument("--seq", type=int, default=256)
    ap.add_argument("--threads", type=int, default=8)
    args = ap.parse_args()

    tok, model, ckpt, cfg = load(args)
    params = sum(p.numel() for p in model.parameters())
    print(f"benchmarking {ckpt} | params {params:,} | d_model {cfg.d_model} blocks {cfg.n_blocks}", flush=True)

    ppl = val_ppl(model, tok, args.val, args.val_batches, args.batch, args.seq)
    hits, total = probe_hits(model, tok)
    tok_s = speed(model, tok)
    smpls = samples(model, tok)

    metrics = {
        "checkpoint": str(ckpt),
        "params": params,
        "val_loss": round(float(np.log(ppl)), 4),
        "val_ppl": round(ppl, 4),
        "probe_hits": f"{hits}/{total}",
        "probe_accuracy": round(hits / total, 3),
        "gen_speed_tok_per_s": round(tok_s, 1),
        "hardware": "8-core ARM, no GPU",
        "samples": smpls,
    }
    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(metrics, indent=2), encoding="utf-8")
    print(json.dumps(metrics, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()