#!/usr/bin/env python3
"""Score the licensed benchmark with REPO-CLUSTERED confidence intervals, and emit per-row verdicts.

Why clustered. Rows drawn from the same repository are not independent: they share an author, a
house style, and often helper functions, so a plain binomial interval understates uncertainty.
The old 400-row benchmark took 15% of its rows from a single repo, which made clustering so
severe that a clustered interval would have been the only honest one to quote — and none was.
The licensed benchmark caps any repo at ~1%, so the cluster effect is small, but it is reported
rather than assumed small.

Method: the cluster bootstrap. Resample REPOSITORIES with replacement (not rows), recompute the
accuracy over the resampled repos, and take the 2.5th/97.5th percentiles. This is the standard
non-parametric interval for clustered binary data and needs no normality assumption. The design
effect (clustered variance / binomial variance) is reported so a reader can see how much the
clustering actually cost.

    ./analyze_scores.py --bench ../benchmarks/csn-3.12-licensed/bench.jsonl \\
        --greedy ../generations/gen_v3_csn600.jsonl \\
        --samples ../generations/boN_v3_csn600.jsonl \\
        --base ../generations/gen_base_csn600.jsonl \\
        --out ../results/scores_csn600.json --rows-out ../results/rows_csn600.jsonl
"""
from __future__ import annotations

import argparse
import json
import math
import random
import sys
from collections import defaultdict
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

from common import load_bench, load_jsonl, ours_ok, require_sound, self_test, strip_fences  # noqa: E402

BOOT = 10000
SEED = 20260804


def cluster_bootstrap(per_repo: dict[str, list[int]], boot: int = BOOT, seed: int = SEED):
    """95% CI for the mean of clustered binary outcomes, by resampling repos with replacement."""
    repos = list(per_repo)
    rng = random.Random(seed)
    n_all = sum(len(v) for v in per_repo.values())
    point = sum(sum(v) for v in per_repo.values()) / n_all
    means = []
    for _ in range(boot):
        num = den = 0
        for _ in range(len(repos)):
            v = per_repo[repos[rng.randrange(len(repos))]]
            num += sum(v)
            den += len(v)
        if den:
            means.append(num / den)
    means.sort()
    lo = means[int(0.025 * len(means))]
    hi = means[int(0.975 * len(means))]

    # design effect vs the naive binomial interval
    binom_se = math.sqrt(point * (1 - point) / n_all) if 0 < point < 1 else 0.0
    clust_se = (hi - lo) / (2 * 1.96) if hi > lo else 0.0
    deff = (clust_se / binom_se) ** 2 if binom_se > 0 else float("nan")
    return {
        "point": round(100 * point, 2),
        "ci95_lo": round(100 * lo, 2),
        "ci95_hi": round(100 * hi, 2),
        "half_width_pp": round(100 * (hi - lo) / 2, 2),
        "binomial_ci95_lo": round(100 * max(0.0, point - 1.96 * binom_se), 2),
        "binomial_ci95_hi": round(100 * min(1.0, point + 1.96 * binom_se), 2),
        "design_effect": round(deff, 2) if deff == deff else None,
        "n": n_all,
        "clusters": len(repos),
    }


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--bench", required=True)
    ap.add_argument("--greedy", required=True)
    ap.add_argument("--samples")
    ap.add_argument("--base", help="untuned-base control generations")
    ap.add_argument("--out", required=True)
    ap.add_argument("--rows-out")
    ap.add_argument("--label", default="")
    a = ap.parse_args()

    bench, _ = load_bench(a.bench)
    st = self_test(bench, "ours")
    print(f"self-test: preflight {st['preflight_pct']}%  mutation kill "
          f"{st['mutation_kill_rate_pct']}%", file=sys.stderr, flush=True)
    require_sound(st)

    greedy = {r["i"]: r for r in load_jsonl(a.greedy)}
    base = {r["i"]: r for r in load_jsonl(a.base)} if a.base else {}
    samples: dict[int, dict[int, str]] = defaultdict(dict)
    if a.samples and Path(a.samples).exists():
        for r in load_jsonl(a.samples):
            samples[r["i"]][r["s"]] = r["got"]

    Ns = (1, 2, 4, 8, 16, 32)
    rows = []
    for i in sorted(bench):
        exp = bench[i]["expected"]
        g_ok = i in greedy and ours_ok(strip_fences(greedy[i]["got"]), exp)
        b_ok = i in base and ours_ok(strip_fences(base[i]["got"]), exp)
        first = None
        if not g_ok:
            for s in sorted(samples.get(i, {})):
                if ours_ok(strip_fences(samples[i][s]), exp):
                    first = s
                    break
        rows.append({
            "i": i,
            "repo": bench[i]["provenance"]["repo"],
            "func": bench[i]["provenance"]["func_name"],
            "commit_sha": bench[i]["provenance"]["commit_sha"],
            "spdx": bench[i]["license"]["spdx"],
            "n_instr": bench[i]["n_instr"],
            "tuned_greedy_certified": bool(g_ok),
            "base_greedy_certified": bool(b_ok),
            "first_passing_sample": first,
            "certified_at_32": bool(g_ok or (first is not None and first <= 30)),
        })

    def by_repo(key):
        d = defaultdict(list)
        for r in rows:
            d[r["repo"]].append(1 if r[key] else 0)
        return d

    report = {
        "label": a.label,
        "bench": str(a.bench),
        "n": len(rows),
        "repos": len({r["repo"] for r in rows}),
        "oracle": "strict L1: byte-identical code object incl. docstrings and co_exceptiontable",
        "self_test": {k: st[k] for k in ("preflight_pct", "mutation_kill_rate_pct", "SOUND")},
        "tuned_greedy": cluster_bootstrap(by_repo("tuned_greedy_certified")),
        "certified_at_32": cluster_bootstrap(by_repo("certified_at_32")),
    }
    if base:
        report["base_greedy_UNTUNED_CONTROL"] = cluster_bootstrap(by_repo("base_greedy_certified"))
        t = sum(r["tuned_greedy_certified"] for r in rows)
        b = sum(r["base_greedy_certified"] for r in rows)
        # paired: rows where exactly one of the two certified
        b01 = sum(1 for r in rows if r["tuned_greedy_certified"] and not r["base_greedy_certified"])
        b10 = sum(1 for r in rows if r["base_greedy_certified"] and not r["tuned_greedy_certified"])
        report["fine_tune_effect"] = {
            "tuned_certified": t, "base_certified": b,
            "absolute_gain_pp": round(100 * (t - b) / len(rows), 2),
            "tuned_only": b01, "base_only": b10,
            "note": "paired discordant counts; McNemar exact p below",
        }
        # exact McNemar (binomial on the discordant pairs)
        n_d = b01 + b10
        if n_d:
            # Exact binomial two-sided p on the discordant pairs. Computed in LOG space: with a
            # few hundred discordant pairs the direct ratio underflows to 0.0, and reporting a
            # p-value of exactly zero would be a floating-point artefact presented as a result.
            log2_p = math.log2(2.0) + math.log2(
                sum(math.comb(n_d, k) for k in range(min(b01, b10) + 1))) - n_d
            p = 2 ** log2_p
            if p >= 1e-12:
                report["fine_tune_effect"]["mcnemar_exact_p"] = round(min(1.0, p), 12)
            else:
                # report the magnitude honestly instead of collapsing it to 0
                report["fine_tune_effect"]["mcnemar_exact_p"] = f"< 1e-12 (log10 p ~ {log2_p * math.log10(2):.0f})"

    # N is a budget of N TOTAL attempts: the greedy decode, then N-1 sampled candidates. So the
    # sampled candidates admitted at budget N are indices 0..N-2, and certified@1 is greedy alone.
    # (Indexing this as `<= N - 1` counts N+1 attempts and makes certified@1 report a two-attempt
    # number, which overstates every point on the curve.)
    curve = {}
    for N in Ns:
        c = sum(1 for r in rows if r["tuned_greedy_certified"]
                or (r["first_passing_sample"] is not None and r["first_passing_sample"] <= N - 2))
        curve[str(N)] = {"certified": c, "pct": round(100 * c / len(rows), 2),
                         "budget": f"1 greedy + {N - 1} sampled"}
    report["certified_curve"] = curve

    Path(a.out).parent.mkdir(parents=True, exist_ok=True)
    Path(a.out).write_text(json.dumps(report, indent=2))
    if a.rows_out:
        Path(a.rows_out).write_text("".join(json.dumps(r) + "\n" for r in rows))
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()