"""Score an eval run's traces.jsonl: per-task reward, mean, Wilson ci95."""
import json
import math
import sys


def wilson(k, n, z=1.96):
    if n == 0:
        return (0.0, 0.0, 0.0)
    p = k / n
    denom = 1 + z * z / n
    center = (p + z * z / (2 * n)) / denom
    half = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / denom
    return (p, max(0.0, center - half), min(1.0, center + half))


def score(path):
    rewards = []
    with open(path) as f:
        for line in f:
            t = json.loads(line)
            for tr in t.get("traces", []):
                r = tr.get("rewards")
                if r is None:
                    continue
                if isinstance(r, dict):
                    # {'solved': {'score': 1.0, 'weight': 1.0}} or {'name': 0.5}
                    vals = []
                    for v in r.values():
                        if isinstance(v, dict) and isinstance(v.get("score"), (int, float)):
                            vals.append(v["score"] * v.get("weight", 1.0))
                        elif isinstance(v, (int, float)):
                            vals.append(v)
                    rewards.append(sum(vals) if vals else 0.0)
                elif isinstance(r, (int, float)):
                    rewards.append(float(r))
                elif isinstance(r, list):
                    rewards.append(sum(r) / len(r) if r else 0.0)
    n = len(rewards)
    k = sum(1 for r in rewards if r > 0.5)
    p, lo, hi = wilson(k, n)
    print(f"{path}")
    print(f"  episodes: {n}, solved: {k}, score: {p:.3f}, ci95: [{lo:.3f}, {hi:.3f}]")
    return p, lo, hi


if __name__ == "__main__":
    for p in sys.argv[1:]:
        score(p)