"""単体の「初稿 → 見直し → 見直し」で、各段の後に完全合格だった数を数える（調べもの課題）。
R1・R2・R5 の最初の呼び出しは同じ依頼文なので、初稿は72本まとめて見る。"""
import json, re
import materials as m
from stats import fisher

gold = m.research_gold("L")


def strip_fence(t):
    t = t.strip()
    mm = re.match(r"^```[a-z]*\n(.*?)\n```$", t, re.S)
    return mm.group(1) if mm else t


def pj(t):
    t = strip_fence(t)
    try:
        return json.loads(t[t.find("{"):t.rfind("}") + 1])
    except Exception:
        return None


def perfect(ans):
    rows = (ans or {}).get("rows") or {}
    for n in [p[0] for p in m.products("L")]:
        for f, g in gold[n].items():
            v = (rows.get(n) or {}).get(f)
            if isinstance(g, str):
                v = str(v).replace("リージョン", "").strip() if v is not None else None
            else:
                try:
                    v = int(round(float(str(v).replace(",", ""))))
                except Exception:
                    v = None
            if v != g:
                return False
    return (ans or {}).get("q1") == gold["_q1"] and (ans or {}).get("q2") == gold["_q2"]


rows = [json.loads(l) for l in open("raw_haiku.jsonl")]
stage = {0: [], 1: [], 2: []}
for r in rows:
    if r["cond"] in ("R1", "R2", "R5"):
        for i, c in enumerate(r["calls"]):
            stage[i].append(perfect(pj(c["text"])))
        if r["cond"] == "R1":
            pass
print("単体の初稿(R1,R2,R5)", sum(stage[0]), len(stage[0]))
print("見直し1回(R2,R5)", sum(stage[1]), len(stage[1]))
print("見直し2回(R5)", sum(stage[2]), len(stage[2]))
r3 = [all(r["checks"].values()) for r in rows if r["cond"] == "R3"]
print("分担(R3)", sum(r3), len(r3))
for name, a, na in [("初稿", sum(stage[0]), len(stage[0])), ("見直し1回", sum(stage[1]), len(stage[1])), ("見直し2回", sum(stage[2]), len(stage[2]))]:
    print(f"分担 vs {name}: p={fisher(sum(r3), len(r3), a, na):.4f}")