"""Surface syllables, v2: scripts/syl_align.py with two alignment bugs fixed (FIX 4, 2026-09-28 data audit).

v1 stays untouched so every model trained on it remains reproducible; models trained with --seg aligned2
use this module.

Bug (a): TLTK writes the velar nasal onset ง as "N", v1's tables said "ng", so no syllable starting with ง
or หง could ever be matched; every multi-syllable word containing one (ทำงาน, สงบ, เหงาหงอย) failed to
align and the whole บท fell out of training at the count filter (poems with ทำงาน were dropped 96.5% of the
time against 14.5% overall).

Bug (b): v1 cut at the FIRST letter spelling the next syllable's onset. When the previous syllable ends in
a consonant that is the same letter as the next onset (วันนี้, สัญญา, จากกัน, ศรัทธา) the cut landed one
letter early: วั|นนี้, สั|ญญา, จา|กกัน - count-correct, so the filter kept them, and the vocabulary learned
fake units (วั, นนี้, กกัน). v2 first requires the piece before a cut to contain a letter that can spell the
previous syllable's coda (or a vowel spelling that carries it: เ-า = w, ไ/ใ = j, -ำ = m); if no cut
satisfies that for the whole word (a letter shared by coda and onset, e.g. ชลบุรี), it falls back to v1's
rule for that word, so v2 never aligns fewer words than v1.

ๆ / ฯ are not changed here: klon8.prosody drops them from the pronunciation, so a line containing ๆ has a
cell count that differs from its pronounced count whichever way the mark is handled; changing that belongs
to the prosody checker, not to the segmenter (one change class per experiment).
"""
import sys, re, functools
sys.path.insert(0, "."); sys.path.insert(0, "scripts")
from klon8 import prosody as P
import syl_align as V1

LEAD_VOWEL = V1.LEAD_VOWEL
ONSET = dict(V1.ONSET, **{"ง": "N"})
DIGRAPH = dict(V1.DIGRAPH, **{"หง": "N"})
# อ นำ ย (อย่า อยู่ อย่าง อยาก) is a digraph only before ่ ู or า; elsewhere อ is a vowel (หรือ|ยัง, ขอ|ยืม)
AY_NEXT = set("ู่า")
# letters that can spell each coda phone TLTK writes
CODA_LETTERS = {
    "k": set("กขคฆ"), "t": set("ดตทธฎฏฐฑฒถจชซศษส"), "p": set("บปพฟภ"),
    "n": set("นณญรลฬ"), "m": set("ม"), "N": set("ง"), "j": set("ย"), "w": set("ว"),
}


def onset_at(w, i):
    out = []
    if w[i:i + 2] == "อย" and i + 2 < len(w) and w[i + 2] in AY_NEXT: out.append(("j", 2))
    if i + 1 < len(w) and w[i:i + 2] in DIGRAPH: out.append((DIGRAPH[w[i:i + 2]], 2))
    if w[i] in ONSET: out.append((ONSET[w[i]], 1))
    return out


def spells_coda(piece, coda):
    """Can this written piece carry the coda phone? The first letter is the onset and does not count."""
    if not coda or coda == "?": return True
    q = 1 if piece and piece[0] in LEAD_VOWEL else 0          # skip a leading vowel
    q += 2 if (piece[q:q + 2] in DIGRAPH or piece[q:q + 2] == "อย") else 1   # then the whole onset cluster
    body = piece[q:]
    if CODA_LETTERS.get(coda, set()) & set(body): return True
    if coda == "w" and "เ" in piece and "า" in piece: return True          # เ-า
    if coda == "j" and (piece[:1] in "ไใ"): return True                   # ไ, ใ
    if coda == "m" and "ำ" in piece: return True                          # -ำ
    return False


def _align(w, onsets, codas, strict):
    starts, k, i = [0], 1, 1
    while i < len(w) and k < len(onsets):
        for cand, n in onset_at(w, i):
            if cand != onsets[k]: continue
            begin = i - 1 if (i and w[i - 1] in LEAD_VOWEL) else i
            if begin <= starts[-1]: continue
            if codas[k - 1] and begin - starts[-1] < 2: continue
            if strict and not spells_coda(w[starts[-1]:begin], codas[k - 1]): continue
            starts.append(begin); k += 1; i = begin + n
            break
        else:
            i += 1
            continue
    if len(starts) != len(onsets): return ()
    return tuple(w[a:b] for a, b in zip(starts, starts[1:] + [len(w)]))


@functools.lru_cache(maxsize=200000)
def split_word(w):
    syls = P.g2p_line(w)
    if not syls: return ()
    onsets = [x.onset for x in syls]; codas = [x.coda for x in syls]
    if len(onsets) <= 1: return (w,)
    return _align(w, onsets, codas, True) or _align(w, onsets, codas, False)


def syls_aligned(line, limit=None):
    from pythainlp.tokenize import word_tokenize
    out = []
    for w in word_tokenize(line, engine="newmm"):
        if not re.search(r"[ก-๙]", w): continue
        got = split_word(w)
        out.extend(got if got else [w])
    return out[:limit] if limit else out


CASES = {  # word -> hand-checked split
    "ทำงาน": "ทำ|งาน", "สงบ": "ส|งบ", "เหงาหงอย": "เหงา|หงอย", "วันนี้": "วัน|นี้", "ความหมาย": "ความ|หมาย",
    "สัญญา": "สัญ|ญา", "จากกัน": "จาก|กัน", "ศรัทธา": "ศรัท|ธา", "บ้านนอก": "บ้าน|นอก", "เนิ่นนาน": "เนิ่น|นาน",
    "ชีวิต": "ชี|วิต", "เสมอ": "เส|มอ", "ออดอ้อน": "ออด|อ้อน", "ถวิล": "ถ|วิล", "พันธนา": "พัน|ธ|นา",
    "ปรารถนา": "ปราร|ถ|นา", "อารมณ์": "อา|รมณ์", "เลิกกัน": "เลิก|กัน", "กลับบ้าน": "กลับ|บ้าน",
    "หลายอย่าง": "หลาย|อย่าง", "อ้อยอิ่ง": "อ้อย|อิ่ง", "สงวน": "ส|งวน", "ฉงน": "ฉ|งน",
    "หรือยัง": "หรือ|ยัง", "ขอยืม": "ขอ|ยืม", "ยื้อยุด": "ยื้อ|ยุด", "ถือยศ": "ถือ|ยศ", "ปรนนิบัติ": "ปรน|นิ|บัติ",
    "กลืนน้ำ": "กลืน|น้ำ",
}

if __name__ == "__main__":
    ok = 0
    for w, want in CASES.items():
        got = "|".join(split_word(w)); old = "|".join(V1.split_word(w))
        ok += got == want
        print(f"  {'OK ' if got == want else 'BAD'} {w:10s} v2 {got:14s} v1 {old:14s} want {want}")
    print(f"{ok}/{len(CASES)} hand cases")