#!/usr/bin/env python3
"""
peer_score.py -- score peer generations, matching longeval_score.py EXACTLY.

The scoring rule is copied deliberately, not reimplemented: norm() strips every
non-alphanumeric character and lowercases, and a prompt scores 1 only when EVERY
expected string appears in the normalised generation (RULER's recall criterion).
Any divergence here would make peer numbers incomparable to ours, which is the
entire point of the exercise.

Writes summary.csv and depth_summary.csv in the SAME shape run_longeval.sh emits,
so the existing notebook builders read peer results without modification.

  usage: peer_score.py --gen-dir gen/ --meta-dir prompts/ --out-dir results/
                       --models qwen3-4b,spark-4b --tasks niah_single,val_backward
"""

import argparse
import csv
import json
import math
import os


def norm(s):
    return "".join(ch for ch in s.lower() if ch.isalnum())


def wilson(h, n):
    if not n:
        return float("nan")
    z, p = 1.96, h / n
    return 100.0 * z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / (1 + z * z / n)


def score_one(gen_path, meta_path):
    """-> (hits, n, {depth: [hits, n]})"""
    meta = json.load(open(meta_path))
    items = {it["i"]: it for it in meta["items"]}
    hits = n = 0
    by_depth = {}
    with open(gen_path) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            g = json.loads(line)
            it = items.get(g["i"])
            if it is None:
                continue
            tn = norm(g["generated"])
            ok = all(norm(e) in tn for e in it["expected"])
            hits += ok
            n += 1
            d = it["depth"]
            by_depth.setdefault(d, [0, 0])
            by_depth[d][0] += ok
            by_depth[d][1] += 1
    return hits, n, by_depth


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--gen-dir", required=True)
    ap.add_argument("--meta-dir", required=True)
    ap.add_argument("--out-dir", required=True)
    ap.add_argument("--length", required=True,
                    help="prompt length; metas live at <meta-dir>/<model>/<length>/")
    ap.add_argument("--models", required=True, help="comma list of model keys")
    ap.add_argument("--tasks", required=True, help="comma list of task names")
    args = ap.parse_args()

    models = args.models.split(",")
    tasks = args.tasks.split(",")
    os.makedirs(args.out_dir, exist_ok=True)

    per_task = {}     # task -> model -> (hits, n)
    per_depth = {}    # depth -> model -> [hits, n]
    for m in models:
        for t in tasks:
            g = os.path.join(args.gen_dir, f"{m}_{t}.jsonl")
            # layout is <meta-dir>/<model>/<length>/<task>_meta.json -- the
            # length level was missing here, so every meta lookup failed and the
            # scorer wrote a header-only summary instead of erroring.
            md = os.path.join(args.meta_dir, m, args.length, f"{t}_meta.json")
            if not (os.path.exists(g) and os.path.exists(md)):
                print(f"  skip {m}/{t}: missing {'gen' if not os.path.exists(g) else 'meta'}")
                continue
            h, n, bd = score_one(g, md)
            if not n:
                print(f"  skip {m}/{t}: 0 scorable"); continue
            per_task.setdefault(t, {})[m] = (h, n)
            for d, (dh, dn) in bd.items():
                cell = per_depth.setdefault(d, {}).setdefault(m, [0, 0])
                cell[0] += dh; cell[1] += dn
            print(f"  {m:<16} {t:<16} {h}/{n} = {100.0*h/n:5.1f}% +-{wilson(h,n):4.1f}")

    have = [m for m in models if any(m in v for v in per_task.values())]
    if not have:
        raise SystemExit(f"  FATAL: nothing scorable for length {args.length}. "
                         f"Checked gen={args.gen_dir}/<model>_<task>.jsonl and "
                         f"meta={args.meta_dir}/<model>/{args.length}/<task>_meta.json")
    with open(os.path.join(args.out_dir, "summary.csv"), "w", newline="") as f:
        w = csv.writer(f); w.writerow(["task"] + have)
        for t in sorted(per_task):
            w.writerow([t] + [f"{100.0*per_task[t][m][0]/per_task[t][m][1]:.1f}"
                              if m in per_task[t] else "" for m in have])
    with open(os.path.join(args.out_dir, "depth_summary.csv"), "w", newline="") as f:
        w = csv.writer(f); w.writerow(["depth", "model", "hits", "n", "pct", "wilson95"])
        for d in sorted(per_depth):
            for m in have:
                if m in per_depth[d]:
                    h, n = per_depth[d][m]
                    w.writerow([d, m, h, n, f"{100.0*h/n:.1f}", f"{wilson(h,n):.1f}"])
    print(f"\n  wrote {args.out_dir}/summary.csv and depth_summary.csv")


if __name__ == "__main__":
    main()