#!/usr/bin/env python3
"""
longeval_score.py -- score CP_GENERATE output against longeval_gen.py metadata.

This is REAL benchmark scoring (substring match on generated text), not the NLL
proxy needle_score.py computes. A point is awarded only when every expected
string appears in the generation, which is RULER's recall-style criterion.

  usage: longeval_score.py --gen <generations.csv> --meta <task_meta.json>
                           [--out results.csv] [--show 3]

Reports overall accuracy plus a per-depth breakdown -- the depth axis is the
lost-in-the-middle curve, so a model that only retrieves from the end of the
window shows a strong depth gradient even at high overall accuracy.
"""

import argparse
import csv
import json
import sys

try:
    import tiktoken
except ImportError:
    sys.exit("longeval_score.py needs tiktoken:  pip install tiktoken")

ENC = tiktoken.get_encoding("gpt2")


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


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--gen", required=True, help="CP_GEN_OUT csv: i,tok,tok,...")
    ap.add_argument("--meta", required=True, help="<task>_meta.json")
    ap.add_argument("--out", default="", help="optional per-item csv")
    ap.add_argument("--show", type=int, default=3, help="print N sample generations")
    args = ap.parse_args()

    meta = json.load(open(args.meta))
    items = {it["i"]: it for it in meta["items"]}

    gens = {}
    with open(args.gen) as f:
        for row in csv.reader(f):
            if not row:
                continue
            i = int(row[0])
            gens[i] = [int(t) for t in row[1:] if t.strip() != ""]

    missing = sorted(set(items) - set(gens))
    if missing:
        print(f"  WARN: {len(missing)} prompts have no generation "
              f"(first: {missing[:5]}) -- scoring only what is present")

    rows, n_ok = [], 0
    by_depth = {}
    for i in sorted(gens):
        if i not in items:
            continue
        it = items[i]
        text = ENC.decode(gens[i])
        tn = norm(text)
        hits = [e for e in it["expected"] if norm(e) in tn]
        ok = len(hits) == len(it["expected"])
        n_ok += ok
        d = it["depth"]
        by_depth.setdefault(d, [0, 0])
        by_depth[d][0] += ok
        by_depth[d][1] += 1
        rows.append({"i": i, "depth": d, "ok": int(ok),
                     "n_expected": len(it["expected"]), "n_hit": len(hits),
                     "expected": "|".join(it["expected"]),
                     "generated": text.replace("\n", "\\n")[:200]})

    n = len(rows)
    if n == 0:
        sys.exit("  no scorable prompts -- check --gen and --meta pair up")

    print(f"\n=== {meta['task']}  (prompt_len={meta['prompt_len']}, n={n}) ===")
    print(f"  accuracy: {n_ok}/{n} = {100.0 * n_ok / n:.1f}%")
    print(f"  {'depth':>7} {'acc':>8}   (needle position in context)")
    for d in sorted(by_depth):
        ok, tot = by_depth[d]
        print(f"  {d:>7.2f} {ok:>3}/{tot:<3} {100.0 * ok / tot:5.1f}%")

    if args.show:
        print(f"\n  --- {min(args.show, n)} sample generations ---")
        for r in rows[:args.show]:
            mark = "HIT " if r["ok"] else "MISS"
            print(f"  [{mark}] expect={r['expected']}")
            print(f"         got={r['generated'][:150]}")

    if args.out:
        with open(args.out, "w", newline="") as f:
            w = csv.DictWriter(f, fieldnames=list(rows[0]))
            w.writeheader()
            w.writerows(rows)
        print(f"\n  wrote {args.out}")


if __name__ == "__main__":
    main()