"""Production pipeline pieces for the bug-fixed base comparison (runs/CSEL/FIXED_BASE_PREREG.md, job 170).

Per arm (SPLIT = runs/CANVAS_SPLIT, FIXED = runs/CANVAS_FIXBASE) and seed block b (seeds 64b .. 64b+63), in D:
  merge D NT SEED0 [NS]  D/cands/c{s}.jsonl, s = SEED0..SEED0+NS-1 (NS = 64) -> D/cands.jsonl (flag and coverage asserts)
  picks D              D/typhoon_top8.jsonl rank 0 -> D/pick.jsonl (id, prompt, poem)
  locmap D             D/pick_scores.jsonl -> D/line.json (argmin copy-masked canvas line mean; vetomore.locmap)
  variants D           D/pick.jsonl + form-passing D/ref/r*_*.jsonl -> D/rcands.jsonl (105's >= 8 per pick floor)
  accept D             vetomore.accept verbatim (+0.5 margin, topic-keep) -> D/final.jsonl
Across arms:
  gates ROOT B         pre-registered drift gates on block B, FIXED vs SPLIT; writes ROOT/jp_{split,fixed}_b{B}.jsonl
                       (judge inputs, id = topic_s{B}) and ROOT/GATES_b{B}.json; exit 2 = killed before judging
  score ROOT NB        after judging: paired form / clean / aligned / all-three, topic-clustered bootstrap, decision
  test                 CPU unit tests of the pure functions
"""
import os, re, sys, json, glob, random, collections
sys.path.insert(0, "."); sys.path.insert(0, "scripts")

DARGS_FLAGS = {"topic_zero": "row", "boost_mode": "flat"}


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


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


def merge(d, nt, s0, ns=64):
    nt, s0, ns = int(nt), int(s0), int(ns)
    rs = []
    for s in range(s0, s0 + ns):
        f = f"{d}/cands/c{s}.jsonl"; x = rows(f)
        assert len(x) == nt, f"{f}: {len(x)} rows, want {nt}"
        rs += [dict(r, seed=s) for r in x]
    bad = [r["id"] for r in rs if any(r.get(k) != v for k, v in DARGS_FLAGS.items()) or r.get("topic_boost") not in (1.5, 3.0)]
    assert not bad, f"{len(bad)} rows decoded with the wrong flags"
    tb = collections.defaultdict(set)
    for r in rs: tb[r["id"]].add(r["topic_boost"])
    assert len(tb) == nt and all(len(v) == 1 for v in tb.values()), "topic coverage / topic_boost differs across seeds"
    wjsonl(f"{d}/cands.jsonl", rs)
    print(f"  {d}: {len(rs)} candidates for {nt} topics, seeds {s0}-{s0 + ns - 1}; form {100 * sum(r['all_pass'] for r in rs) / len(rs):.1f}%")


def picks(d):
    top = [r for r in rows(f"{d}/typhoon_top8.jsonl") if r["sel_rank"] == 0]
    assert len({r["topic_id"] for r in top}) == len(top), "two rank-0 picks for one topic"
    wjsonl(f"{d}/pick.jsonl", [{"id": r["topic_id"], "prompt": r["prompt"], "poem": r["poem"]} for r in sorted(top, key=lambda r: r["topic_id"])])
    print(f"  {d}: {len(top)} Typhoon picks")


def locmap(d):
    P_, S = rows(f"{d}/pick.jsonl"), rows(f"{d}/pick_scores.jsonl")
    assert [p["id"] for p in P_] == [s["id"] for s in S] and all(p["poem"] == s["poem"] for p, s in zip(P_, S)), "scores do not belong to the picks"
    m = {s["id"]: min(range(4), key=lambda i: s["score"]["line_mean"][i]) + 1 for s in S if s["score"]}
    miss = [s["id"] for s in S if not s["score"]]
    print(f"  {d}: {len(m)} picks located, {len(miss)} without a score (kept unchanged, as round 1): {miss}")
    assert len(miss) <= max(3, len(S) // 20), "too many picks fail to round-trip"
    json.dump(m, open(f"{d}/line.json", "w"), indent=0)


def variants(d):
    base = rows(f"{d}/pick.jsonl"); loc = json.load(open(f"{d}/line.json"))
    out = [dict(r, variant="orig") for r in base]; n = 0
    for f in sorted(glob.glob(f"{d}/ref/r*_*.jsonl")):
        for r in rows(f):
            if r.get("refill_line") and r["all_pass"]: out.append(dict(r, variant=os.path.basename(f))); n += 1
    wjsonl(f"{d}/rcands.jsonl", out)
    print(f"  {d}: {len(base)} picks + {n} form-passing variants")
    assert n >= len(loc) * 8, "fewer than half the expected variants (105's floor)"


def accept(d):
    import vetomore
    best, uns = vetomore.accept(rows(f"{d}/rcands.jsonl"), rows(f"{d}/rcands_scores.jsonl"),
                                {k: int(v) for k, v in json.load(open(f"{d}/line.json")).items()})
    from klon8 import prosody as P
    fin = [{"id": i, "prompt": best[i]["prompt"], "poem": best[i]["poem"], "variant": best[i]["variant"],
            "form": bool(P.score(best[i]["poem"], n_bot_expected=1)["all_pass"])} for i in sorted(best)]
    assert len(fin) == len(rows(f"{d}/pick.jsonl"))
    wjsonl(f"{d}/final.jsonl", fin); P.save_cache()
    print(f"  {d}: {len(fin)} final poems, {sum(r['variant'] != 'orig' for r in fin)} refilled, {len(uns)} unscored, "
          f"form {100 * sum(r['form'] for r in fin) / len(fin):.1f}%")


# ---- drift metrics (same segmenter for all arms: syl_align v1, as every earlier gate) ----
ARMS = ("SPLIT", "CTL12", "FIXED")
_REDUP = None


def redup():
    """Syllables the corpus writes with ไม้ยมก (Xๆ); excluded from the stutter count in EVERY arm, since FIXED
    writes them out (X X) internally."""
    global _REDUP
    if _REDUP is None:
        import fixbase_data
        _REDUP = fixbase_data.redup_syllables()
    return _REDUP


def stutters(p, exclude=frozenset()):
    from syl_align import syls_aligned as SA
    n = 0
    for ln in p.split("\n"):
        s = SA(ln, 12); n += sum(a == b and a not in exclude for a, b in zip(s, s[1:]))
    return n


def xrep(p):
    from syl_align import syls_aligned as SA
    from canvas_mlm import STOP_SYL
    ls = [set(x for x in SA(l, 12) if x not in STOP_SYL and len(x) > 1) for l in p.split("\n")]
    c = collections.Counter(x for s_ in ls for x in s_); return sum(v - 1 for v in c.values() if v > 1)


def distinct2(ps):
    from syl_align import syls_aligned as SA
    bg = [(a, b) for p in ps for ln in p.split("\n") for s in [SA(ln, 12)] for a, b in zip(s, s[1:])]
    return len(set(bg)) / max(len(bg), 1)


def judged_text(poem, _w=[]):
    """The text the judge sees: reduplication written back as ๆ (human anchors use ๆ), the same rule for every
    arm. Form, stutters and repeats are measured on the text the model produced (before this)."""
    import fixbase_data, syl_align_v2
    if not _w:
        from pythainlp.corpus import thai_words
        _w.append(thai_words())
    return "\n".join(fixbase_data.contract_yamok(ln, lambda t: syl_align_v2.syls_aligned(t, 12), _w[0]) for ln in poem.split("\n"))


def halves(ids):
    """Two pools per block (pinned recipe: <= 5 system poems per anchor, 20 anchors): 3 arms x 49 topics do not fit
    one pool, so topics split by sorted id; the three arms of a topic are always in the same pool."""
    ids = sorted(ids); h = (len(ids) + 1) // 2
    return [ids[:h], ids[h:]]


def gates(root, b):
    """Pre-registered (FIXED_BASE_PREREG.md), FIXED vs its comparator CTL12: FIXED form >= 98% and >= CTL12 - 2pp;
    stutters (corpus reduplications excluded) <= CTL12 + 0.10 per poem; cross-line repeats <= CTL12 + 0.50;
    distinct-2 >= CTL12 - 0.05. SPLIT is reported alongside. Then the judge inputs (ๆ written back)."""
    R = {a: {r["id"]: r for r in rows(f"{root}/{a}_b{b}/final.jsonl")} for a in ARMS}
    ids = sorted(R["FIXED"])
    assert all(set(R[a]) == set(ids) and all(R[a][i]["prompt"] == R["FIXED"][i]["prompt"] for i in ids) for a in ARMS), \
        "arms cover different topics"
    ex = redup()
    m = lambda f, a: sum(f(R[a][i]["poem"]) for i in ids) / len(ids)
    g = {"n": len(ids)}
    for k, f in (("form", lambda a: 100 * sum(R[a][i]["form"] for i in ids) / len(ids)),
                 ("stutters", lambda a: m(lambda p: stutters(p, ex), a)), ("stutters_raw", lambda a: m(stutters, a)),
                 ("xrep", lambda a: m(xrep, a)), ("distinct2", lambda a: distinct2([R[a][i]["poem"] for i in ids])),
                 ("refilled", lambda a: sum(R[a][i]["variant"] != "orig" for i in ids))):
        g[k] = {a: f(a) for a in ARMS}
    F_, C_ = "FIXED", "CTL12"
    fail = []
    if g["form"][F_] < 98 or g["form"][F_] < g["form"][C_] - 2: fail.append("form")
    if g["stutters"][F_] > g["stutters"][C_] + 0.10: fail.append("stutters")
    if g["xrep"][F_] > g["xrep"][C_] + 0.50: fail.append("cross-line repeats")
    if g["distinct2"][F_] < g["distinct2"][C_] - 0.05: fail.append("distinct-2")
    g["fail"] = fail
    for k, v in g.items():
        if isinstance(v, dict): print(f"  block {b} {k:13s} " + "  ".join(f"{a} {v[a]:.3f}" for a in ARMS))
    changed = 0
    for h, part in enumerate(halves(ids)):
        for a in ARMS:
            out = []
            for i in part:
                t = judged_text(R[a][i]["poem"]); changed += t != R[a][i]["poem"]
                out.append({"id": f"{i}_s{b}", "prompt": R[a][i]["prompt"], "poem": t})
            wjsonl(f"{root}/jp_{a.lower()}_b{b}h{h}.jsonl", out)
    g["yamok_written_back"] = changed
    json.dump(g, open(f"{root}/GATES_b{b}.json", "w"), indent=1)
    print(f"  block {b}: {changed} judged poems carry reduplication written back as ๆ")
    if fail:
        print(f"  KILLED before judging, block {b}: {fail}"); sys.exit(2)
    print(f"  block {b}: drift gates pass; judge inputs {root}/jp_*_b{b}h{{0,1}}.jsonl")


def paired_stats(a, b, seed=0, n_boot=10000):
    """a, b: {key: bool}, key = (topic, block). Difference b - a in pp, up/down, topic-clustered percentile CI."""
    keys = sorted(set(a) & set(b))
    by = collections.defaultdict(list)
    for k in keys: by[k[0]].append(k)
    T = sorted(by); rng = random.Random(seed); boots = []
    for _ in range(n_boot):
        ks = [k for t in (rng.choice(T) for _ in T) for k in by[t]]
        boots.append(100 * (sum(b[k] for k in ks) - sum(a[k] for k in ks)) / len(ks))
    boots.sort()
    d = 100 * (sum(b[k] for k in keys) - sum(a[k] for k in keys)) / len(keys)
    return {"n": len(keys), "a": 100 * sum(a[k] for k in keys) / len(keys), "b": 100 * sum(b[k] for k in keys) / len(keys),
            "diff": d, "up": sum(b[k] and not a[k] for k in keys), "down": sum(a[k] and not b[k] for k in keys),
            "ci": (boots[int(0.025 * n_boot)], boots[int(0.975 * n_boot) - 1])}


def decide(allthree):
    """Pre-registered: CONTINUE if all-three >= +5pp with the lower 95% bound > 0; GREY if >= +5pp with the CI
    including 0; otherwise STOP. Applied to the PRIMARY contrast FIXED vs CTL12 on all blocks."""
    if allthree["diff"] >= 5 and allthree["ci"][0] > 0: return "CONTINUE"
    if allthree["diff"] >= 5: return "GREY"
    return "STOP"


def score(root, nb):
    lab = {"clean": {}, "aligned": {}, "form": {}}
    for b in range(int(nb)):
        for a in ARMS:
            for r in rows(f"{root}/{a}_b{b}/final.jsonl"): lab["form"][(a.lower(), r["id"], b)] = r["form"]
        for h in (0, 1):
            for kind, pre in (("clean", "FBJ_CLEAN"), ("aligned", "FBJ_ALIGN")):
                d = f"runs/{pre}{b}{h}"
                km = json.load(open(f"{d}/keymap.json", encoding="utf-8")); res = {r["item"]: r for r in json.load(open(f"{d}/run1.json", encoding="utf-8"))}
                assert set(res) == set(km), f"{d}: judge output does not cover the pool"
                for k, v in km.items():
                    if v["src"] == "human": continue
                    r = res[k]
                    val = (not any(e.get("severity") in ("critical", "major") for e in r["errors"])) if kind == "clean" else int(r["align"]) >= 4
                    t, bb = v["id"].rsplit("_s", 1); assert int(bb) == b
                    assert (v["src"], t, b) not in lab[kind], f"{v} judged twice"
                    lab[kind][(v["src"], t, b)] = val
    for kind in ("clean", "aligned"):
        assert set(lab[kind]) == set(lab["form"]), f"{kind}: judged poems differ from the pipeline output"
    lab["all-three"] = {k: lab["form"][k] and lab["clean"][k] and lab["aligned"][k] for k in lab["form"]}
    out = {}
    for name, (A, B) in (("PRIMARY fixes (FIXED vs CTL12)", ("ctl12", "fixed")), ("schedule (CTL12 vs SPLIT)", ("split", "ctl12")),
                         ("whole (FIXED vs SPLIT)", ("split", "fixed"))):
        out[name] = {}
        print(f"  {name}")
        for kind, L in lab.items():
            a = {k[1:]: v for k, v in L.items() if k[0] == A}; bb = {k[1:]: v for k, v in L.items() if k[0] == B}
            st = out[name][kind] = paired_stats(a, bb)
            print(f"    {kind:9s} {A} {st['a']:5.1f}%  {B} {st['b']:5.1f}%  diff {st['diff']:+.1f}pp  ({st['up']} up / {st['down']} down, "
                  f"n={st['n']})  95% CI [{st['ci'][0]:+.1f}, {st['ci'][1]:+.1f}]")
    out["decision"] = decide(out["PRIMARY fixes (FIXED vs CTL12)"]["all-three"])
    json.dump(out, open(f"{root}/RESULT.json", "w"), indent=1)
    print(f"  DECISION (pre-registered, primary): {out['decision']}")


def _test():
    # paired stats and the decision rule
    a = {(f"T{i}", 0): i % 2 == 0 for i in range(40)}
    b = {(f"T{i}", 0): (i % 2 == 0) or i < 20 for i in range(40)}
    s = paired_stats(a, b, n_boot=2000)
    assert s["n"] == 40 and s["up"] == 10 and s["down"] == 0 and abs(s["diff"] - 25) < 1e-9 and s["ci"][0] > 0
    assert decide(s) == "CONTINUE"
    assert decide({"diff": 6.0, "ci": (-2.0, 14.0)}) == "GREY" and decide({"diff": 4.9, "ci": (1.0, 9.0)}) == "STOP"
    same = paired_stats(a, a, n_boot=500); assert same["diff"] == 0 and same["ci"] == (0, 0)
    # clustering: two blocks of the same topic resample together
    a2 = {**a, **{(t, 1): v for (t, _), v in a.items()}}; b2 = {**b, **{(t, 1): v for (t, _), v in b.items()}}
    s2 = paired_stats(a2, b2, n_boot=2000); assert s2["n"] == 80 and abs(s2["diff"] - 25) < 1e-9
    # stutters: written-out reduplication is excluded only when asked
    ex = redup()
    assert ex, "no reduplication found in the training poems"
    w = next(iter(sorted(ex)))
    p = f"{w}{w}ไปแล้ว\nก\nข\nค"
    assert stutters(p) >= 1 and stutters(p, ex) == stutters(p) - 1
    assert xrep("ความรักไป\nความรักมา\nก\nข") >= 1
    # pools: every topic in exactly one half, halves differ by at most one
    h = halves([f"T{i:03d}" for i in range(49)]); assert [len(x) for x in h] == [25, 24] and not set(h[0]) & set(h[1])
    # judged text: reduplication written back, nothing else changed, round trip to the model's cells
    import fixbase_data, syl_align_v2
    assert judged_text("ค่อยช้าช้าไปดีดี\nนานาชาติ") == "ค่อยช้าๆไปดีๆ\nนานาชาติ"
    assert judged_text("สายลมพัดผ่าน") == "สายลมพัดผ่าน"
    t = "ทุกทุกวันคิดถึง"; seg = lambda x: syl_align_v2.syls_aligned(x, 12)
    assert fixbase_data.expand_yamok(seg(judged_text(t))) == seg(t)
    print(f"fixbase_pipe: tests pass ({len(ex)} corpus reduplication syllables)")


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