#!/usr/bin/env python3
"""Export the reviewed Kuku Yalanji corpus as TAGGED, STRATIFIED, SPLIT training data.

Design (informed by the v20 failure = wrong conditional distribution, not too few rows):
  * METADATA tags  -> stored in every row's `meta` object. Rich. Used for
                      filtering / balancing / stratified split / route-specific eval.
  * INPUT CONTROL tag -> a SINGLE stable source token (`<synth>` here; the human
                      anchor rows carry `<gold>`). Everything else stays metadata,
                      because at INFERENCE the user only has an English sentence — they
                      will not hand-label negation/tense, so those must NOT be input tags.
                      Emitted as a separate `src_tagged` field so the trainer can run the
                      tagged-vs-untagged ablation without re-exporting.

Split policy:
  * Split is stratified BY DOMAIN and deterministic (stable hash of batch+seq), so every
    domain appears in train/dev/synthtest — no whole-domain leakage, no random-template
    leakage (our rows are hand-authored distinct arcs, not templated near-duplicates).
  * Exact-duplicate en/ku are forced into the SAME split (guard against soft leakage).
  * The ~500 HUMAN pairs (elder 43 + DB usage 449) are the LOCKED anchor/test and live in
    separate files — NEVER generated from, NEVER dumped whole into train. This exporter
    only touches synthetic.db; the human split is asserted, not produced here.
"""
import sqlite3, json, re, hashlib, os, collections

DB = "/mnt/donto-data/donto-resources/research/translation-training/kuku-yalanji-runpod-2026-06-30/synthetic/claude-synthetic-v1-2026-07-02/synthetic.db"
OUT = "/mnt/donto-data/donto-resources/research/translation-training/kuku-yalanji-runpod-2026-06-30/synthetic/claude-synthetic-v1-2026-07-02/export"

# batch -> domain (coarse; recent long-form expansion is domain-coherent per batch)
DOMAIN = {23:"body-coastal",24:"social-emotion",25:"craft-activity",26:"hunting",
          27:"sea-fishing",28:"plant-healing",29:"camp-life",30:"weather-sky"}
def domain_for(bid, label): return DOMAIN.get(bid, "core-daily-life")

GRAMMAR_PROBES = [
    ("negation",      lambda K,A,E: bool(re.search(r'\bkari\b|\bkariku\b', K))),
    ("question",      lambda K,A,E: E.rstrip().endswith('?') or bool(re.search(r'babaji|\bwanya\b|wanjawanja|wanyurringku|\bwanyu\b', K))),
    ("imperative",    lambda K,A,E: 'IMP' in A),
    ("purposive",     lambda K,A,E: 'PURP' in A or '-nka' in A),
    ("precautionary", lambda K,A,E: 'PRECAUT' in A or '-nyji' in A),
    ("successive",    lambda K,A,E: 'SUCC' in A),
    ("subordinate",   lambda K,A,E: 'SUB' in A),
    ("reciprocal",    lambda K,A,E: 'RECIP' in A),
    ("reflexive",     lambda K,A,E: 'REFL' in A or 'middle' in A),
    ("comitative",    lambda K,A,E: 'COMIT' in A),
    ("ablative",      lambda K,A,E: 'ABL' in A),
    ("inchoative",    lambda K,A,E: 'INCHO' in A),
    ("causative",     lambda K,A,E: 'CAUS' in A or '-bunga' in A or '-mani-' in A),
    ("ergative",      lambda K,A,E: 'ERG' in A),
    ("temp3",         lambda K,A,E: 'TEMP3' in A or '-ny-arr' in A),
    ("conditional",   lambda K,A,E: 'kaki' in K or 'COND' in A),
]

def grammar_tags(ku, an, en):
    K, A, E = ku.lower(), an, en
    return [name for name, fn in GRAMMAR_PROBES if fn(K, A, E)]

def split_for(bid, seq, dup_key, dup_first):
    # exact-dup en/ku pinned to first-seen split
    if dup_key in dup_first: return dup_first[dup_key]
    h = int(hashlib.sha1(f"{bid}:{seq}".encode()).hexdigest(), 16) % 100
    return "dev" if h < 8 else "synthtest" if h < 16 else "train"

def main():
    os.makedirs(OUT, exist_ok=True)
    db = sqlite3.connect(DB)
    rows = db.execute("""SELECT b.id, b.label, s.seq, s.english, s.kuku, s.analysis, s.frame,
                                s.tier, s.words_used, s.confidence, s.status
                         FROM sentences s JOIN batches b ON b.id=s.batch_id
                         WHERE s.status IN ('accepted','revised') ORDER BY b.id, s.seq""").fetchall()
    dup_first, seen = {}, {}
    for bid,label,seq,en,ku,an,frame,tier,wu,conf,status in rows:
        seen.setdefault((en.strip().lower(), ku.strip().lower()), []).append((bid,seq))
    for k,v in seen.items():
        if len(v) > 1:
            b0,s0 = v[0]; dup_first[k] = split_for(b0,s0,None,{})

    buckets = collections.defaultdict(list)
    cov = collections.defaultdict(collections.Counter)
    for bid,label,seq,en,ku,an,frame,tier,wu,conf,status in rows:
        dom = domain_for(bid,label)
        gt = grammar_tags(ku,an,en)
        try: lex = json.loads(wu) if wu else []
        except Exception: lex = []
        dk = (en.strip().lower(), ku.strip().lower())
        split = dup_first.get(dk) or split_for(bid,seq,dk,dup_first)
        nwords = len(ku.split())
        rec = {
            "en": en, "ku": ku, "analysis": an,
            "src_tagged": f"<synth> {en}",            # input-control ablation field
            "meta": {
                "source": "synthetic-claude-v1", "register": "elder-patz",
                "domain": dom, "frame": frame, "grammar": gt,
                "lexical_targets": lex, "quality_tier": {"high":"A","medium":"B","low":"C"}.get(conf,"B"),
                "template_family": label, "n_kuku_words": nwords, "status": status,
            },
        }
        buckets[split].append(rec)
        cov[split]["_rows"] += 1; cov[split]["_domain:"+dom] += 1
        for g in gt: cov[split]["g:"+g] += 1

    for split, recs in buckets.items():
        with open(os.path.join(OUT, f"{split}.jsonl"), "w") as f:
            for r in recs: f.write(json.dumps(r, ensure_ascii=False)+"\n")

    # policy + manifest
    manifest = {
        "corpus": "claude-synthetic-v1", "total_rows": len(rows),
        "splits": {k: len(v) for k,v in buckets.items()},
        "input_control_tags": ["<synth> (synthetic) / <gold> (human anchor, applied at merge time)"],
        "metadata_only_tags": ["domain","frame","grammar[]","lexical_targets[]","quality_tier","template_family"],
        "human_anchor_policy": {
            "elder_gold_pairs": 43, "db_usage_pairs": 449,
            "locked_test": "50-100 elder/gold rows, held out, NEVER trained, NEVER in generation prompts",
            "train_anchor": "remaining human rows, tagged <gold>, OVERSAMPLED to ~10-15% of train exposures",
            "leakage_guard": "synthetic rows authored from grammar.md/dictionary, NOT paraphrased from locked test",
        },
        "split_policy": "stratified-by-domain, deterministic hash(batch:seq), exact-dup pinned same split",
    }
    with open(os.path.join(OUT, "MANIFEST.json"), "w") as f:
        json.dump(manifest, f, indent=2, ensure_ascii=False)

    print(f"Exported {len(rows)} rows -> {OUT}/")
    for split in ("train","dev","synthtest"):
        c = cov[split]; n = c["_rows"]
        if not n: continue
        print(f"\n[{split}]  {n} rows")
        doms = sorted([(k[8:],v) for k,v in c.items() if k.startswith("_domain:")], key=lambda x:-x[1])
        print("  domains: " + ", ".join(f"{d}:{v}" for d,v in doms))
    # global grammar coverage
    tot = collections.Counter()
    for split in buckets:
        for k,v in cov[split].items():
            if k.startswith("g:"): tot[k[2:]] += v
    print("\nGrammar-construction coverage (all splits):")
    for g,v in tot.most_common():
        print(f"  {g:14s} {v:5d}  ({100*v/len(rows):.0f}%)")

if __name__ == "__main__":
    main()
