"""Two-layer canvas data (user idea 2026-09-29): a SOUND layer of pronounced syllables and a WORD layer of written
words spanning 1+ sound cells.

User: "There got to be two layers. Syllable layer that are pronounciation of วรรค such as ไฉน => ฉะ | ไหน,
สลาย => สะ | หลาย, อภัย => อะ | ไพ. Rhyme cleanness should focus on this layer. Naturalness and misspell layer
should be judged on other layer."

Units. A วรรค is cut into written words by newmm and each word is read by TLTK g2p word by word - exactly the path
klon8.prosody.g2p_line takes, so the sound cells ARE the syllables the form checker counts, rhymes and tones.
One cell = one pronounced syllable. Every cell carries
  sound: the TLTK syllable string (e.g. cha1, naj4) -> onset / rime family / tone, exact, in context
  word:  the written word it belongs to, and its index k inside that word (0..n-1)
The model's token at a cell is the pair (word, k) - "word-slot" - so a word is never cut into written pieces:
ไฉน is two cells (ไฉน,0)(ไฉน,1) with sounds cha1 | naj4, never ไฉ | น.

ๆ repeats the whole preceding word (ดีๆ -> ดี ดี); ฯ / ๏ are dropped (fixbase_data.prep_line, same choice as FIXED).

Commands (CPU only, no GPU):
  twolayer_data.py measure DATA [--n N]       numbers for the design (fragments, inventories, coverage, drops)
  twolayer_data.py build DATA OUT_DIR [--words 30000]   poems.jsonl + lexicon.json for the prototype
  twolayer_data.py test                       round trip / label / hand cases
"""
import os, sys, re, json, collections, functools, argparse
sys.path[:0] = [".", "scripts"]
from klon8 import prosody as P

C, TL, L = 10, 12, 4                     # cells per วรรค, title cells, วรรค per บท (same grid as the canvas)
THAI = re.compile(r"[ก-๙]")


def prep_line(line):
    """ฯ / ๏ are not pronounced (FIXED's stated choice, fixbase_data.prep_line)."""
    return line.replace("ฯลฯ", "").replace("ฯ", "").replace("๏", "")


def rime_key(s):
    return P._MERGED_RIME.get(s.rime, s.rime)


@functools.lru_cache(maxsize=300000)
def word_sounds(w):
    """TLTK syllables of one written word, parsed; () if any piece does not parse (the word is unusable)."""
    raw = P.g2p_word(w)
    if not raw: return ()
    out = []
    for ps in raw.split("~"):
        s = P._parse_syl(w, ps)
        if s is None: return ()
        out.append(s)
    return tuple(out)


def line_words(line):
    """newmm words of one วรรค, Thai only, with ๆ expanded to a copy of the previous WORD (ดีๆ -> ดี ดี)."""
    from pythainlp.tokenize import word_tokenize
    line = prep_line(line)
    clean = P._CLEAN.sub(" ", line)
    words = []
    for w in word_tokenize(re.sub(r"\s+", " ", clean).strip(), engine="newmm"):
        w = w.strip()
        if not w: continue
        parts = w.split("ๆ")
        for j, part in enumerate(parts):
            if j and words: words.append(words[-1])
            if part and THAI.search(part): words.append(part)
    return words


def line_cells(line):
    """[(word, k, n, Syl)] one per pronounced syllable, or None if a word has no parseable reading."""
    cells = []
    for w in line_words(line):
        ss = word_sounds(w)
        if not ss: return None
        for k, s in enumerate(ss): cells.append((w, k, len(ss), s))
    return cells


def checker_count(words):
    """What klon8.prosody counts for the JOINED words (the model writes words with no spaces; newmm may re-cut)."""
    return len(P.g2p_line("".join(words)))


def checker_phons(words):
    return [s.phon for s in P.g2p_line("".join(words))]


def poem_record(row, title_clean_fn=None):
    """One บท -> {"title": [cells], "grid": [[cells] x 4]} with cell = [word, k, n, phon], or (None, reason)."""
    lines = row["poem"].split("\n")
    if len(lines) != L: return None, "not4"
    grid = []
    for ln in lines:
        cs = line_cells(ln)
        if cs is None: return None, "g2p_unparsed"
        if not 7 <= len(cs) <= C: return None, "count"
        words = [w for (w, k, n, s) in cs if k == 0]
        # the checker re-cuts the JOINED words with newmm; its syllables must be ours exactly (not just their number:
        # ฝาก|ว่า rejoined reads ฝา|กว่า - same count, different sounds; test 2026-09-29)
        if checker_phons(words) != [s.phon for (w, k, n, s) in cs]: return None, "rejoin"
        grid.append([[w, k, n, s.phon] for (w, k, n, s) in cs])
    t = row.get("title", "") or ""
    if title_clean_fn: t = title_clean_fn(t)
    tc = line_cells(t) or []
    # title: whole words only, up to TL sound cells (a word that would cross cell 12 is dropped, not cut)
    title = []
    for c in tc:
        if c[1] == 0 and len(title) + c[2] > TL: break
        title.append([c[0], c[1], c[2], c[3].phon])
    return {"title": title, "grid": grid}, "ok"


# ------------------------------------------------------------------ measurement

def measure(path, n_max=0):
    import syl_align_v2 as A2
    from pythainlp.tokenize import word_tokenize
    rows = [json.loads(l) for l in open(path, encoding="utf-8") if l.strip()]
    if n_max: rows = rows[:n_max]
    st = collections.Counter()
    word_tok = collections.Counter(); end_word = collections.Counter(); phon_tok = collections.Counter()
    reasons = collections.Counter()
    a2_keep = 0
    for r in rows:
        lines = r["poem"].split("\n")
        if len(lines) != L: reasons["not4"] += 1; continue
        # ---- the FIXED base's cells (aligned2 + ๆ/ฯ), grouped by newmm word, as fixbase load_poems builds them
        a2_ok = True; a2_lines = []
        for ln in lines:
            ln2 = prep_line(ln)
            per_word = []
            for w in word_tokenize(ln2, engine="newmm"):
                if not re.search(r"[ก-๙]", w): continue
                got = A2.split_word(w) or (w,)
                per_word.append((w, got))
            cells = [p for _, g in per_word for p in g]
            if "ๆ" in ln2 or not (7 <= len(cells) <= C) or len(cells) != len(P.g2p_line("".join(cells))):
                a2_ok = False; break
            a2_lines.append(per_word)
        if a2_ok:
            a2_keep += 1
            for per_word in a2_lines:
                last_multi = False
                for wi, (w, pieces) in enumerate(per_word):
                    ws = word_sounds(w)
                    for k, pc in enumerate(pieces):
                        st["a2_cells"] += 1
                        multi = len(pieces) > 1
                        st["a2_in_multi"] += multi
                        # misread: the piece read ALONE (how a per-type label sees it) != its in-context syllable
                        alone = P.g2p_line(pc)
                        if ws and len(ws) == len(pieces):
                            st["a2_cmp"] += 1
                            ctx = ws[k]
                            bad = (len(alone) != 1 or alone[0].phon != ctx.phon)
                            st["a2_misread"] += bad
                            st["a2_misread_rt"] += (len(alone) != 1 or (rime_key(alone[0]), alone[0].tone) != (rime_key(ctx), ctx.tone))
                    if wi == len(per_word) - 1:
                        st["a2_final"] += 1; st["a2_final_multi"] += len(pieces) > 1
        # ---- two-layer cells
        rec, why = poem_record(r)
        reasons[why] += 1
        if rec is None: continue
        for g in rec["grid"]:
            for (w, k, n, ph) in g:
                st["tl_cells"] += 1; phon_tok[ph] += 1
                if k == 0: word_tok[w] += 1
            w, k, n, ph = g[-1]
            end_word[w] += 1; st["tl_final_multi"] += n > 1; st["tl_final"] += 1
    kept = reasons["ok"]
    out = {"rows": len(rows), "a2_kept_poems": a2_keep, "twolayer_kept_poems": kept, "drop_reasons": dict(reasons)}
    out["a2_cells_in_multisyllable_word"] = st["a2_in_multi"] / max(st["a2_cells"], 1)
    out["a2_cells_misread_alone_phon"] = st["a2_misread"] / max(st["a2_cmp"], 1)
    out["a2_cells_misread_alone_rime_tone"] = st["a2_misread_rt"] / max(st["a2_cmp"], 1)
    out["a2_linefinal_in_multisyllable_word"] = st["a2_final_multi"] / max(st["a2_final"], 1)
    out["tl_linefinal_multisyllable_word"] = st["tl_final_multi"] / max(st["tl_final"], 1)
    tot_p = sum(phon_tok.values())
    out["phon_types"] = len(phon_tok)
    for k in (1000, 2000, 3000):
        out[f"phon_cov_top{k}"] = sum(c for _, c in phon_tok.most_common(k)) / tot_p
    out["phon_types_ge3"] = sum(1 for c in phon_tok.values() if c >= 3)
    tot_w = sum(word_tok.values()); tot_e = sum(end_word.values())
    out["word_types"] = len(word_tok); out["word_tokens"] = tot_w
    out["word_types_ge2"] = sum(1 for c in word_tok.values() if c >= 2)
    out["word_types_ge3"] = sum(1 for c in word_tok.values() if c >= 3)
    for k in (20000, 30000, 50000):
        top = dict(word_tok.most_common(k))
        out[f"word_unk_tok_{k}"] = 1 - sum(top.values()) / tot_w
        out[f"end_word_unk_{k}"] = 1 - sum(c for w, c in end_word.items() if w in top) / tot_e
        # poems with >= 1 UNK word
    out["syl_per_word_tokens"] = st["tl_cells"] / max(tot_w, 1)
    nsyl = collections.Counter(len(word_sounds(w)) for w in word_tok.elements())
    out["word_len_dist_tokens"] = {k: v / tot_w for k, v in sorted(nsyl.items())}
    return out, word_tok, phon_tok


KMAX = 8          # longest word the model represents (syllables); the longest training-poem word has 8


def extra_words():
    """Words beyond the poems (user 2026-09-29: the pretrained LM knows far more words than the poems use):
    pythainlp thai_words + WangchanBERTa vocabulary pieces that are whole Thai words. A piece counts as a word when it is
    Thai only, starts with a consonant or a leading vowel, and newmm keeps it as one token. Returns {word: source}."""
    from pythainlp.corpus import thai_words
    from pythainlp.tokenize import word_tokenize
    from transformers import AutoTokenizer
    out = {w: "dict" for w in thai_words()}
    tk = AutoTokenizer.from_pretrained("airesearch/wangchanberta-base-att-spm-uncased")
    for p in tk.get_vocab():
        p = p.replace("▁", "")
        if p and re.fullmatch(r"[ก-๛]+", p) and re.match(r"[ก-ฮเแโใไ]", p) and word_tokenize(p, engine="newmm") == [p]:
            out[p] = "wc" if p not in out else "dict+wc"
    return out


def usable(w):
    """a word the sound layer can carry: Thai only, g2p reads every syllable, at most KMAX syllables, and newmm reads it
    back as itself (so the checker re-cuts it as one word when it stands alone)"""
    from pythainlp.tokenize import word_tokenize
    if not re.fullmatch(r"[ก-๛]+", w) or "ๆ" in w or "ฯ" in w: return False
    ss = word_sounds(w)
    return bool(ss) and len(ss) <= KMAX and word_tokenize(w, engine="newmm") == [w]


def build(path, out_dir, n_words=0, clean_titles=True, vocab="union"):
    import title_clean
    os.makedirs(out_dir, exist_ok=True)
    recs, reasons = [], collections.Counter()
    for l in open(path, encoding="utf-8"):
        if not l.strip(): continue
        r = json.loads(l)
        rec, why = poem_record(r, title_clean.clean_title if clean_titles else None)
        reasons[why] += 1
        if rec: rec["id"] = r.get("id"); recs.append(rec)
    cnt = collections.Counter(c[0] for rec in recs for g in rec["grid"] for c in g if c[1] == 0)
    # vocab=poems: the POEM words only. vocab=union (default, user 2026-09-29): poem words + thai_words + whole-word
    # WangchanBERTa pieces, filtered by usable(); poem words first (by count), then the rest sorted. A title word outside
    # the lexicon is read as UNK plus its exact sound channels.
    words = [w for w, _ in (cnt.most_common(n_words) if n_words else cnt.most_common())]
    src = {w: "poems" for w in words}
    if vocab == "union":
        ex = extra_words(); seen = set(words); n0 = len(words)
        for w in sorted(ex):
            if w not in seen and usable(w): words.append(w); seen.add(w); src[w] = ex[w]
        print(f"  lexicon: {n0:,} poem words + {len(words) - n0:,} unseen words "
              f"({collections.Counter(src[w] for w in words[n0:])})")
    lex = {w: [s.phon for s in word_sounds(w)] for w in words}
    with open(os.path.join(out_dir, "poems.jsonl"), "w", encoding="utf-8") as f:
        for rec in recs: f.write(json.dumps(rec, ensure_ascii=False) + "\n")
    json.dump({"words": words, "sounds": lex, "count": {w: cnt[w] for w in words}, "source": src},
              open(os.path.join(out_dir, "lexicon.json"), "w", encoding="utf-8"), ensure_ascii=False)
    print(f"  {len(recs):,} poems kept of {sum(reasons.values()):,} {dict(reasons)}; lexicon {len(words):,} words")
    return recs, words


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

HAND = {"ไฉน": ["cha1", "naj4"], "สลาย": ["sa1", "laaj4"], "อภัย": ["?a1", "phaj0"], "ขออภัย": ["khOO4", "?a1", "phaj0"],
        "ทำงาน": ["tham0", "Naan0"], "สงบ": ["sa1", "Nop1"], "วันนี้": ["wan0", "nii3"]}


def test():
    fails = 0
    for w, want in HAND.items():
        got = [s.phon for s in word_sounds(w)]
        ok = got == want
        fails += not ok
        print(f"  {'OK ' if ok else 'BAD'} {w} -> {'|'.join(got)} (want {'|'.join(want)})")
    # round trip: word -> sound cells -> words -> text; checker agrees with the cell count
    for ln in ["ไฉนเลยใจสลายไม่อภัย", "ดีๆ ใจเราเฝ้ารอคอยเธอ", "แว่วแว่วมาว่าใครส่งใจหา"]:
        cs = line_cells(ln)
        words = [w for (w, k, n, s) in cs if k == 0]
        text = "".join(words)
        ok = (text == prep_line(ln).replace(" ", "").replace("ๆ", "") or "ๆ" in ln) and checker_count(words) == len(cs)
        fails += not ok
        print(f"  {'OK ' if ok else 'BAD'} {ln} -> {' | '.join(f'{w}:{s.phon}' for w, k, n, s in cs)}")
    print("PASS" if not fails else f"FAIL {fails}")
    return fails


if __name__ == "__main__":
    ap = argparse.ArgumentParser(); sub = ap.add_subparsers(dest="cmd")
    m = sub.add_parser("measure"); m.add_argument("data"); m.add_argument("--n", type=int, default=0); m.add_argument("--out", default=None)
    b = sub.add_parser("build"); b.add_argument("data"); b.add_argument("out_dir"); b.add_argument("--words", type=int, default=0, help="cap on poem words (0 = all)")
    b.add_argument("--vocab", choices=["poems", "union"], default="union")
    sub.add_parser("test")
    a = ap.parse_args()
    if a.cmd == "measure":
        res, wt, pt = measure(a.data, a.n)
        print(json.dumps(res, ensure_ascii=False, indent=1))
        if a.out: json.dump({"summary": res, "word_count": dict(wt), "phon_count": dict(pt)}, open(a.out, "w", encoding="utf-8"), ensure_ascii=False)
    elif a.cmd == "build":
        build(a.data, a.out_dir, a.words, vocab=a.vocab)
    elif a.cmd == "test":
        sys.exit(1 if test() else 0)
    P.save_cache()