"""Numeric equivalence for Thai ASR scoring.  SCORING ONLY -- never touch training targets.

WHY
---
FLEURS-th references write numbers as Arabic DIGITS in 21.5% of utterances.
Whisper-family models do inverse text normalisation and emit digits too
(Pathumma 19.2%, distill-whisper 20.1%, faster-whisper 25.8%, Qwen3 21.5%),
while NeMo RNNT models emit SPOKEN THAI (ChindaASR ~0%, typhoon-asr-realtime 0.0%).
Raw scoring therefore rewards output CONVENTION, not recognition.  Measured on
v3-arm/FLEURS: CER 20.34 on digit-bearing refs vs 8.49 on the rest.

CRITICAL IMPLEMENTATION NOTE
----------------------------
Thai has no word delimiters, so REGEX SUBSTRING MATCHING IS CATASTROPHIC here.
A first version matched number words inside ordinary words --
  สามารถ(able)->'3ารถ'  พันธุ์(species)->'1000ธุ์'  โกหก(lie)->'โก6'
  สัมพันธ์(relationship)->'สัม1000ธ์'
and made 239/1021 utterances WORSE.  Conversion MUST run on word-segmenter
tokens and only on tokens that are EXACTLY numerals.  pythainlp newmm keeps
สามารถ / พันธุ์พืช / โกหก / สัมพันธ์ whole while splitting สี่|หมื่น|คน.
"""
import re
from pythainlp.tokenize import word_tokenize

_D = {'ศูนย์':0,'หนึ่ง':1,'เอ็ด':1,'สอง':2,'สาม':3,'สี่':4,'ห้า':5,
      'หก':6,'เจ็ด':7,'แปด':8,'เก้า':9}
_M = {'สิบ':10,'ร้อย':100,'พัน':1000,'หมื่น':10000,'แสน':100000,'ล้าน':1000000}
# newmm emits numeral COMPOUNDS as single tokens (ยี่สิบ, หนึ่งร้อย, สิบเอ็ด...).
# A token counts as numeric only if numeral morphemes cover it ENTIRELY -- that
# full-match requirement is what makes สามารถ / พันธุ์ / โกหก / สัมพันธ์ safe.
_MORPH = sorted(list(_D) + list(_M) + ['ยี่'], key=len, reverse=True)
_FULL = re.compile('^(?:' + '|'.join(map(re.escape, _MORPH)) + ')+$')
_SPLIT = re.compile('(' + '|'.join(map(re.escape, _MORPH)) + ')')
_THAI_DIGIT = str.maketrans('๐๑๒๓๔๕๖๗๘๙', '0123456789')


def _is_num(tok):
    return bool(_FULL.match(tok))


def _morphs(tok):
    return [m for m in _SPLIT.split(tok) if m]


def _seq(tokens):
    """Parse a run of numeral tokens -> canonical digit string.
    positional  สี่ หมื่น        -> 40000
    digit-wise  สอง ศูนย์ ศูนย์ หก -> 2006  (years/IDs: no multiplier, >=2 digits)
    """
    tokens = [m for t in tokens for m in _morphs(t)]
    tokens = ['สอง' if t == 'ยี่' else t for t in tokens]
    if not tokens:
        return ''
    if not any(t in _M for t in tokens):
        if len(tokens) >= 2:
            return ''.join(str(_D[t]) for t in tokens)
        return str(_D[tokens[0]])
    total = 0; cur = 0; last = None
    for t in tokens:
        if t in _D:
            cur = _D[t]; last = 'd'
        else:
            m = _M[t]
            if m == 1000000:
                total = (total + max(cur, 1)) * m; cur = 0
            else:
                total += (cur if cur else 1) * m; cur = 0
            last = 'm'
    return str(total + cur)


def normalize_numbers(text):
    text = text.translate(_THAI_DIGIT)
    toks = word_tokenize(text, engine='newmm')
    out = []; run = []
    for t in toks:
        if _is_num(t):
            run.append(t)
        else:
            if run:
                out.append(_seq(run)); run = []
            out.append(t)
    if run:
        out.append(_seq(run))
    return ''.join(out)


if __name__ == '__main__':
    CASES = [
        ('สี่หมื่น','40000'), ('ยี่สิบ','20'), ('สองศูนย์ศูนย์หก','2006'),
        ('หนึ่งร้อยยี่สิบสาม','123'), ('สิบเอ็ด','11'), ('หนึ่งร้อย','100'), ('สองพันยี่สิบห้า','2025'), ('ห้าล้าน','5000000'),
        ('สามพันสองร้อย','3200'), ('เก้า','9'), ('๑๒๓','123'), ('40000','40000'),
        # REGRESSION GUARDS -- these broke the regex version
        ('สามารถตรวจจับ','สามารถตรวจจับ'), ('พันธุ์พืช','พันธุ์พืช'),
        ('โกหกหรือไม่','โกหกหรือไม่'), ('ความสัมพันธ์','ความสัมพันธ์'),
        ('ไม่มีตัวเลข','ไม่มีตัวเลข'),
    ]
    bad = 0
    for src, want in CASES:
        got = normalize_numbers(src); ok = got == want; bad += not ok
        print(f"{'ok ' if ok else 'FAIL'} {src!r} -> {got!r}" + ('' if ok else f"  want {want!r}"))
    print(f"\n{len(CASES)-bad}/{len(CASES)} pass")
    raise SystemExit(1 if bad else 0)