#!/usr/bin/env python3
import argparse
import collections
import json
import math


CODE_SUBSTRINGS = [
    "computer",
    "programming",
    "operating_system",
    "architecture",
    "network",
    "security",
]

MATH_SUBSTRINGS = [
    "math",
    "physics",
    "chemistry",
    "biology",
    "statistics",
    "probability",
    "astronomy",
    "anatomy",
    "medical",
    "electrical",
    "engineering",
    "machine_learning",
]


def load_jsonl(path):
    rows = {}
    with open(path, encoding="utf-8") as f:
        for line in f:
            if line.strip():
                row = json.loads(line)
                rows[row["id"]] = row
    return rows


def bucket(row):
    category = row.get("category")
    subset = str(row.get("subset") or "").lower()
    meta = row.get("metadata") or {}
    domain = str(meta.get("high_level_domain") or "").lower()
    subdomain = str(meta.get("subdomain") or "").lower()
    text = " ".join([subset, domain, subdomain])

    if any(x in text for x in CODE_SUBSTRINGS):
        return "coding_or_cs"
    if category == "gpqa" or any(x in text for x in MATH_SUBSTRINGS):
        return "math_or_science_reasoning"
    if category == "ceval":
        return "chinese_exam"
    if category == "mmlu":
        return "english_general_mmlu"
    return "other"


def mean(xs):
    xs = [x for x in xs if x is not None]
    return sum(xs) / len(xs) if xs else None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--benchmark", required=True)
    ap.add_argument("--eval", required=True)
    args = ap.parse_args()

    bench = load_jsonl(args.benchmark)
    eval_rows = load_jsonl(args.eval)
    groups = collections.defaultdict(list)
    for item_id, ev in eval_rows.items():
        b = bench[item_id]
        groups[bucket(b)].append((b, ev))

    out = {}
    for name, rows in sorted(groups.items()):
        out[name] = {
            "n": len(rows),
            "mcq_acc_avg_norm": mean([ev.get("mcq_correct_avg_norm") for _, ev in rows]),
            "mcq_acc_sum": mean([ev.get("mcq_correct_sum") for _, ev in rows]),
            "ppl_token_mean": mean([ev.get("ppl", {}).get("ppl") for _, ev in rows if ev.get("ppl")]),
            "nll_per_token_mean": mean([ev.get("ppl", {}).get("nll_per_token") for _, ev in rows if ev.get("ppl")]),
            "bits_per_byte_mean": mean([ev.get("ppl", {}).get("bits_per_byte") for _, ev in rows if ev.get("ppl")]),
            "subsets_top": dict(collections.Counter(str(b.get("subset")) for b, _ in rows).most_common(20)),
        }
    print(json.dumps(out, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()