#!/usr/bin/env python3
"""exp15 — assemble the FINAL 30-节 trap set: merge deterministic straddles (straddle_ledger.json:
23 T-ent control + 7 T-ana) with kimi enrichment (ana_enrichment.json: T-ana/T-ell), dedup by
(type, boundary), require each to STRADDLE a flat-A0-30 greedy cut (common denominator), and map each
to the dependent chunk (i+1) so the judge can score each arm's rendering of that chunk.

Output final_traps.json: the trap set for Q1b (oracle vs flat) and Q3a (carryover) judging.
Run: python assemble_traps.py
"""
from __future__ import annotations
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from chunker import greedy_chunks

SD = Path("/home/ubuntu/books/gu-zhenren/exp15")
FLAT = SD / "S2prime_big_flat.txt"
OUT = SD / "final_traps.json"


def flat_boundaries(text):
    ch = greedy_chunks(text)
    offs, starts, pos = [], [0], 0
    for i, c in enumerate(ch):
        idx = text.find(c[:25], max(0, pos - 5))
        if idx < 0:
            idx = text.find(c[:25])
        if i > 0 and idx >= 0:
            offs.append(idx); starts.append(idx)
        pos = idx + len(c) if idx >= 0 else pos + len(c)
    return ch, offs, starts


def chunk_of_offset(starts, off):
    c = 0
    for i, s in enumerate(starts):
        if off >= s:
            c = i
    return c


def straddled_boundary(text, offs, a_span, d_span):
    ai = text.find(a_span[:25]); di = text.find(d_span[:25])
    if ai < 0 or di < 0:
        return None, None, None
    lo, hi = min(ai, di), max(ai + len(a_span), di + len(d_span))
    crossed = [k for k, cut in enumerate(offs) if lo < cut < hi]
    if not crossed:
        return None, ai, di
    return crossed[0], ai, di   # boundary index (0-based cut = between chunk k and k+1)


def main():
    text = FLAT.read_text(encoding="utf-8")
    ch, offs, starts = flat_boundaries(text)
    traps = []
    seen = set()   # (type, boundary_cut)

    def add(t, typ, control, a_span, d_span, phen, exp):
        bcut, ai, di = straddled_boundary(text, offs, a_span, d_span)
        if bcut is None:
            return "not_straddling"
        key = (typ, bcut)
        if key in seen:
            return "dup"
        seen.add(key)
        dep_chunk = chunk_of_offset(starts, di)      # chunk containing the dependent (scored chunk)
        anc_chunk = chunk_of_offset(starts, ai)
        traps.append({"id": f"{typ}-{len(traps)+1}", "type": typ, "control": control,
                      "boundary_cut": bcut, "anc_chunk": anc_chunk, "dep_chunk": dep_chunk,
                      "anchor_span": a_span, "dependent_span": d_span,
                      "phenomenon": phen, "expected_ru": exp})
        return "ok"

    stats = {"ok": 0, "dup": 0, "not_straddling": 0}
    # deterministic
    det = json.load(open(SD / "straddle_ledger.json"))["traps"]
    for t in det:
        r = add(t, t["type"], t.get("control", False), t["anchor_span"], t["dependent_span"],
                t["phenomenon"], t["expected_ru"])
        stats[r] += 1
    # enrichment
    enr = json.load(open(SD / "ana_enrichment.json"))["straddles"]
    for t in enr:
        r = add(t, t.get("type", "T-ana"), False, t.get("antecedent_span", ""), t.get("dependent_span", ""),
                t.get("phenomenon", ""), t.get("expected_ru", ""))
        stats[r] += 1

    from collections import Counter
    by = Counter(t["type"] for t in traps)
    prim = Counter(t["type"] for t in traps if not t["control"])
    OUT.write_text(json.dumps({
        "n": len(traps), "by_type": dict(by),
        "quota": {"T-ana": {"have": prim.get("T-ana", 0), "need": 20, "meets": prim.get("T-ana", 0) >= 20},
                  "T-ent": {"have": by.get("T-ent", 0), "need": 20, "meets": by.get("T-ent", 0) >= 20, "role": "negative-control"},
                  "T-ell": {"have": by.get("T-ell", 0), "need": 10, "meets": by.get("T-ell", 0) >= 10, "note": "secondary; short = honest-power"}},
        "undecidable_a_priori": ["T-tense", "T-cat", "T-gen", "T-reg"],
        "merge_stats": stats,
        "traps": traps,
    }, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"final trap set: {dict(by)} (total {len(traps)})  merge {stats}")
    print(f"  T-ana (primary): {prim.get('T-ana',0)} {'MEETS>=20' if prim.get('T-ana',0)>=20 else 'SHORT'}")
    print(f"  T-ent (control): {by.get('T-ent',0)} {'MEETS>=20' if by.get('T-ent',0)>=20 else 'SHORT'}")
    print(f"  T-ell (secondary): {by.get('T-ell',0)} (>=10? {by.get('T-ell',0)>=10})")
    print(f"  undecidable a-priori: T-tense/T-cat/T-gen/T-reg -> {OUT}")


if __name__ == "__main__":
    main()