"""Summarize inert result JSON; paired problem bootstrap, never execute candidates."""

import argparse
import json
from collections import defaultdict
from pathlib import Path

import numpy as np

CONFIGS = ["base", "rE-t050-default", "rE-t050-low", "rE-t050-high", "rA-t030-default"]
ROOT = Path(__file__).resolve().parents[1] / "data/eval/code"


def problem_means(records, bench, metric):
    grouped = defaultdict(list)
    for record in records:
        if record["bench"] == bench:
            grouped[record["task_id"]].append(record)
    result = {}
    for task_id, rows in grouped.items():
        assert {r["sample_index"] for r in rows} == {0, 1, 2, 3} and len(rows) == 4
        result[task_id] = sum(r[metric] for r in rows) / 4
    return result


def paired_ci(base, other, seed=20261030, replicates=20000):
    assert base.keys() == other.keys(), "Paired comparisons require identical problem IDs"
    keys = sorted(base)
    differences = np.array([other[k] - base[k] for k in keys], dtype=np.float64)
    rng = np.random.default_rng(seed)
    bootstrap = np.empty(replicates)
    for start in range(0, replicates, 1000):
        size = min(1000, replicates - start)
        ix = rng.integers(0, len(keys), size=(size, len(keys)))
        bootstrap[start:start + size] = differences[ix].mean(axis=1)
    lo, hi = np.quantile(bootstrap, [0.025, 0.975])
    return {"difference": float(differences.mean()), "ci95": [float(lo), float(hi)],
            "problems": len(keys), "bootstrap_replicates": replicates, "seed": seed}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--run", required=True)
    args = parser.parse_args()
    run = ROOT / args.run
    records, summaries, manifests = {}, {}, {}
    for config in CONFIGS:
        folder = run / config
        manifests[config] = json.loads((folder / "manifest.json").read_text())
        assert manifests[config]["generation_complete"], f"Incomplete configuration: {config}"
        records[config] = [json.loads(x) for x in (folder / "results.jsonl").read_text().splitlines()]
        summaries[config] = json.loads((folder / "summary.json").read_text())
        assert len(records[config]) == 4 * (164 + 378)
        manifest = manifests[config]
        assert (manifest["loops"], manifest["temperature"], manifest["top_p"],
                manifest["samples_per_problem"], manifest["max_tokens"],
                manifest["max_model_len"]) == (2, 1.0, 0.95, 4, 65536, 69632)
        assert all(not r["plus_pass"] or r["base_pass"] for r in records[config])
        assert all(r["thinking_closed"] or not (r["base_pass"] or r["plus_pass"])
                   for r in records[config])
        for bench, expected_n in (("humaneval", 164), ("mbpp", 378)):
            rs = [r for r in records[config] if r["bench"] == bench]
            means = problem_means(records[config], bench, "plus_pass")
            assert len(means) == expected_n and len(rs) == expected_n * 4
            summary = summaries[config]["benches"][bench]
            for field, metric in (("base_pass_at_1", "base_pass"), ("plus_pass_at_1", "plus_pass"),
                                  ("mean_output_tokens", "tokens")):
                assert abs(summary[field] - sum(r[metric] for r in rs) / len(rs)) < 1e-10
        assert manifest["dataset_hashes"] == manifests["base"]["dataset_hashes"]
    prompts = [json.loads((run / c / "prompts.jsonl").read_text().splitlines()[0])["sha256"] for c in CONFIGS]
    assert len(set(prompts)) == 1
    # Check every problem's actual user prompt, not only the first.
    prompt_maps = [{r["task_id"]: r["sha256"] for r in map(json.loads, (run / c / "prompts.jsonl").read_text().splitlines())} for c in CONFIGS]
    assert all(m == prompt_maps[0] for m in prompt_maps)
    comparisons = {}
    for config in CONFIGS[1:]:
        comparisons[config] = {}
        for bench in ("humaneval", "mbpp"):
            comparisons[config][bench] = {}
            for metric in ("base_pass", "plus_pass"):
                comparisons[config][bench][metric] = paired_ci(
                    problem_means(records["base"], bench, metric),
                    problem_means(records[config], bench, metric),
                )
    result = {"run": args.run, "summaries": summaries, "paired_vs_base": comparisons}
    (ROOT / "summary.json").write_text(json.dumps(result, indent=2) + "\n")
    lines = [
        "# Dory coding ability after Beam reasoning distillation", "",
        "All configurations use 2 recurrent loops. Pass@1 is the mean of each problem's four sampled correctness indicators. Percentages below are not pass@4.", "",
        "`base` is NVIDIA Dory1.0-4B (`/weights/dory-4b`); `rE-t050` is Dory-Beam-4B v2 (`/ckpt/merges/rE-t050`); `rA-t030` is Dory-Beam-4B-Accurate (`/ckpt/merges/rA-t030`).", "",
        "| Configuration | HumanEval base | HumanEval+ | MBPP base | MBPP+ | HE mean tokens | MBPP mean tokens | HE truncated | MBPP truncated |",
        "|---|---:|---:|---:|---:|---:|---:|---:|---:|",
    ]
    for config in CONFIGS:
        h = summaries[config]["benches"]["humaneval"]
        m = summaries[config]["benches"]["mbpp"]
        lines.append(f"| {config} | {100*h['base_pass_at_1']:.2f}% | {100*h['plus_pass_at_1']:.2f}% | {100*m['base_pass_at_1']:.2f}% | {100*m['plus_pass_at_1']:.2f}% | {h['mean_output_tokens']:.0f} | {m['mean_output_tokens']:.0f} | {100*h['truncation_rate']:.2f}% | {100*m['truncation_rate']:.2f}% |")
    lines += ["", "## Paired differences against base", "",
              "Differences and 95% percentile bootstrap confidence intervals are in percentage points. Each replicate resamples problems with replacement, keeping the two models paired and keeping all four samples within a problem together. 20,000 replicates; NumPy seed 20261030. Intervals are marginal, without multiple-comparison correction.", "",
              "| Configuration | HumanEval base Δ [95% CI] | HumanEval+ Δ [95% CI] | MBPP base Δ [95% CI] | MBPP+ Δ [95% CI] |",
              "|---|---:|---:|---:|---:|"]
    for config in CONFIGS[1:]:
        cells = []
        for bench in ("humaneval", "mbpp"):
            for metric in ("base_pass", "plus_pass"):
                stat = comparisons[config][bench][metric]
                cells.append(f"{100*stat['difference']:+.2f} [{100*stat['ci95'][0]:+.2f}, {100*stat['ci95'][1]:+.2f}]")
        lines.append("| " + config + " | " + " | ".join(cells) + " |")
    lines += ["", "## Protocol and artifacts", "",
              "- [EvalPlus 0.3.1](https://github.com/evalplus/evalplus/releases/tag/v0.3.1), full HumanEval+ v0.1.10 (164 problems) and MBPP+ v0.2.0 (378 problems). Four samples/problem/config, 10,840 scored samples total. No mini/noextreme filtering.",
              "- Temperature 1.0, top_p 0.95, max_tokens 65,536, max_model_len 69,632; per-problem seed 20261030 + problem index, shared across configurations. One user message containing the unchanged EvalPlus problem prompt and the identical code-block instruction, verified by SHA-256 for every problem. Chat template enable_thinking=true and medium_effort=false. Only the requested low/high settings supply reasoning_effort.",
              "- Only text following the final </think> is eligible. Missing </think> fails. Extract the first Python fenced block (or the answer text if no fence), then apply EvalPlus's sanitizer. Use its check_correctness, canonical ground truths, special oracles, default 1-second minimum test limit and 4× reference-time factor, fast_check=true. Plus pass requires both base and extra tests to pass.",
              "- Inference ran on Modal H100s with vLLM PR #60149, commit ce9dd0ed884b7cb4014396c1aa715a8b7c9fc4a4, recurrent_kv_cache_mode=per_loop. A separate CPU container per config performed EvalPlus execution inside Modal with networking blocked. No model-generated code was executed on the Mac.",
              f"- Run: `{args.run}`. Each configuration directory contains `results.jsonl` (problem ID, sample index, tokens, finish reason, base/plus correctness), `generations.jsonl`, `sanitized.jsonl`, `prompts.jsonl`, `manifest.json`, and `summary.json`. Aggregate machine-readable scores and CIs are in `summary.json` alongside this report.",
              "- These benchmarks measure sampled function synthesis with this extraction protocol. They do not establish a causal effect of distillation independent of weight blending, reasoning length, or checkpoint selection, and do not measure repository-level coding ability.", ""]
    (ROOT / "REPORT.md").write_text("\n".join(lines))
    print("\n".join(lines[:12]))


if __name__ == "__main__":
    main()