#!/usr/bin/env python3
"""
Reproduce the UDI v0.3.1 LongMemEval-V2 Small result — from frozen
official-scorer judgments, offline, in one command.

    python3 reproduce.py

What this proves
----------------
The sealed UDI adapter already produced its 451 answers, and the benchmark's
official scorer already judged each one correct/incorrect. This script recomputes
the headline numbers straight from those frozen per-question judgments. It needs
no network, no API key, and — importantly — no access to the protected adapter:
the mechanism stays sealed; only its scored OUTPUT is checked here.

If the recomputed numbers match the published claim (they do), then the claim is
arithmetically faithful to the frozen scoring log. To verify the JUDGE itself
(re-run the official LLM grader over the answers), see README.md — that step
needs the official harness and an API key, and is a deeper, separate check.
"""
from __future__ import annotations
import hashlib
import json
import os
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
DATA = os.path.join(HERE, "data", "judged_frozen.jsonl")

# Published claim (from the frozen submission_overview.json, tier=small, op=fast)
EXPECTED = {
    "total": 451,
    "correct": 256,
    "accuracy": 0.5676274944567627,          # 56.762749%
    "memory_query_avg_seconds": 0.187002020206575,
    "lafs_gain": 3.5807026308063215,          # derived; see note below
    "reference_lafs": 55.76484693638005,
    "submission_lafs": 59.345549567186374,
}
ACC_TOL = 1e-9
LAT_TOL = 1e-6


def sha256(path: str) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(65536), b""):
            h.update(chunk)
    return h.hexdigest()


def main() -> int:
    if not os.path.isfile(DATA):
        print(f"FAIL: missing frozen data file: {DATA}")
        return 2

    total = correct = 0
    latency_sum = 0.0
    by_cat: dict[str, list[int]] = {}
    for line in open(DATA, encoding="utf-8"):
        line = line.strip()
        if not line:
            continue
        d = json.loads(line)
        total += 1
        hit = 1 if d.get("score_bool") else 0
        correct += hit
        latency_sum += float(d.get("memory_query_duration_seconds") or 0.0)
        cat = d.get("category", "?")
        by_cat.setdefault(cat, [0, 0])
        by_cat[cat][0] += hit
        by_cat[cat][1] += 1

    accuracy = correct / total if total else 0.0
    latency = latency_sum / total if total else 0.0

    checks = [
        ("questions scored", total, EXPECTED["total"], total == EXPECTED["total"]),
        ("answers correct", correct, EXPECTED["correct"], correct == EXPECTED["correct"]),
        ("accuracy", accuracy, EXPECTED["accuracy"], abs(accuracy - EXPECTED["accuracy"]) < ACC_TOL),
        ("avg memory-query latency (s)", latency, EXPECTED["memory_query_avg_seconds"],
         abs(latency - EXPECTED["memory_query_avg_seconds"]) < LAT_TOL),
    ]

    print("=" * 66)
    print("UDI v0.3.1 — LongMemEval-V2 Small — offline reproduction")
    print("=" * 66)
    print(f"frozen judgments : data/judged_frozen.jsonl")
    print(f"sha256           : {sha256(DATA)}")
    print("-" * 66)
    ok = True
    for name, got, exp, passed in checks:
        ok = ok and passed
        g = f"{got:.12f}" if isinstance(got, float) else str(got)
        e = f"{exp:.12f}" if isinstance(exp, float) else str(exp)
        print(f"  [{'OK' if passed else 'XX'}] {name:<32} {g:>18}  (published {e})")
    print("-" * 66)
    print(f"  accuracy = {correct}/{total} = {accuracy*100:.6f}%   "
          f"avg latency = {latency:.6f}s")
    print(f"  LAFS gain (derived by official leaderboard) = "
          f"+{EXPECTED['lafs_gain']:.6f} pts "
          f"({EXPECTED['submission_lafs']:.4f} vs {EXPECTED['reference_lafs']:.4f})")
    print("-" * 66)
    print("  category breakdown (accuracy):")
    for cat in sorted(by_cat):
        c, n = by_cat[cat]
        print(f"    {cat:<26} {c:>3}/{n:<3} = {100*c/n:6.2f}%")
    print("=" * 66)
    if ok:
        print("PASS — recomputed headline numbers match the published claim exactly.")
        print("Note: this checks the SCORE against frozen judgments. It is not an")
        print("independent re-run of the judge, and not an accepted leaderboard entry.")
        return 0
    print("FAIL — recomputed numbers do NOT match the published claim.")
    return 1


if __name__ == "__main__":
    sys.exit(main())