"""Un-blind the scored responses, total them per model and per category, and record the run.

Purpose: turn the filled-in scores.csv from run-blind-comparison.py into a result: a table of
         scores per model overall and by category, the gap between each pair of models set
         against the pre-registered margin, an exact sign-permutation test that says how often
         a gap that large arises between two equally good models, the median decode speed each
         local model reached, and one JSON line appended to the lab notebook so the run can be
         compared with later ones.
Platform: all (spark, strix, mac, nvidia). Pure standard library: no packages to install.
Minimum memory: 12 GB, the same floor as the run it scores; this script itself needs almost none.
Assumes: run-blind-comparison.py has been run in this directory, every row of scores.csv has a
         score of 0, 1 or 2, and labbook.md is the notebook started in Part 1.

Usage: python score-blind-comparison.py --labbook ../labbook.md --note "first run"
       python score-blind-comparison.py --per-task --no-labbook
"""
import argparse
import csv
import json
import statistics
import sys
from collections import defaultdict
from pathlib import Path


def load_scores(path: Path) -> dict:
    scores = {}
    blank = []
    try:
        fh = path.open(encoding="utf-8", newline="")
    except FileNotFoundError:
        raise SystemExit(f"{path} not found: run run-blind-comparison.py first.")
    with fh:
        for row in csv.DictReader(fh):
            rid = (row.get("response_id") or "").strip()
            raw = (row.get("score") or "").strip()
            if not rid:
                continue
            if raw == "":
                blank.append(rid)
                continue
            try:
                value = int(raw)
            except ValueError:
                raise SystemExit(f"{path}: response {rid} has a score of {raw!r}; use 0, 1 or 2.")
            if value not in (0, 1, 2):
                raise SystemExit(f"{path}: response {rid} scored {value}; use 0, 1 or 2.")
            scores[rid] = value
    if blank:
        raise SystemExit(
            f"{path}: {len(blank)} response(s) have no score yet, starting with {blank[0]}. "
            "Score every response before running this, or the comparison is not like for like."
        )
    return scores


def load_json(path: Path, what: str) -> dict:
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except FileNotFoundError:
        raise SystemExit(f"{path} not found: {what}")
    except json.JSONDecodeError as exc:
        raise SystemExit(f"{path} is not valid JSON: {exc}")


def decode_tokens_per_second(record: dict):
    """Ollama reports eval_count tokens generated in eval_duration nanoseconds."""
    count, duration = record.get("eval_count"), record.get("eval_duration")
    if not count or not duration:
        return None
    return count / (duration / 1_000_000_000)


def sign_permutation_p(differences: list) -> tuple:
    """Exact two-sided sign-permutation test on paired per-task differences.

    If the two models were equally good, each task's difference would be as likely to have
    gone the other way. Enumerate every way of flipping the signs of the non-zero differences
    (2^k patterns, counted exactly by convolution) and return the fraction whose total is at
    least as far from zero as the observed total, plus k, the number of tasks that differed.
    """
    nonzero = [d for d in differences if d != 0]
    observed = abs(sum(nonzero))
    dist = {0: 1}
    for d in nonzero:
        nxt = defaultdict(int)
        for total, count in dist.items():
            nxt[total + d] += count
            nxt[total - d] += count
        dist = nxt
    patterns = 2 ** len(nonzero)
    extreme = sum(count for total, count in dist.items() if abs(total) >= observed)
    return extreme / patterns, len(nonzero)


def bar(fraction: float, width: int = 20) -> str:
    filled = round(fraction * width)
    return "#" * filled + "." * (width - filled)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--scores", default="scores.csv")
    parser.add_argument("--key", default="blind-key.json")
    parser.add_argument("--tasks", default="tasks.json")
    parser.add_argument("--labbook", default="labbook.md")
    parser.add_argument("--margin", type=int, default=None,
                        help="pre-registered margin in points (default: scoring.preregistered_margin "
                             "in the task file, else 4)")
    parser.add_argument("--per-task", action="store_true",
                        help="also print every task's score per source, so you can see where a gap came from")
    parser.add_argument("--no-labbook", action="store_true", help="print the result but record nothing")
    parser.add_argument("--note", default="", help="a line of your own to store with the result")
    args = parser.parse_args()

    scores = load_scores(Path(args.scores))
    key = load_json(Path(args.key), "run run-blind-comparison.py first.")
    tasks = load_json(Path(args.tasks), "the task file the run used must be present.")
    task_order = [t["id"] for t in tasks["tasks"]]
    category_of = {t["id"]: t["category"] for t in tasks["tasks"]}
    max_score = tasks.get("scoring", {}).get("max_score", 2)
    margin = args.margin
    if margin is None:
        margin = key.get("preregistered_margin")
    if margin is None:
        margin = tasks.get("scoring", {}).get("preregistered_margin", 4)
    responses = key["responses"]

    unknown = sorted(set(scores) - set(responses))
    if unknown:
        raise SystemExit(f"{args.scores} scores responses that are not in {args.key}: {unknown[:5]}")
    unscored = sorted(set(responses) - set(scores))
    if unscored:
        raise SystemExit(f"{len(unscored)} response(s) in {args.key} were never scored, e.g. {unscored[0]}.")

    totals = defaultdict(int)
    per_category = defaultdict(lambda: defaultdict(int))
    counts = defaultdict(int)
    per_category_counts = defaultdict(lambda: defaultdict(int))
    per_task = defaultdict(dict)
    speeds = defaultdict(list)
    cut_off = defaultdict(int)
    thinking_chars = defaultdict(list)

    for rid, score in scores.items():
        record = responses[rid]
        source, task_id = record["source"], record["task"]
        category = category_of.get(task_id, "uncategorised")
        totals[source] += score
        counts[source] += 1
        per_category[source][category] += score
        per_category_counts[source][category] += 1
        per_task[task_id][source] = score
        speed = decode_tokens_per_second(record)
        if speed is not None:
            speeds[source].append(speed)
        if record.get("done_reason") == "length":
            cut_off[source] += 1
        if record.get("thinking_chars"):
            thinking_chars[source].append(record["thinking_chars"])

    sources = sorted(totals, key=lambda s: (-totals[s], s))
    categories = tasks.get("categories") or sorted({c for c in category_of.values()})
    width = max(len(s) for s in sources)

    print()
    print(f"Scored {len(scores)} responses from {len(sources)} sources, {max_score} points each.")
    if key.get("ollama_version"):
        print(f"Ollama {key['ollama_version']}, thinking {'on' if key.get('think') else 'off'}, "
              f"options {json.dumps(key.get('options'))}")
    print()
    for source in sources:
        possible = counts[source] * max_score
        fraction = totals[source] / possible if possible else 0.0
        print(f"  {source:<{width}}  {totals[source]:>3} / {possible:<3}  {bar(fraction)}  {fraction:6.1%}")
    print()
    print("  By category (score out of the category maximum):")
    header = "    " + " " * width + "".join(f"  {c[:14]:>14}" for c in categories)
    print(header)
    for source in sources:
        cells = ""
        for category in categories:
            got = per_category[source][category]
            possible = per_category_counts[source][category] * max_score
            cells += f"  {f'{got}/{possible}':>14}" if possible else f"  {'-':>14}"
        print(f"    {source:<{width}}{cells}")

    if args.per_task:
        print()
        print("  Per task (score per source; a task where every source scored the same tells you nothing):")
        print("    " + f"{'task':<8}" + "".join(f"  {s[:14]:>14}" for s in sources))
        for task_id in task_order:
            row = per_task.get(task_id, {})
            cells = "".join(f"  {row.get(s, '-'):>14}" for s in sources)
            mark = "" if len({row.get(s) for s in sources}) > 1 else "   (same)"
            print(f"    {task_id:<8}{cells}{mark}")

    print()
    print(f"  Pre-registered margin: {margin} points out of {len(task_order) * max_score}.")
    pairs = []
    for i, a in enumerate(sources):
        for b in sources[i + 1:]:
            diffs = [per_task[t].get(a, 0) - per_task[t].get(b, 0) for t in task_order]
            gap = sum(diffs)
            p_value, differing = sign_permutation_p(diffs)
            within = abs(gap) <= margin
            pairs.append({"a": a, "b": b, "gap": gap, "tasks_differing": differing,
                          "p_two_sided": round(p_value, 4), "within_margin": within})
            print(f"  {a} vs {b}: gap {gap:+d} point(s) over {differing} task(s) that differed; "
                  f"{'within' if within else 'outside'} the margin.")
            print(f"    If the two were equally good, a gap at least this large would arise in "
                  f"{p_value:.1%} of runs ({'not ' if p_value > 0.05 else ''}rare by the usual 5% rule).")

    if speeds:
        print()
        print("  Median decode speed reported by Ollama, tokens per second:")
        for source in sources:
            if speeds[source]:
                print(f"    {source:<{width}}  {statistics.median(speeds[source]):.1f}")
            else:
                print(f"    {source:<{width}}  not measured (answers came from elsewhere)")
    if any(cut_off.values()) or any(thinking_chars.values()):
        print()
        for source in sources:
            notes = []
            if cut_off[source]:
                notes.append(f"{cut_off[source]} answer(s) cut off by num_predict")
            if thinking_chars[source]:
                notes.append(f"median thinking trace {statistics.median(thinking_chars[source]):.0f} characters")
            if notes:
                print(f"    {source:<{width}}  " + "; ".join(notes))

    print()
    print(f"  {len(task_order)} tasks is a small sample. A margin of one or two points is not a finding.")
    print()

    record = {
        "lab": "part-03/reality-check-small-local-versus-frontier",
        "generated": key.get("generated"),
        "ollama_version": key.get("ollama_version"),
        "think": key.get("think"),
        "options": key.get("options"),
        "models": key.get("models"),
        "tasks_file": args.tasks,
        "tasks": len(task_order),
        "max_score": max_score,
        "preregistered_margin": margin,
        "totals": {s: {"score": totals[s], "possible": counts[s] * max_score} for s in sources},
        "by_category": {s: {c: {"score": per_category[s][c],
                               "possible": per_category_counts[s][c] * max_score}
                            for c in categories if per_category_counts[s][c]} for s in sources},
        "pairs": pairs,
        "median_decode_tokens_per_second": {
            s: round(statistics.median(v), 2) for s, v in speeds.items() if v},
        "cut_off_by_num_predict": {s: cut_off[s] for s in sources if cut_off[s]},
        "note": args.note,
    }

    if args.no_labbook:
        print(json.dumps(record))
        return
    labbook = Path(args.labbook)
    if not labbook.exists():
        print(f"{labbook} does not exist; creating it.", file=sys.stderr)
    with labbook.open("a", encoding="utf-8") as fh:
        fh.write(json.dumps(record) + "\n")
    print(f"  recorded in {labbook}")


if __name__ == "__main__":
    main()