"""Compares the self-aggregation module with the weighted vote over the same K attempts.
Split problems (attempts disagree) get the aggregator's answer, unanimous ones keep the shared answer.
usage: agg_eval.py data/agg/m180.jsonl agg_m180_dm   (the second argument is a graded results tag)"""
import collections
import json
import os
import random
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from grade import grade_math  # noqa: E402

items = {json.loads(l)["id"]: json.loads(l) for l in open(sys.argv[1], encoding="utf-8")}
agg = collections.defaultdict(list)
for l in open(os.path.join("results", sys.argv[2], "graded.jsonl"), encoding="utf-8"):
    g = json.loads(l)
    agg[g["id"]].append(g["correct"])


def vote_ok(it):
    v = it["vote"]
    if v is None:
        return False
    return v == it["answer"] if it["task"] == "mmlupro" else bool(grade_math("\\boxed{" + v + "}", it["answer"]))


vote = {i: vote_ok(it) for i, it in items.items()}
split = [i for i, it in items.items() if not it["unanimous"] and i in agg]
diff = [sum(agg[i]) / len(agg[i]) - vote[i] for i in split]
rng = random.Random(0)
boots = sorted(sum(rng.choice(diff) for _ in diff) / len(diff) for _ in range(2000)) if diff else [0] * 2000
n = len(items)
hyb = sum(vote[i] for i, it in items.items() if it["unanimous"] or i not in agg) + sum(sum(agg[i]) / len(agg[i]) for i in split)
print(f"{sys.argv[2]}: {n} problems, {len(split)} split, samples/problem {sum(len(agg[i]) for i in split) / max(len(split), 1):.1f}")
print(f"  split problems: vote {100 * sum(vote[i] for i in split) / max(len(split), 1):.1f} -> aggregator "
      f"{100 * sum(sum(agg[i]) / len(agg[i]) for i in split) / max(len(split), 1):.1f}")
print(f"  all problems:   vote {100 * sum(vote.values()) / n:.1f} -> vote+aggregator {100 * hyb / n:.1f}  "
      f"d={100 * sum(diff) / n:+.1f} [{100 * boots[50] * len(split) / n:+.1f}, {100 * boots[1949] * len(split) / n:+.1f}]")