"""Split a Thai word into surface syllables by aligning the writing to the pronunciation.

The unit of กลอน is the syllable, so a grid cell must be one. No surface segmenter we have gets this
right: pythainlp's han_solo, dict and ssg engines and TLTK's syl_segment all return ถวิล, พันธนา and
ปรารถนา whole, and each invents splits of its own (ถวิ|ล, ปรา|รถ, แก้|มหอม). Only the pronunciation
is correct, because the count follows the sound — ถวิล is ถะ-หวิน.

TLTK gives the pronunciation one syllable at a time (ถวิล -> tha1~win4), so each syllable's ONSET is
known. This walks the written form left to right and cuts wherever the next syllable's onset letter
appears, carrying any leading vowel (เ แ โ ใ ไ) back to the syllable it belongs to.
"""
import sys, re, functools
sys.path.insert(0, ".")
from klon8 import prosody as P

LEAD_VOWEL = "เแโใไ"
# Thai consonant letters by the onset phone TLTK writes for them
ONSET = {
    "ก": "k", "ข": "kh", "ฃ": "kh", "ค": "kh", "ฅ": "kh", "ฆ": "kh", "ง": "ng",
    "จ": "c", "ฉ": "ch", "ช": "ch", "ฌ": "ch", "ซ": "s", "ศ": "s", "ษ": "s", "ส": "s",
    "ญ": "j", "ย": "j", "ฎ": "d", "ด": "d", "ฏ": "t", "ต": "t",
    "ฐ": "th", "ฑ": "th", "ฒ": "th", "ถ": "th", "ท": "th", "ธ": "th",
    "ณ": "n", "น": "n", "บ": "b", "ป": "p", "ผ": "ph", "พ": "ph", "ภ": "ph",
    "ฝ": "f", "ฟ": "f", "ม": "m", "ร": "r", "ล": "l", "ฬ": "l", "ว": "w",
    "ห": "h", "ฮ": "h", "อ": "?",
}
# ห as a silent tone marker, and the true clusters
DIGRAPH = {"หน": "n", "หม": "m", "หญ": "j", "หย": "j", "หง": "ng", "หล": "l", "หว": "w", "หร": "r",
           "กร": "kr", "กล": "kl", "กว": "kw", "ขร": "khr", "ขล": "khl", "ขว": "khw",
           "คร": "khr", "คล": "khl", "คว": "khw", "ปร": "pr", "ปล": "pl", "ผล": "phl",
           "พร": "phr", "พล": "phl", "ตร": "tr", "ทร": "s", "ดร": "dr", "บร": "br", "บล": "bl"}


def onset_at(w, i):
    """Every onset the letters at position i could spell, longest first."""
    out = []
    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


@functools.lru_cache(maxsize=200000)
def split_word(w):
    """Surface syllables of one word, count-matched to its pronunciation. () if alignment fails."""
    # the onset comes from prosody's own parser, not from a regex over the phone string:
    # "tha1" is th + a + tone, and [a-zA-Z?]+ grabs "tha", which matches no letter.
    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,)
    starts, k = [0], 1
    i = 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
            # a syllable that is pronounced with a coda needs at least an onset letter and a coda
            # letter written, so a one-character piece cannot be it. Without this, ออดอ้อน (?OOt~?OOn)
            # cuts at the second อ - which is the vowel of syllable 1, not the onset of syllable 2.
            if codas[k - 1] and begin - starts[-1] < 2: 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)]))


def syls_aligned(line, limit=None):
    """Surface syllables of one วรรค. Falls back to the whole word when alignment fails."""
    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


if __name__ == "__main__":
    tests = [("มาเหนี่ยวรั้งร้อยขวัญพันธนา", "มา|เหนี่ยว|รั้ง|ร้อย|ขวัญ|พัน|ธ|นา"),
             ("หอมกลิ่นคำร่ำรินถวิลนัก", "หอม|กลิ่น|คำ|ร่ำ|ริน|ถ|วิล|นัก"),
             ("ปรารถนาก็เวียนวกเกินยกพ้น", "ปราร|ถ|นา|ก็|เวียน|วก|เกิน|ยก|พ้น"),
             ("ไม่ได้จูบแก้มหอมดอมเสนอ", "ไม่|ได้|จูบ|แก้ม|หอม|ดอม|เส|นอ"),
             ("เผลอ-เผยความอาวรณ์ออดอ้อนหา", "เผลอ|เผย|ความ|อา|วรณ์|ออด|อ้อน|หา"),
             ("อารมณ์ผู้เพ็ญถวิล ก็รินไหล", "อา|รมณ์|ผู้|เพ็ญ|ถ|วิล|ก็|ริน|ไหล")]
    nc = ns = 0
    for t, hand in tests:
        got = syls_aligned(t)
        want = hand.split("|")
        cnt = len(got) == len(P.g2p_line(t))
        exact = got == tuple(want) or list(got) == want
        nc += cnt; ns += exact
        print(f"  {'EXACT' if exact else ('count' if cnt else '     '):5s} {len(got)} vs {len(want)}   " + " | ".join(got))
        if not exact: print(f"        hand: {' | '.join(want)}")
    print(f"\n  count correct {nc}/6, split identical to the hand segmentation {ns}/6")
    P.save_cache()