#!/usr/bin/env python3
"""The text front-end of Swar: raw text -> the letters the model reads.

    from swar.text import normalize, to_ids, VOCAB
    normalize("तपाईंको balance Rs. 250 छ।")  ->  "तपाईंको balance दुई सय पचास rupees छ."

RULES, all enforced here so training and inference cannot drift:
  * Output holds ONLY Devanagari letters/signs, Latin letters, space, apostrophe and
    ordinary punctuation. No tags, no markers, no language switches -- the script
    a letter is written in IS its language.
  * Numbers become words in the language of the text around them (nearest letter
    wins): Devanagari context -> Nepali words, Latin context -> English words.
    Digit runs of 6+ are identifiers (phone, account, OTP) and are read digit by
    digit with "double" (डबल in Nepali).
  * Anything else is dropped (normalize) or makes the sentence unusable
    (is_clean), so a stray symbol can never reach the model as an unknown id.

This file is the vocabulary too: VOCAB is fixed here, and ids are positions in it,
so a checkpoint is only valid with the VOCAB it was trained on (saved inside it).
"""
from __future__ import annotations

import re
import unicodedata

# ---- vocabulary ----------------------------------------------------------------
PAD, BOS, EOS = "<pad>", "<s>", "</s>"            # internal ids, never in text
PUNCT = list(".,?!;:'\"-()…—")
LATIN = [chr(c) for c in range(ord("a"), ord("z") + 1)] + \
        [chr(c) for c in range(ord("A"), ord("Z") + 1)]
# Every Devanagari code point Nepali text uses: U+0900..U+0963 (letters, signs,
# vowel signs, virama, nukta) plus U+0970..U+097F; digits U+0966..U+096F are
# excluded because numbers are always verbalised first; dandas fold to "."
DEVANAGARI = [chr(c) for c in range(0x0900, 0x0964)] + [chr(c) for c in range(0x0970, 0x0980)]
MASK = "<mask>"                                   # stage A (PL-BERT) only; appended LAST so
VOCAB = [PAD, BOS, EOS, " "] + PUNCT + LATIN + DEVANAGARI + [MASK]   # no other id moves
assert len(VOCAB) <= 256, len(VOCAB)
STOI = {s: i for i, s in enumerate(VOCAB)}
ALLOWED = set(VOCAB[3:-1])

# ---- character folding -----------------------------------------------------------
_FOLD = {
    "।": ".", "॥": ".", "’": "'", "‘": "'", "`": "'", "“": '"', "”": '"',
    "–": "-", "‐": "-", "‑": "-", "...": "…", " ": " ",
}
_DROP = dict.fromkeys(map(ord, "​‌‍‎‏"), None)

# ---- English numbers ---------------------------------------------------------------
_ONES = ("zero one two three four five six seven eight nine ten eleven twelve thirteen "
         "fourteen fifteen sixteen seventeen eighteen nineteen").split()
_TENS = "_ _ twenty thirty forty fifty sixty seventy eighty ninety".split()
_SCALES = [(10 ** 9, "billion"), (10 ** 6, "million"), (1000, "thousand")]
_ORD = {"one": "first", "two": "second", "three": "third", "five": "fifth",
        "eight": "eighth", "nine": "ninth", "twelve": "twelfth"}


def en_cardinal(n: int) -> str:
    if n < 20:
        return _ONES[n]
    if n < 100:
        return _TENS[n // 10] + ("" if n % 10 == 0 else "-" + _ONES[n % 10])
    if n < 1000:
        rest = n % 100
        return _ONES[n // 100] + " hundred" + ("" if rest == 0 else " and " + en_cardinal(rest))
    for v, w in _SCALES:
        if n >= v:
            rest = n % v
            head = en_cardinal(n // v) + " " + w
            if rest == 0:
                return head
            return head + (" and " if rest < 100 else " ") + en_cardinal(rest)
    raise ValueError(n)


def en_ordinal(n: int) -> str:
    w = en_cardinal(n)
    last = re.split(r"([ -])", w)
    tail = last[-1]
    if tail in _ORD:
        tail = _ORD[tail]
    elif tail.endswith("y"):
        tail = tail[:-1] + "ieth"
    else:
        tail += "th"
    return "".join(last[:-1]) + tail


def en_digits(run: str) -> str:
    """Identifier reading: 9851005996 -> nine eight five one double zero ..."""
    out, i = [], 0
    while i < len(run):
        if i + 1 < len(run) and run[i] == run[i + 1]:
            out.append("double " + _ONES[int(run[i])])
            i += 2
        else:
            out.append(_ONES[int(run[i])])
            i += 1
    return " ".join(out)


# ---- Nepali numbers ------------------------------------------------------------------
_NE_ONES = ("शून्य एक दुई तीन चार पाँच छ सात आठ नौ दश एघार बाह्र तेह्र चौध पन्ध्र सोह्र सत्र "
            "अठार उन्नाइस बीस एक्काइस बाइस तेइस चौबीस पच्चीस छब्बीस सत्ताइस अठ्ठाइस उनन्तीस "
            "तीस एकतीस बत्तीस तेत्तीस चौँतीस पैँतीस छत्तीस सैँतीस अठतीस उनन्चालीस चालीस "
            "एकचालीस बयालीस त्रिचालीस चवालीस पैँतालीस छयालीस सच्चालीस अठचालीस उनन्चास पचास "
            "एकाउन्न बाउन्न त्रिपन्न चवन्न पचपन्न छपन्न सन्ताउन्न अन्ठाउन्न उनन्साठी साठी "
            "एकसट्ठी बयसट्ठी त्रिसट्ठी चौसट्ठी पैँसट्ठी छयसट्ठी सतसट्ठी अठसट्ठी उनन्सत्तरी सत्तरी "
            "एकहत्तर बहत्तर त्रिहत्तर चौहत्तर पचहत्तर छयहत्तर सतहत्तर अठहत्तर उनासी असी "
            "एकासी बयासी त्रियासी चौरासी पचासी छयासी सतासी अठासी उनान्नब्बे नब्बे "
            "एकानब्बे बयानब्बे त्रियानब्बे चौरानब्बे पन्चानब्बे छयानब्बे सन्तानब्बे अन्ठानब्बे "
            "उनान्सय").split()
assert len(_NE_ONES) == 100, len(_NE_ONES)
_NE_SCALES = [(10 ** 7, "करोड"), (10 ** 5, "लाख"), (1000, "हजार"), (100, "सय")]


def ne_cardinal(n: int) -> str:
    if n < 100:
        return _NE_ONES[n]
    for v, w in _NE_SCALES:
        if n >= v:
            rest = n % v
            head = ne_cardinal(n // v) + " " + w
            return head if rest == 0 else head + " " + ne_cardinal(rest)
    raise ValueError(n)


_NE_DIGIT = dict(zip("0123456789", _NE_ONES[:10]))


def ne_digits(run: str) -> str:
    """'9851005996' -> 'नौ आठ पाँच एक डबल शून्य पाँच डबल नौ छ'"""
    d = run.translate(str.maketrans("०१२३४५६७८९", "0123456789"))
    words, i = [], 0
    while i < len(d):
        if i + 1 < len(d) and d[i] == d[i + 1]:
            words.append("डबल " + _NE_DIGIT[d[i]])
            i += 2
        else:
            words.append(_NE_DIGIT[d[i]])
            i += 1
    return " ".join(words)


# ---- script context -------------------------------------------------------------------
def _is_dev(ch: str) -> bool:
    return 0x0900 <= ord(ch) <= 0x097F and not ("०" <= ch <= "९")


def _context_is_nepali(s: str, start: int, end: int) -> bool:
    """The sentence's majority script decides; the nearest letter only breaks ties.

    "तपाईंको balance Rs. 250 छ" is a Nepali sentence with one English loanword,
    and a Nepali speaker reads its number in Nepali -- nearest-letter alone would
    say English, because "balance" is closer than "छ".
    """
    dev = sum(1 for ch in s if _is_dev(ch))
    lat = sum(1 for ch in s if ch.isascii() and ch.isalpha())
    if dev != lat:
        return dev > lat
    for d in range(1, max(start, len(s) - end) + 1):
        for j in (start - d, end - 1 + d):
            if 0 <= j < len(s) and s[j].isalpha():
                return _is_dev(s[j])
    return False


_CURRENCY = re.compile(r"(?:\b(?:Rs|NPR|NRs)\.?|रु\.?|\$)\s*(?=[0-9०-९])")
_NUM = re.compile(r"[0-9०-९][0-9०-९,]*(?:\.[0-9०-९]+)*(?:st|nd|rd|th)?")
_TIME = re.compile(r"(?<![0-9०-९:])([0-9०-९]{1,2}):([0-9०-९]{2})(?![0-9०-९:])")
_DEV2ASCII = str.maketrans("०१२३४५६७८९", "0123456789")


def _verbalise_numbers(s: str) -> str:
    # currency prefix -> spoken unit after the number ("Rs 250" -> "250 rupees")
    # Placeholders are control characters, NOT letters: a placeholder made of
    # letters would itself be counted by the script-context check.
    def cur(m):
        return "\x01" if not m.group(0).startswith("$") else "\x02"
    s = _CURRENCY.sub(cur, s)
    s = _TIME.sub(lambda m: _time(s, m), s)

    out, last = [], 0
    for m in _NUM.finditer(s):
        try:
            word_or_none = _one_number(s, m, last, out)
        except (ValueError, IndexError):
            word_or_none = None
        if word_or_none is None:
            # Unreadable ("1,4th", "3.2.1"): leave the digits. They are not in
            # VOCAB, so is_clean() rejects the sentence instead of the model
            # hearing a made-up reading.
            out.append(s[last:m.end()])
        last = m.end()
    out.append(s[last:])
    return "".join(out).replace("\x01", "").replace("\x02", "")


def _time(s, m):
    """5:30 -> 'five thirty' / 'पाँच बजेर तीस मिनेट'; 9:00 -> "nine o'clock" / 'नौ बजे'."""
    h, mi = (int(x.translate(_DEV2ASCII)) for x in m.groups())
    if h > 24 or mi > 59:
        return m.group(0)
    if _context_is_nepali(s, m.start(), m.end()):
        return ne_cardinal(h) + " बजे" + ("" if mi == 0 else "र " + ne_cardinal(mi) + " मिनेट")
    if mi == 0:
        return en_cardinal(h) + " o'clock"
    return en_cardinal(h) + " " + ("oh " + _ONES[mi] if mi < 10 else en_cardinal(mi))


def _one_number(s, m, last, out):
    """Append the spoken form of number match m to out; return True, or None."""
    if m.group(0).translate(_DEV2ASCII).count(".") > 1 or re.match(r"[.:/][0-9०-९]", s[m.end():m.end() + 2]):
        return None                       # 3.2.1, 1/2, 12:30:15 -- no safe reading
    tok = m.group(0)
    ne = _context_is_nepali(s, m.start(), m.end())
    raw = tok.translate(_DEV2ASCII)
    suffix = ""
    mo = re.match(r"^(.*?)(st|nd|rd|th)$", raw)
    if mo and mo.group(1).isdigit():
        raw, suffix = mo.group(1), "ord"
    digits = raw.replace(",", "")
    pre = s[last:m.start()]
    money = pre.endswith(("\x01", "\x02"))
    if money and "." in digits and len(digits.split(".", 1)[1]) == 2:
        a, b = (int(x) for x in digits.split(".", 1))
        usd = pre.endswith("\x02")
        if ne:
            big, small = ("डलर", "सेन्ट") if usd else ("रुपैयाँ", "पैसा")
            word = f"{ne_cardinal(a)} {big}" + (f" {ne_cardinal(b)} {small}" if b else "")
        else:
            big, small = ("dollars", "cents") if usd else ("rupees", "paisa")
            word = f"{en_cardinal(a)} {big}" + (f" and {en_cardinal(b)} {small}" if b else "")
        out.append(pre[:-1] + word)
        return True
    if "." in digits:
        a, b = digits.split(".", 1)
        word = ((ne_cardinal(int(a)) + " दशमलव " + " ".join(_NE_ONES[int(c)] for c in b))
                if ne else
                (en_cardinal(int(a)) + " point " + " ".join(_ONES[int(c)] for c in b)))
    elif len(digits) >= 6 or (digits.startswith("0") and len(digits) > 1):
        word = ne_digits(digits) if ne else en_digits(digits)
    elif suffix == "ord" and not ne:
        word = en_ordinal(int(digits))
    else:
        word = ne_cardinal(int(digits)) if ne else en_cardinal(int(digits))
    unit = ""
    if pre.endswith("\x01"):
        pre, unit = pre[:-1], (" रुपैयाँ" if ne else " rupees")
    elif pre.endswith("\x02"):
        pre, unit = pre[:-1], (" डलर" if ne else " dollars")
    out.append(pre + word + unit)
    return True


# ---- public API -------------------------------------------------------------------------
def normalize(text: str) -> str:
    """Raw text -> model letters. Same function for training data and live requests."""
    s = unicodedata.normalize("NFC", text or "").translate(_DROP)
    for a, b in _FOLD.items():
        s = s.replace(a, b)
    s = _verbalise_numbers(s)
    s = s.replace("%", " percent").replace("&", " and ")
    s = "".join(ch if ch in ALLOWED else (" " if ch.isspace() else "") for ch in s)
    s = re.sub(r"\s+", " ", s).strip()
    s = re.sub(r"\s+([.,?!;:…])", r"\1", s)                 # no space before punctuation
    s = re.sub(r"([.?!])\1+", r"\1", s)
    return s


def is_clean(raw: str, normalized: str) -> bool:
    """True if normalising dropped nothing but whitespace and known foldings --
    i.e. the sentence can be trained on without its text silently differing from
    what the audio says."""
    s = unicodedata.normalize("NFC", raw or "").translate(_DROP)
    for a, b in _FOLD.items():
        s = s.replace(a, b)
    s = _verbalise_numbers(s).replace("%", " percent").replace("&", " and ")
    return all(ch in ALLOWED or ch.isspace() for ch in s) and bool(normalized)


def to_ids(normalized: str) -> list[int]:
    return [STOI[BOS]] + [STOI[c] for c in normalized] + [STOI[EOS]]


if __name__ == "__main__":
    print(f"VOCAB {len(VOCAB)} symbols: {len(PUNCT)} punct, {len(LATIN)} Latin, "
          f"{len(DEVANAGARI)} Devanagari, 4 internal/space")
    for t in ["तपाईंको balance Rs. 250 छ।",
              "Your balance is Rs 1,250.50, valid until the 15th.",
              "Please call 9851005996 or 015574574.",
              "मेरो नम्बर ९८५१००५९९६ हो, र म ५ बजे आउँछु।",
              "रु ५०० को रिचार्ज गर्नुहोस्",
              "A & B give 20% off… ‘really’?",
              "emoji 😀 and © symbols are dropped"]:
        n = normalize(t)
        print(f"{'clean' if is_clean(t, n) else 'DIRTY'}  {t!r}\n       -> {n!r}")