"""raw_*.jsonl を集計する。使い方: python3 analyze.py raw_haiku.jsonl [raw_sonnet.jsonl ...]"""
import json, sys, statistics as st
from collections import defaultdict
from materials import row_checks

# 呼び出し列のうち、同時に走る段（クリティカルパス計算用）
def stages(cond, n):
    if cond in ("E1", "R1"): return [n]
    if cond in ("E2", "R2"): return [1, 1]
    if cond in ("E7", "R5"): return [1, 1, 1]
    if cond in ("E3", "E5"): return [3]
    if cond in ("E4", "E6"): return [3, 1]
    if cond == "R3": return [n - 1, 1]
    if cond == "R4": return [n - 2, 1, 1]

def tokens(c): return c["input"] + c["cache_read"] + c["cache_write"] + c["output"]

def summarize(path):
    rows = [json.loads(l) for l in open(path)]
    by = defaultdict(list)
    for r in rows: by[r["cond"]].append(r)
    out = {}
    for cond, rs in sorted(by.items()):
        checks = [row_checks(r) for r in rs]
        scores = [sum(c.values()) / len(c) for c in checks]
        full = [all(c.values()) for c in checks]
        tot = [sum(tokens(c) for c in r["calls"]) for r in rs]
        outt = [sum(c["output"] for c in r["calls"]) for r in rs]
        cost = [sum(c["cost_usd"] for c in r["calls"]) for r in rs]
        cp = []
        for r in rs:
            ms, i = 0, 0
            for k in stages(cond, len(r["calls"])):
                ms += max(c["ms"] for c in r["calls"][i:i + k]); i += k
            cp.append(ms / 1000)
        errs = sum(1 for r in rs for c in r["calls"] if c.get("error"))
        fails = defaultdict(int)
        for c in checks:
            for k, v in c.items():
                if not v: fails[k] += 1
        out[cond] = dict(n=len(rs), score=st.mean(scores), full=sum(full), tokens=st.mean(tot),
                         output=st.mean(outt), cost=st.mean(cost), calls=st.mean(len(r["calls"]) for r in rs),
                         path_s=st.mean(cp), errors=errs, fails=dict(sorted(fails.items(), key=lambda x: -x[1])),
                         scores=sorted(round(s, 3) for s in scores))
    return out

if __name__ == "__main__":
    for p in sys.argv[1:]:
        print("==", p)
        for c, d in summarize(p).items():
            print(c, {k: (round(v, 3) if isinstance(v, float) else v) for k, v in d.items()})