"""Paired comparison of two graded runs on the same items.

For every task: accuracy of A and B, difference, 95% paired-bootstrap CI of the difference (resampling items,
samples of one item averaged), and a two-sided paired permutation p-value (random sign flips of per-item deltas).

usage: compare.py <tagA> <tagB> [--json out.json]
"""
import json
import os
import random
import sys
from collections import defaultdict

ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")


def load(tag):
    per = defaultdict(lambda: defaultdict(list))
    meta = defaultdict(lambda: defaultdict(list))
    for l in open(os.path.join(ROOT, "results", tag, "graded.jsonl"), encoding="utf-8"):
        g = json.loads(l)
        per[g["task"]][g["id"]].append(float(g["correct"]))
        meta[g["task"]]["tok"].append(g["n_tokens"])
        meta[g["task"]]["trunc"].append(float(g["truncated"]))
    return per, meta


def main():
    a_tag, b_tag = sys.argv[1], sys.argv[2]
    A, Am = load(a_tag)
    B, Bm = load(b_tag)
    rng = random.Random(0)
    out = {}
    print(f"{'task':10s} {'n':>5s} {a_tag[:12]:>12s} {b_tag[:12]:>12s} {'delta':>7s} {'95% CI':>17s} {'p':>7s}"
          f" {'tokA':>6s} {'tokB':>6s} {'trA%':>5s} {'trB%':>5s}")
    macro = []
    for t in sorted(set(A) & set(B)):
        ids = sorted(set(A[t]) & set(B[t]))
        a = [sum(A[t][i]) / len(A[t][i]) for i in ids]
        b = [sum(B[t][i]) / len(B[t][i]) for i in ids]
        d = [y - x for x, y in zip(a, b)]
        n = len(ids)
        mean_d = sum(d) / n
        boots = sorted(sum(d[rng.randrange(n)] for _ in range(n)) / n for _ in range(2000))
        lo, hi = boots[50], boots[1949]
        obs = abs(sum(d))
        perm = sum(abs(sum(x if rng.random() < 0.5 else -x for x in d)) >= obs - 1e-12 for _ in range(4000))
        p = (perm + 1) / 4001
        acc_a, acc_b = 100 * sum(a) / n, 100 * sum(b) / n
        tok_a = sum(Am[t]["tok"]) / len(Am[t]["tok"])
        tok_b = sum(Bm[t]["tok"]) / len(Bm[t]["tok"])
        tr_a = 100 * sum(Am[t]["trunc"]) / len(Am[t]["trunc"])
        tr_b = 100 * sum(Bm[t]["trunc"]) / len(Bm[t]["trunc"])
        print(f"{t:10s} {n:5d} {acc_a:12.2f} {acc_b:12.2f} {100 * mean_d:+7.2f} [{100 * lo:+6.2f},{100 * hi:+6.2f}]"
              f" {p:7.4f} {tok_a:6.0f} {tok_b:6.0f} {tr_a:5.1f} {tr_b:5.1f}")
        out[t] = {"n": n, "acc_a": acc_a, "acc_b": acc_b, "delta": 100 * mean_d, "ci": [100 * lo, 100 * hi], "p": p,
                  "tok_a": tok_a, "tok_b": tok_b, "trunc_a": tr_a, "trunc_b": tr_b}
        macro.append((acc_a, acc_b))
    ma = sum(x for x, _ in macro) / len(macro)
    mb = sum(y for _, y in macro) / len(macro)
    print(f"{'MACRO':10s} {'':5s} {ma:12.2f} {mb:12.2f} {mb - ma:+7.2f}")
    if "--json" in sys.argv:
        json.dump(out, open(sys.argv[sys.argv.index("--json") + 1], "w"), indent=1)


if __name__ == "__main__":
    main()