"""Dynamic Accent Flash text front-end.

    from frontend import FrontEnd
    fe = FrontEnd(diac_model="models/diacritics", mt=True)
    fe("<eng>The meeting is at 3:30 on 21/10/2026.</eng>", mode="codeswitch")          # tags kept, numbers verbalised
    fe("Ojo naa dara pupo", mode="native", lang="yor")                                # adds Yorùbá tone marks/underdots
    fe("Please come home early today.", mode="translate", lang="ibo")                 # MADLAD-400 eng/pcm -> ibo
    fe("I go call you tomorrow", mode="accent", lang="ibo")                           # English/Pidgin words, Igbo accent

Modes
  codeswitch : input already tagged <ibo>…</ibo> <eng>…</eng> (or untagged = one language `lang`); each segment is
               normalised in its own language; Yorùbá/Igbo segments that carry no marks are diacritised.
  native     : the whole text is one language (`lang`); normalise + diacritise.
  translate  : `lang` is the target (ibo|yor|hau); the input (English or Pidgin) is translated first with
               MADLAD-400-3B-MT (Apache-2.0), then normalised + diacritised (MADLAD Yorùbá is often unmarked).
  accent     : keep the words as written but read them with the `lang` accent (whole text tagged <lang>), i.e.
               "English words with a native accent". Numbers are verbalised in English.
Output: a tagged string for infer.py / flash_lm.py.
"""
import re
import unicodedata

try:
    from .normalize import normalize
except ImportError:  # run as a script
    from normalize import normalize

SEG = re.compile(r"<(\w{3})>(.*?)</\1>", re.S)


def _strip(s):
    return "".join(c for c in unicodedata.normalize("NFD", s) if unicodedata.category(c) != "Mn")


def unmarked(text, thresh=0.05):
    w = re.findall(r"\w+", unicodedata.normalize("NFC", text))
    if not w:
        return False
    return sum(_strip(x) != unicodedata.normalize("NFC", x) for x in w) / len(w) < thresh


class FrontEnd:
    def __init__(self, diac_model=None, mt=False, device="cuda"):
        self.diac_dir, self.device = diac_model, device
        self._d = self._mt = None
        self.use_mt = mt

    # ---- diacritics (ByT5-small fine-tune, word-guarded)
    def diacritise(self, text, lang):
        if lang not in ("yor", "ibo") or not self.diac_dir:
            return text
        import torch
        from transformers import AutoTokenizer, T5ForConditionalGeneration
        if self._d is None:
            self._d = (AutoTokenizer.from_pretrained(self.diac_dir),
                       T5ForConditionalGeneration.from_pretrained(self.diac_dir).to(self.device).eval())
        tok, m = self._d
        src = unicodedata.normalize("NFC", _strip(text))
        x = tok([f"{lang}: {src}"], return_tensors="pt").to(self.device)
        with torch.no_grad():
            y = m.generate(**x, max_length=int(x.input_ids.shape[1] * 1.6) + 16)
        hyp = unicodedata.normalize("NFC", tok.decode(y[0], skip_special_tokens=True))
        sw, hw = src.split(), hyp.split()
        if len(sw) != len(hw):
            return text
        return " ".join(h if unicodedata.normalize("NFC", _strip(h)) == s else s for s, h in zip(sw, hw))

    # ---- MT (MADLAD-400-3B-MT)
    def translate(self, text, tgt):
        import torch
        from transformers import T5ForConditionalGeneration, T5Tokenizer
        if self._mt is None:
            self._mt = (T5Tokenizer.from_pretrained("google/madlad400-3b-mt"),
                        T5ForConditionalGeneration.from_pretrained("google/madlad400-3b-mt", torch_dtype=torch.bfloat16).to(self.device).eval())
        tok, m = self._mt
        code = {"ibo": "ig", "yor": "yo", "hau": "ha", "eng": "en"}[tgt]
        out = []
        for sent in re.split(r"(?<=[.!?])\s+", text.strip()):
            if not sent:
                continue
            x = tok([f"<2{code}> {sent}"], return_tensors="pt").to(self.device)
            with torch.no_grad():
                y = m.generate(**x, max_new_tokens=256, num_beams=4)
            out.append(tok.decode(y[0], skip_special_tokens=True))
        return " ".join(out)

    def _seg(self, lang, s):
        s = normalize(s, lang)
        if lang in ("yor", "ibo") and unmarked(s):
            s = self.diacritise(s, lang)
        return f"<{lang}>{s}</{lang}>"

    def __call__(self, text, mode="codeswitch", lang="eng"):
        text = unicodedata.normalize("NFC", text).strip()
        if mode == "translate":
            if lang not in ("ibo", "yor", "hau"):
                raise ValueError("translate mode needs lang in ibo|yor|hau")
            plain = re.sub(r"</?\w{3}>", "", text)
            return self._seg(lang, self.translate(normalize(plain, "eng"), lang))
        if mode == "accent":
            plain = re.sub(r"</?\w{3}>", "", text)
            return f"<{lang}>{normalize(plain, 'eng')}</{lang}>"
        if mode == "native" or not SEG.search(text):
            return self._seg(lang, re.sub(r"</?\w{3}>", "", text))
        parts, pos = [], 0
        for m in SEG.finditer(text):
            if text[pos:m.start()].strip():
                parts.append(self._seg(lang, text[pos:m.start()]))
            parts.append(self._seg(m.group(1), m.group(2)))
            pos = m.end()
        if text[pos:].strip():
            parts.append(self._seg(lang, text[pos:]))
        return " ".join(parts)


if __name__ == "__main__":
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("text"); ap.add_argument("--mode", default="codeswitch"); ap.add_argument("--lang", default="eng")
    ap.add_argument("--diac_model", default=None); ap.add_argument("--cpu", action="store_true")
    a = ap.parse_args()
    print(FrontEnd(a.diac_model, device="cpu" if a.cpu else "cuda")(a.text, a.mode, a.lang))