"""Error bars for the paper, from saved per-question results (eval.py). No model is run here.

For each model and each benchmark: the score and its 95% range (bootstrap: re-draw the same number of questions with
replacement 10,000 times; the middle 95% of the re-drawn scores is the range).
For each pair of arms: the difference on the SAME questions and its 95% range (paired bootstrap: the same re-drawn
questions are used for both arms, so luck in which questions were asked cancels). If the range of the difference does
not include 0, the difference is real at 95%; otherwise the two arms are tied on that benchmark.
Validation loss: the same, with 1,024-token validation sequences in place of questions.

  python bars.py reports/eval_armA_*.json reports/eval_armB_*.json ... [--ref armA_...]
Writes reports/ERROR_BARS.md and reports/error_bars.json.
"""
import argparse, json, pathlib
import numpy as np

N_BOOT = 10_000


def ci(x, rng):
    x = np.asarray(x, float)
    idx = rng.integers(0, len(x), (N_BOOT, len(x)))
    b = x[idx].mean(1)
    return float(x.mean()), float(np.percentile(b, 2.5)), float(np.percentile(b, 97.5))


def paired(x, y, rng):
    x, y = np.asarray(x, float), np.asarray(y, float)
    assert len(x) == len(y), "paired comparison needs the same questions"
    d = y - x
    idx = rng.integers(0, len(d), (N_BOOT, len(d)))
    b = d[idx].mean(1)
    return float(d.mean()), float(np.percentile(b, 2.5)), float(np.percentile(b, 97.5))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("files", nargs="+")
    ap.add_argument("--ref", default="", help="name of the reference model for differences (default: the first file)")
    a = ap.parse_args()
    rng = np.random.default_rng(0)
    runs = [json.load(open(f)) for f in a.files]
    names = [r["name"] for r in runs]
    ref = a.ref or names[0]
    R = runs[names.index(ref)]
    tasks = list(R["per_question"]) + ["val_loss"]

    def series(r, t):
        return r["val_loss_per_seq"] if t == "val_loss" else r["per_question"][t]

    out = {"scores": {}, "diffs": {}}
    L = ["# Error bars (95%, bootstrap, 10,000 re-draws)", "",
         "Benchmarks: accuracy in %, higher is better. Validation loss: nats/token, lower is better.", "",
         "| model | " + " | ".join(tasks) + " |", "|---|" + "---|" * len(tasks)]
    for r in runs:
        cells = []
        for t in tasks:
            m, lo, hi = ci(series(r, t), rng)
            out["scores"].setdefault(r["name"], {})[t] = [m, lo, hi]
            k = 1 if t == "val_loss" else 100
            cells.append(f"{m * k:.{3 if t == 'val_loss' else 1}f} ({lo * k:.{3 if t == 'val_loss' else 1}f}-{hi * k:.{3 if t == 'val_loss' else 1}f})")
        L.append(f"| {r['name']} | " + " | ".join(cells) + " |")
    L += ["", f"## Differences against {ref}, same questions (paired)", "",
          "| model minus " + ref + " | " + " | ".join(tasks) + " |", "|---|" + "---|" * len(tasks)]
    for r in runs:
        if r["name"] == ref:
            continue
        cells = []
        for t in tasks:
            x, y = series(R, t), series(r, t)
            if len(x) != len(y):
                cells.append("n/a"); continue
            d, lo, hi = paired(x, y, rng)
            out["diffs"].setdefault(r["name"], {})[t] = [d, lo, hi]
            k = 1 if t == "val_loss" else 100
            nd = 3 if t == "val_loss" else 1
            verdict = "tied" if lo <= 0 <= hi else ("better" if (d > 0) != (t == "val_loss") else "worse")
            cells.append(f"{d * k:+.{nd}f} ({lo * k:+.{nd}f} to {hi * k:+.{nd}f}) {verdict}")
        L.append(f"| {r['name']} | " + " | ".join(cells) + " |")
    L += ["", "Reading: 'tied' = the 95% range of the difference includes 0, so the two cannot be told apart on that "
          "benchmark; 'better'/'worse' = a real difference at 95%."]
    pathlib.Path("reports").mkdir(exist_ok=True)
    open("reports/ERROR_BARS.md", "w").write("\n".join(L) + "\n")
    json.dump(out, open("reports/error_bars.json", "w"), indent=1)
    print("\n".join(L))


if __name__ == "__main__":
    main()