"""Casual / chat-romanization keys (AksharaLP v2).

Two lossy keys are computed for a roman token and for a Devanagari word; a
casual spelling matches the intended word when the keys collide.
  K1: consonant classes (aspiration, nukta, sibilants merged), vowel classes
      (length merged, ai->e, au->o), every 'a'/schwa removed except word
      initial, nasal dropped, repeats collapsed.
  K2: K1 with all vowels removed (chat consonant-only spellings: bht, thk).
The index is built at load time from the lexicon (0 bytes stored).
Same algorithm is ported 1:1 to JS (release/runtime) - keep in sync.
"""
import re

ASP = set("kgjtdpbc")
# Devanagari -> key class (nukta ignored except ड़/ढ़ -> R)
_DC = {}
for chars, k in [("कख", "k"), ("गघ", "g"), ("चछ", "c"), ("जझ", "j"), ("टठतथ", "t"),
                 ("डढदध", "d"), ("णनऩ", "n"), ("पफ", "p"), ("बभ", "b"), ("म", "m"),
                 ("य", "y"), ("रऱ", "r"), ("लळ", "l"), ("व", "v"), ("शषस", "s"), ("ह", "h")]:
    for c in chars: _DC[c] = k
for c, k in zip("\u0958\u0959\u095a\u095b\u095c\u095d\u095e\u095f", "kkgjRRpy"): _DC[c] = k
_DV = {}
for chars, k in [("अआा", "a"), ("इईिी", "i"), ("उऊुू", "u"), ("एऐेैॅऍ", "e"), ("ओऔऑोौॉ", "o"), ("ऋृ", "r")]:
    for c in chars: _DV[c] = k
VOW = set("aeiou")


def _fin(s):
    # drop non-initial a, collapse repeats
    out = []
    for i, ch in enumerate(s):
        if ch == "a" and i > 0: continue
        if out and out[-1] == ch: continue
        out.append(ch)
    return "".join(out)


def k2_of(k1):
    s = k1[0] + "".join(c for c in k1[1:] if c not in VOW) if k1 else ""
    if s and s[0] in VOW: s = s[1:] or s
    out = []
    for ch in s:
        if out and out[-1] == ch: continue
        out.append(ch)
    return "".join(out)


def dev_keys(w):
    """K1 keys of a Devanagari word (2 if it has ड़/ढ़: as d and as r)."""
    raw = []
    for i, ch in enumerate(w):
        if ch == "\u093c":
            if raw and raw[-1] == "d" and i > 0 and w[i - 1] in "डढ": raw[-1] = "R"
            continue
        if ch in _DC: raw.append(_DC[ch])
        elif ch in _DV: raw.append(_DV[ch])
    s = "".join(raw)
    if not s: return []
    ks = {_fin(s.replace("R", "d")), _fin(s.replace("R", "r"))}
    return sorted(k for k in ks if k)


def _norm(t):
    t = t.lower()
    t = re.sub(r"[^a-z]", "", t)
    return t


def rom_keys(t, cap=12):
    """K1 keys of a roman token, expanded over ambiguous casual rules."""
    t = _norm(t)
    if not t: return [], []
    t = re.sub(r"(.)\1{2,}", r"\1", t)          # emphasis repeats
    # deterministic letter merges
    t = t.replace("chh", "c").replace("cch", "c").replace("ch", "c").replace("sh", "s")
    t = t.replace("x", "ks").replace("q", "k").replace("z", "j").replace("f", "p").replace("w", "v")
    # (unit list) each position: list of alternatives
    alts = []
    i = 0; L = len(t)
    while i < L:
        c = t[i]; nx = t[i + 1] if i + 1 < L else ""
        prev = t[i - 1] if i > 0 else ""
        if c == "h" and prev in ASP and i > 0:
            alts.append(("", "h")); i += 1; continue          # aspiration / separate h
        if nx and c in "ae" and nx in "iy" and (i + 2 >= L or t[i + 2] not in VOW):
            if c == "a" and nx == "i": alts.append(("e", "ai")); i += 2; continue
            if c == "a" and nx == "y" and i + 2 >= L: alts.append(("e",)); i += 2; continue
            if c == "e" and nx == "y" and i + 2 >= L: alts.append(("e",)); i += 2; continue
        if nx and c in "ao" and nx in "uvw" and (i + 2 >= L or t[i + 2] not in VOW) and c + nx != "av":
            alts.append(("o", c + "u")); i += 2; continue
        if c == "o" and nx == "u": alts.append(("o", "u")); i += 2; continue
        if c == "e" and nx == "e": alts.append(("i",)); i += 2; continue
        if c == "o" and nx == "o": alts.append(("u", "o")); i += 2; continue
        if c == "e" and nx == "i" and (i + 2 >= L or t[i + 2] not in VOW): alts.append(("e",)); i += 2; continue
        if c == "e" and nx == "h" and i + 2 < L and t[i + 2] not in VOW and i > 0: alts.append(("e", "")); i += 1; continue
        if nx and c == "i" and nx in "aeou" and i > 0: alts.append(("y", "i")); i += 1; continue
        if c == "y" and i == L - 1 and prev and prev not in VOW: alts.append(("e", "i")); i += 1; continue
        if c == "n" and prev in VOW and i > 0 and (i == L - 1 or (nx and nx not in VOW and nx != "y")):
            alts.append(("n", "")); i += 1; continue
        if c in "iu" and prev == "r" and i >= 2 and t[i - 2] not in VOW:
            alts.append((c, "")); i += 1; continue
        alts.append((c,)); i += 1
    outs = [""]
    for a in alts:
        nxt = []
        for o in outs:
            for x in a: nxt.append(o + x)
        outs = list(dict.fromkeys(nxt))[:64]
    k1 = []; seen = set()
    for o in outs:
        k = _fin(o)
        if k and k not in seen: seen.add(k); k1.append(k)
    k1 = k1[:cap]
    k2 = list(dict.fromkeys(k2_of(k) for k in k1 if k2_of(k)))
    return k1, k2


# Irregular chat abbreviations (not derivable by rules). Values: candidate natives in prior order.
ABBR = {
    "h": ["है", "हैं"], "hy": ["है"], "hn": ["हाँ", "हैं"], "hm": ["हम", "हम्म"],
    "k": ["के", "की", "ओके"], "ki": ["की", "कि"], "n": ["ना", "न"], "ni": ["नहीं", "नी"], "nai": ["नहीं", "नई"],
    "nhi": ["नहीं"], "nhe": ["नहीं"], "nahi": ["नहीं"], "nai": ["नहीं", "नई"],
    "kr": ["कर"], "krna": ["करना"], "kro": ["करो"], "tm": ["तुम"], "tu": ["तू"], "mai": ["मैं", "में"],
    "me": ["में", "मैं"], "m": ["मैं", "में"], "thk": ["ठीक"], "tk": ["ठीक", "तक"],
    "yr": ["यार"], "yrr": ["यार"], "plz": ["प्लीज़"], "pls": ["प्लीज़"], "plzz": ["प्लीज़"],
    "tq": ["थैंक यू"], "ty": ["थैंक यू"], "thx": ["थैंक्स"], "thnx": ["थैंक्स"], "thnks": ["थैंक्स"],
    "gn": ["गुड नाइट"], "gm": ["गुड मॉर्निंग"], "ok": ["ओके"], "okk": ["ओके"], "kk": ["ओके"],
    "sry": ["सॉरी"], "bro": ["ब्रो"], "msg": ["मैसेज"], "bday": ["बर्थडे"], "clg": ["कॉलेज"],
    "ofc": ["ऑफ़िस"], "pic": ["पिक"], "bhai": ["भाई"], "bht": ["बहुत"], "bhot": ["बहुत"], "bohot": ["बहुत"],
    "pta": ["पता"], "kb": ["कब"], "jb": ["जब"], "sb": ["सब"], "mtlb": ["मतलब"], "kch": ["कुछ"], "kuch": ["कुछ"],
}


class CasualIndex:
    """K1/K2 -> lexicon ids sorted by frequency (top `top_n` lexicon words only)."""

    def __init__(self, words, uni=None, top_n=300000, per_key=8):
        order = range(len(words))
        if uni is not None:
            u = uni if callable(uni) else uni.__getitem__
            order = sorted(order, key=lambda i: (-u(i), i))
        self.k1 = {}; self.k2 = {}
        for r, i in enumerate(order):
            if r >= top_n: break
            for k in dev_keys(words[i]):
                b = self.k1.setdefault(k, [])
                if len(b) < per_key and i not in b: b.append(i)
                k2 = k2_of(k)
                if k2:
                    b2 = self.k2.setdefault(k2, [])
                    if len(b2) < per_key and i not in b2: b2.append(i)

    def lookup(self, roman, n1=6, n2=4):
        k1, k2 = rom_keys(roman)
        r1 = []; r2 = []
        for k in k1:
            for i in self.k1.get(k, ()):
                if i not in r1: r1.append(i)
        for k in k2:
            for i in self.k2.get(k, ()):
                if i not in r2: r2.append(i)
        return r1[:n1 * 2], r2[:n2 * 2]


def squeeze_in(t, mode=0):
    """Emphasis-repeat normalization of the model input. mode 0: off; 1: 3+ -> 1; 2: 3+ -> 2 for a/e/o, 1 otherwise."""
    if not mode: return t
    if mode == 1: return re.sub(r"([a-z])\1{2,}", r"\1", t)
    return re.sub(r"([a-z])\1{2,}", lambda m: m.group(1) * (2 if m.group(1) in "aeo" else 1), t)


def squeeze(t):
    """Normalize emphasis repeats (3+ same letters -> keep 1, e.g. pleaseee->please)."""
    return re.sub(r"([a-zA-Z])\1{2,}", r"\1", t)