"""Build a judged pool to the pinned recipe in docs/judge_frozen.md section 2.1.

Pool composition is part of the instrument: the same 20 human poems scored 85% perfect in one pool
and 100% in another, and an unanchored pool runs ~0.45 bad-วรรค against 0.26 anchored. So the recipe
is enforced here rather than remembered: the pinned anchor set, five system poems per anchor, and
every batch carrying the pool's own composition (10 system + 2 anchors).

Usage:
  build_pool.py OUT_DIR LABEL=FILE [LABEL=FILE ...] [--bot]
      --bot builds a 4-วรรค (one บท) pool: the 1-บท anchors and the 4-line judge prompt.
      --no-drop never drops a poem: anchors = ceil(n/5) (<= 20), ratio at most 5:1 (final confirmation only).

Writes batch_00.txt .. (frozen prompt + items) and keymap.json ({key: {src, id}}).
"""
import sys, os, json, random, hashlib

# Two lengths, two instruments. The 8-วรรค pool is the frozen v3.2 setup. The 1-บท pool is the
# 4-วรรค target: the same wording with "8 numbered lines" -> "4 numbered lines" and <1-8> -> <1-4>,
# and the same twenty anchors reduced to their first บท.
FORMS = {
    8: dict(anchors="data/pools/anchor20.jsonl",
            anchor_sha="0dd2a105b6c9b41f8999b7511f47e82dc897d094a9e7d7a03c48e2080577cd26",
            prompt="docs/judge_v3.2_prompt.txt",
            prompt_sha="9f051d5d6c2f05e7cce70770691e35fb5af8e6ed0335221a8f3425bfa5705eec"),
    # The 4-วรรค anchor is twenty poems their authors WROTE as one บท (posts that were four lines and
    # nothing else), not the first บท sliced out of an 8-วรรค poem. A slice is a fragment: its วรรคส่ง
    # is built to hand off to the next บท and its topic may only be developed later, so judging it
    # would penalise the human for our cut.
    4: dict(anchors="data/pools/anchor20_bot_native.jsonl",
            anchor_sha="34335cea2fc210e9b480d48336258d39fd1fbe73a1d8c58afae1d13cecac1e16",
            prompt="docs/judge_v4bot_prompt.txt",
            prompt_sha="60db6eedf758c081c272ff3fc7d0c01a69568d1479c6c9c5558952a66508c6e6"),
}
ALIGN_PROMPT = "docs/judge_align_bot_prompt.txt"
ALIGN_SHA = "f7c7f004e417f18576bed5e7286b03727ff9195f088cf3ad22c2c6501d868d98"
PER_ANCHOR, BATCH_SYS, BATCH_ANC = 5, 10, 2


def sha(path):
    return hashlib.sha256(open(path, "rb").read()).hexdigest()


def rows(path, nlines=8):
    out, seen = [], set()
    for l in open(path, encoding="utf-8"):
        r = json.loads(l)
        if not l.strip() or r["id"] in seen: continue
        seen.add(r["id"])
        poem = r["poem"] if isinstance(r["poem"], str) else "\n".join(r["poem"])
        if len(poem.split("\n")) != nlines: continue
        out.append({"id": r["id"], "topic": r.get("prompt") or r.get("topic") or r.get("text"), "poem": poem})
    return out


def build(out_dir, systems, nlines=8, align=False, no_drop=False):
    F = dict(FORMS[nlines])
    if align:                       # same pool recipe, the alignment instrument instead
        F["prompt"], F["prompt_sha"] = ALIGN_PROMPT, ALIGN_SHA
    assert sha(F["anchors"]) == F["anchor_sha"], f'anchor set has changed: {sha(F["anchors"])}'
    assert sha(F["prompt"]) == F["prompt_sha"], f'judge prompt has changed: {sha(F["prompt"])}'
    anchors = rows(F["anchors"], nlines)
    sysrows = [(lab, r) for lab, f in systems for r in rows(f, nlines)]
    assert sysrows, f"no {nlines}-line poems found in {[f for _, f in systems]}"
    n_anchor = min(len(anchors), max(1, len(sysrows) // PER_ANCHOR))
    if no_drop:
        # final confirmation (runs/CSEL/CONFIRM_PREREG.md): never drop a poem - round the anchor count UP, so the
        # ratio is at most 5 system poems per anchor; the pool must fit the 20 pinned anchors
        n_anchor = -(-len(sysrows) // PER_ANCHOR)
        assert n_anchor <= len(anchors), f"{len(sysrows)} poems exceed {len(anchors)} anchors x {PER_ANCHOR}"
    anchors = sorted(anchors, key=lambda r: r["id"])[:n_anchor]      # subsample by id, never at random
    assert len(sysrows) >= PER_ANCHOR, "too few system poems for even one anchor"

    # round-robin by label FIRST, then cap: capping the concatenation would take the whole quota
    # from the first system and silently drop the second arm of an A/B pool.
    by_lab = {}
    for lab, r in sysrows: by_lab.setdefault(lab, []).append(r)
    order, labs = [], sorted(by_lab)
    while any(by_lab[l] for l in labs):
        for l in labs:
            if by_lab[l]: order.append((l, by_lab[l].pop(0)))
    if no_drop: assert len(order) <= PER_ANCHOR * len(anchors)
    order = order[:PER_ANCHOR * len(anchors)]                        # hold the 5:1 ratio exactly
    if len(labs) > 1:
        c = {l: sum(1 for x, _ in order if x == l) for l in labs}
        assert max(c.values()) - min(c.values()) <= 1, f"A/B pool is unbalanced: {c}"

    nb = max(1, len(order) // BATCH_SYS)
    batches = [[] for _ in range(nb)]
    # deal each arm separately with its own offset. Dealing the interleaved list with i % nb aliases
    # when nb is even: base lands on even batches, B on odd, and every batch is single-arm again.
    for li, lab in enumerate(labs):
        mine = [(lab, r) for l, r in order if l == lab]
        for j, item in enumerate(mine): batches[(j + li) % nb].append(item)
    for i, r in enumerate(anchors): batches[i % nb].append(("human", r))
    for b, items in enumerate(batches):
        present = set(l for l, _ in items)
        assert present >= set(labs) | {"human"}, f"batch {b} is missing {set(labs) | {'human'} - present}"

    os.makedirs(out_dir, exist_ok=True)
    prompt = open(F["prompt"], encoding="utf-8").read().rstrip()
    keymap, kr = {}, random.Random(0)
    for b, items in enumerate(batches):
        assert any(l == "human" for l, _ in items), f"batch {b} carries no anchor"
        random.Random(f"{out_dir}:{b}").shuffle(items)
        lines = [prompt, ""]
        for lab, r in items:
            key = "%06x" % kr.randrange(16 ** 6)
            keymap[key] = {"src": lab, "id": r["id"]}
            lines.append(f'item {key}  topic: {r["topic"]}')
            lines += [f"{i}. {ln}" for i, ln in enumerate(r["poem"].split("\n"), 1)]
            lines.append("")
        open(os.path.join(out_dir, f"batch_{b:02d}.txt"), "w", encoding="utf-8").write("\n".join(lines))
    json.dump(keymap, open(os.path.join(out_dir, "keymap.json"), "w", encoding="utf-8"), ensure_ascii=False, indent=0)
    comp = {l: sum(1 for k, v in keymap.items() if v["src"] == l) for l in set(v["src"] for v in keymap.values())}
    print(f"{len(keymap)} poems in {nb} batches -> {out_dir}")
    print(f"  composition {comp}  (anchors {100*comp.get('human',0)/len(keymap):.0f}% of the pool)")
    print(f"  per batch: " + ", ".join(f"b{b}={len(x)}" for b, x in enumerate(batches)))
    return keymap


if __name__ == "__main__":
    args = [a for a in sys.argv[1:] if not a.startswith("--")]
    n = 4 if "--bot" in sys.argv else 8
    build(args[0], [a.split("=", 1) for a in args[1:]], nlines=n, align="--align" in sys.argv, no_drop="--no-drop" in sys.argv)