"""Pipeline helpers for job 186 (two-layer canvas, runs/CSEL/TWOLAYER_PREREG.md). CPU only, no model code.

  gates OUT A-B CTL=DIR ARM=DIR ...     pre-judge drift gates of every arm vs the control (170's FIXED raw candidates),
                                        seeds A..B; writes OUT/gates.json; prints the arms that pass (exit 0 always)
  pools OUT SET A-B CTL_DIR ARM=DIR ... judge inputs for one pool set (control + <= 3 arms), seeds A..B, per seed and
                                        topic half (25/24): OUT/t{SET}_{name}_s{seed}h{h}.jsonl, RAW text (no ๆ write-back),
                                        a (topic, seed) kept only when the control AND every arm of the set pass form
                                        (equal counts per arm, so build_pool's balance assert holds)
  score OUT A-B SETSPEC ...             after judging runs/TL{SET}_{CLEAN,ALIGN}{seed}{h}/run1.json; SETSPEC =
                                        SET:arm1,arm2,... ; primary = judged clean rate, arm - control on the paired set,
                                        topic-clustered 95% CI, exact McNemar p, Holm across the arms of each MODEL
                                        (h_* and p_*), decision per the prereg; void pool (anchors < 70% clean) -> exit 3
  test                                  CPU tests of the pure parts
"""
import os, sys, json, math, glob, collections
sys.path.insert(0, "."); sys.path.insert(0, "scripts")


def rows(p): return [json.loads(l) for l in open(p, encoding="utf-8") if l.strip()]


def wrows(p, R):
    with open(p + ".tmp", "w", encoding="utf-8") as f:
        for r in R: f.write(json.dumps(r, ensure_ascii=False) + "\n")
    os.replace(p + ".tmp", p)


def halves(ids):   # = fixbase_pipe.halves / fixon_raw.halves
    ids = sorted(ids); h = (len(ids) + 1) // 2
    return [ids[:h], ids[h:]]


def seeds(spec):
    a, b = map(int, spec.split("-")); return range(a, b + 1)


# ---------------------------------------------------------------- gates

WF = None


def word_logfreq():
    """= fixtrain_gates.word_logfreq (copied: that module imports canvas_bot)"""
    global WF
    if WF is None:
        from pythainlp.tokenize import word_tokenize
        c = collections.Counter()
        for l in open("data/modern_bot_train.jsonl", encoding="utf-8"):
            for ln in json.loads(l)["poem"].split("\n"): c.update(w for w in word_tokenize(ln.replace(" ", ""), engine="newmm") if w.strip())
        WF = (c, sum(c.values()))
    return WF


def stats(R):
    """= fixtrain_gates.stats (syllables by the v1 segmenter, as every earlier gate) + cross-line content repeats"""
    from pythainlp.tokenize import word_tokenize
    from syl_align import syls_aligned as SA
    from klon8 import prosody as P
    c_, tot_ = word_logfreq()
    finals = collections.Counter(); bi = set(); nbi = 0; stut = 0; lf = []; rep = 0
    for r in R:
        seen_lines = []
        for ln in r["poem"].split("\n"):
            sy = SA(ln, 12)
            if sy: finals[sy[-1]] += 1
            pairs = list(zip(sy, sy[1:])); nbi += len(pairs); bi.update(pairs); stut += sum(a == b for a, b in pairs)
            ws = [w for w in word_tokenize(ln, engine="newmm") if w.strip()]
            if ws: lf.append(math.log((c_[ws[-1]] + 1) / (tot_ + len(c_))))
            content = {w for w in ws if len(P.g2p_line(w)) >= 2}          # multi-syllable words repeated across lines
            rep += sum(any(w in s for s in seen_lines) for w in content); seen_lines.append(content)
    tot = sum(finals.values())
    return {"n": len(R), "form": 100 * sum(bool(r["all_pass"]) for r in R) / len(R),
            "final_bits": -sum(c / tot * math.log2(c / tot) for c in finals.values()),
            "distinct2": len(bi) / max(nbi, 1), "stutters": stut / len(R), "repeats": rep / len(R),
            "final_word_logfreq": sum(lf) / max(len(lf), 1),
            "final_word_unseen": sum(x < math.log(1.5 / (tot_ + len(c_))) for x in lf) / max(len(lf), 1)}


def gate_fail(c, f):
    """pre-registered (TWOLAYER_PREREG.md): form >= 98%; stutters <= CTL + 0.10/poem; cross-line repeats <= CTL + 0.5;
    distinct-2 >= CTL - 0.05; line-final syllable entropy >= CTL - 0.3 bits; line-final unseen-word share <= CTL + 0.02
    (the open vocabulary must not fill line ends with strange words); line-final word log-frequency within 0.3 of CTL"""
    dead = []
    if f["form"] < 98: dead.append("form < 98%")
    if f["stutters"] > c["stutters"] + 0.10: dead.append("stutters")
    if f["repeats"] > c["repeats"] + 0.5: dead.append("cross-line repeats")
    if f["distinct2"] < c["distinct2"] - 0.05: dead.append("distinct-2")
    if f["final_bits"] < c["final_bits"] - 0.3: dead.append("rhyme-word variety")
    if f["final_word_unseen"] > c["final_word_unseen"] + 0.02: dead.append("line-final unseen words")
    if abs(f["final_word_logfreq"] - c["final_word_logfreq"]) > 0.3: dead.append("line-final word frequency")
    return dead


def gates(out, spec, *arms):
    A = [tuple(x.split("=", 1)) for x in arms]
    assert A[0][0] == "CTL", "first arm must be CTL=DIR"
    S = {n: stats([r for s in seeds(spec) for r in rows(f"{d}/c{s}.jsonl")]) for n, d in A}
    res = {"arms": S, "dead": {}}
    for n, s in S.items():
        print(f"  {n:10s} n {s['n']}, form {s['form']:.1f}%, final bits {s['final_bits']:.2f}, distinct-2 {s['distinct2']:.3f}, "
              f"stutters {s['stutters']:.3f}, repeats {s['repeats']:.2f}, final-word log-freq {s['final_word_logfreq']:.2f}, unseen {s['final_word_unseen']:.3f}")
        if n != "CTL":
            res["dead"][n] = gate_fail(S["CTL"], s)
            print(f"  {n}: " + ("KILLED before judging: " + ", ".join(res["dead"][n]) if res["dead"][n] else "pre-judge gates PASS"))
    os.makedirs(out, exist_ok=True); json.dump(res, open(f"{out}/gates.json", "w"), indent=1)
    print("PASSING " + " ".join(n for n, d in res["dead"].items() if not d))


# ---------------------------------------------------------------- pools

def pools(out, set_name, spec, ctl, *arms):
    A = [("fixbase", ctl)] + [tuple(x.split("=", 1)) for x in arms]
    assert len({n for n, _ in A}) == len(A) and 2 <= len(A) <= 4, "control + 1..3 arms per pool, unique names"
    os.makedirs(out, exist_ok=True); n = dropped = 0
    for s in seeds(spec):
        R = {name: {r["id"]: r for r in rows(f"{d}/c{s}.jsonl")} for name, d in A}
        ids = sorted(R["fixbase"])
        assert all(set(R[k]) == set(ids) for k in R), f"seed {s}: arms cover different topics"
        assert all(R[k][i]["prompt"] == R["fixbase"][i]["prompt"] for k in R for i in ids), f"seed {s}: prompts differ"
        for h, part in enumerate(halves(ids)):
            keep = [i for i in part if all(R[k][i]["all_pass"] for k in R)]
            dropped += len(part) - len(keep)
            for name, _ in A:
                J = [{"id": f"{i}_s{s}", "prompt": R[name][i]["prompt"], "poem": R[name][i]["poem"]} for i in keep]
                bad = [r["id"] for r in J if len(r["poem"].split("\n")) != 4]
                assert not bad, f"{name} seed {s}: poems without 4 lines: {bad}"
                wrows(f"{out}/t{set_name}_{name}_s{s}h{h}.jsonl", J); n += len(J)
    print(f"  pool set {set_name}: arms {[x for x, _ in A]}, seeds {spec}, {n} poems; (topic, seed) dropped for form in any arm: {dropped}")


# ---------------------------------------------------------------- score

def mcnemar_p(up, down):
    """exact two-sided sign test on the discordant pairs"""
    n = up + down
    if n == 0: return 1.0
    k = min(up, down)
    return min(1.0, 2 * sum(math.comb(n, i) for i in range(k + 1)) / 2 ** n)


def holm(ps):
    """Holm step-down adjusted p-values, {name: p}"""
    order = sorted(ps, key=ps.get); m = len(ps); out = {}; run = 0.0
    for r, k in enumerate(order):
        run = max(run, min(1.0, (m - r) * ps[k])); out[k] = run
    return out


def decide(clean, aligned, p_holm):
    """pre-registered: CONTINUE if clean >= +5pp, Holm-adjusted p < 0.05 and aligned >= -3pp; KILL if clean < +2pp;
    otherwise GREY"""
    if clean["diff"] >= 5 and p_holm < 0.05 and aligned["diff"] >= -3: return "CONTINUE"
    if clean["diff"] < 2: return "KILL"
    return "GREY"


FAMILIES = {m: [f"{m}_{u}" for u in ("syllable", "word", "wak", "bot")] for m in ("h", "p")}


def score(out, spec, *setspecs, root_pools="runs"):
    import fixon_raw, fixbase_pipe
    a, b = map(int, spec.split("-"))
    R = {"sets": {}, "arms": {}}; void = {}
    for ss in setspecs:
        set_name, arms = ss.split(":"); arms = arms.split(",")
        lab, anchors = fixon_raw.labels(f"TL{set_name}", a, b, root_pools)
        void.update({d: v for d, v in anchors.items() if v < 0.70}); R["sets"][set_name] = anchors
        for arm in arms:
            st = {}
            for kind in ("clean", "aligned"):
                c = {k[1:]: v for k, v in lab[kind].items() if k[0] == "fixbase"}
                t = {k[1:]: v for k, v in lab[kind].items() if k[0] == arm}
                assert c and set(c) == set(t), f"set {set_name}: {arm} and the control were judged on different poems"
                st[kind] = fixbase_pipe.paired_stats(c, t)
            st["clean"]["p"] = mcnemar_p(st["clean"]["up"], st["clean"]["down"]); st["set"] = set_name
            R["arms"][arm] = st
    # Holm families FIXED in advance (prereg): the 4 units of each model; an arm not judged (gate kill, set C skipped)
    # enters its family with p = 1, so skipping a set can never make the others' correction weaker
    for fam in FAMILIES:
        ps = {arm: (R["arms"][arm]["clean"]["p"] if arm in R["arms"] else 1.0) for arm in FAMILIES[fam]}
        for arm, ph in holm(ps).items():
            if arm not in R["arms"]: continue
            st = R["arms"][arm]; st["clean"]["p_holm"] = ph
            st["decision"] = decide(st["clean"], st["aligned"], ph)
            c, al = st["clean"], st["aligned"]
            print(f"  {arm:12s} clean control {c['a']:5.1f}% arm {c['b']:5.1f}% diff {c['diff']:+.1f}pp ({c['up']} up / {c['down']} down, "
                  f"n={c['n']}) CI [{c['ci'][0]:+.1f}, {c['ci'][1]:+.1f}] p {c['p']:.4f} Holm {ph:.4f}; aligned {al['diff']:+.1f}pp -> {st['decision']}")
    os.makedirs(out, exist_ok=True); R["void_pools"] = sorted(void)
    json.dump(R, open(f"{out}/RESULT.json", "w"), indent=1)
    if void:
        print(f"  VOID pools (anchors < 70% clean): {void} - re-judge those pools before any decision"); sys.exit(3)


# ---------------------------------------------------------------- tests

def _test():
    import tempfile
    t = tempfile.mkdtemp()
    poem = "ฝนตกลงมาที่บ้านเกิด\nลม\nพัด\nมา"
    for d, bad in (("ctl", {3}), ("h_word", {1}), ("h_bot", set())):
        os.makedirs(f"{t}/{d}")
        wrows(f"{t}/{d}/c0.jsonl", [{"id": f"T{i:02d}", "prompt": "ฝน", "poem": poem, "all_pass": i not in bad} for i in range(5)])
    pools(f"{t}/p", "A", "0-0", f"{t}/ctl", f"h_word={t}/h_word", f"h_bot={t}/h_bot")
    # kept only where every arm passes form: T00,T02 in half 0 (T01 fails in h_word), T04 in half 1 (T03 fails in ctl)
    for name in ("fixbase", "h_word", "h_bot"):
        assert [r["id"] for r in rows(f"{t}/p/tA_{name}_s0h0.jsonl")] == ["T00_s0", "T02_s0"], name
        assert [r["id"] for r in rows(f"{t}/p/tA_{name}_s0h1.jsonl")] == ["T04_s0"], name
    try: pools(f"{t}/p", "B", "0-0", f"{t}/ctl", *[f"a{i}={t}/h_bot" for i in range(4)]); raise AssertionError("5 arms accepted")
    except AssertionError as e: assert "control + 1..3" in str(e)
    assert abs(mcnemar_p(8, 0) - 2 / 256) < 1e-12 and mcnemar_p(3, 3) == 1.0
    hp = holm({"a": 0.01, "b": 0.04, "c": 0.03}); assert abs(hp["a"] - 0.03) < 1e-12 and abs(hp["c"] - 0.06) < 1e-12 and abs(hp["b"] - 0.06) < 1e-12, hp
    assert decide({"diff": 6}, {"diff": 0}, 0.01) == "CONTINUE" and decide({"diff": 6}, {"diff": 0}, 0.2) == "GREY"
    assert decide({"diff": 1}, {"diff": 0}, 0.5) == "KILL" and decide({"diff": 6}, {"diff": -4}, 0.01) == "GREY"
    # scoring on fake judged pools: control fails T00 (clean), arms clean everywhere
    for inst in ("CLEAN", "ALIGN"):
        for h in (0, 1):
            d = f"{t}/runs/TLA_{inst}0{h}"; os.makedirs(d); km, res = {}, []
            for name in ("fixbase", "h_word", "h_bot"):
                for r in rows(f"{t}/p/tA_{name}_s0h{h}.jsonl"):
                    k = f"{name}{r['id']}"; km[k] = {"src": name, "id": r["id"]}
                    badj = name == "fixbase" and r["id"] == "T00_s0"
                    res.append({"item": k, "errors": [{"severity": "major"}] if badj else [], "align": 2 if badj else 5})
            for j in range(3):
                km[f"h{j}"] = {"src": "human", "id": f"H{j}"}; res.append({"item": f"h{j}", "errors": [], "align": 5})
            json.dump(km, open(f"{d}/keymap.json", "w")); json.dump(res, open(f"{d}/run1.json", "w"))
    score(f"{t}/o", "0-0", "A:h_word,h_bot", root_pools=f"{t}/runs")
    Rr = json.load(open(f"{t}/o/RESULT.json"))
    c = Rr["arms"]["h_word"]["clean"]
    assert c["n"] == 3 and c["up"] == 1 and c["down"] == 0 and abs(c["diff"] - 100 / 3) < 1e-9, c
    assert abs(c["p_holm"] - 1.0) < 1e-12 and Rr["arms"]["h_word"]["decision"] == "GREY", Rr["arms"]["h_word"]
    # fixed family: 2 judged h arms with p = 0.01 each and 2 unjudged (p = 1): Holm = 4 x 0.01 and 3 x 0.01 -> 0.04
    hp = holm({"h_syllable": 1.0, "h_word": 0.01, "h_wak": 0.01, "h_bot": 1.0}); assert abs(hp["h_word"] - 0.04) < 1e-12, hp
    # gates: a clean arm passes, a stuttering arm fails
    base = {"n": 10, "form": 99, "final_bits": 7, "distinct2": 0.6, "stutters": 0.05, "repeats": 1.0, "final_word_logfreq": -9, "final_word_unseen": 0.01}
    assert gate_fail(base, dict(base)) == [] and gate_fail(base, dict(base, stutters=0.2, form=97)) == ["form < 98%", "stutters"]
    print("twolayer_pipe: tests pass")


if __name__ == "__main__":
    cmd, args = sys.argv[1], sys.argv[2:]
    {"gates": gates, "pools": pools, "score": score, "test": _test}[cmd](*args)