"""Per-word language ID + language-mix estimate for a user turn (and for a reply).
Reuses the Dynamic Accent release tagger (rule-based lexicons/diacritics/Pidgin markers, optional GlotLID clause scores).
estimate("Ndewo, how are you doing today?") ->
  {"mix": {"eng": 0.67, "ibo": 0.33}, "matrix": "eng", "pattern": ["greeting:ibo"], "words": [...], "labels": [...]}"""
import re, unicodedata
from tagger import Tagger, fold

LANGS = ["eng", "pcm", "ibo", "yor", "hau"]
NAMES = {"eng": "English", "pcm": "Nigerian Pidgin", "ibo": "Igbo", "yor": "Yoruba", "hau": "Hausa"}
# greeting / closing / discourse lexicon (diacritics folded) -> language
GREET = {"ibo": ["nde wo", "ndewoo", "ndeewo", "kedu", "nnoo", "ututu oma", "ehihie oma", "mgbede oma", "ka chi fo", "kedu ka i mere", "i bola chi"],
         "yor": ["e kaaro", "e kaasan", "e kurole", "e ku irole", "bawo", "se daadaa ni", "e nle", "pele", "e kaabo", "e ku abo", "ekaaro", "ekaasan", "ekaro", "e karo", "e kaaro o"],
         "hau": ["sannu", "sanu", "sanno", "baka da yamma", "baka da safiya", "barka", "barka da safiya", "barka da yamma", "ina kwana", "ina wuni", "yaya dai", "lafiya lau", "maraba"],
         "pcm": ["how far", "how body", "wetin dey", "how you dey", "how una dey", "i hail"]}
CLOSE = {"ibo": ["daalu", "dalu", "dalo", "daalo", "dalow", "ka o di", "ka emesia", "chukwu gozie"], "yor": ["e se", "o se", "o dabo", "e seun", "ese", "odabo"],
         "hau": ["na gode", "sai an jima", "sai anjima", "nagode", "mun gode", "na godi", "nagodi"], "pcm": ["we go see", "thank you well well", "make we yarn later"]}
DISC = {"pcm": ["abeg", "abi", "sha", "o", "na so", "omo", "wahala"], "yor": ["sha", "o", "ehen", "abi"], "ibo": ["biko", "nna", "nne", "chai", "ehen", "ngwa"],
        "hau": ["to", "kai", "wallahi", "haba"]}

_T = {}
def tagger(lid="glotlid"):
    if lid not in _T:
        try:
            _T[lid] = Tagger(lid if lid != "none" else None)
        except Exception as e:  # no fasttext / GlotLID: rules only
            print("tagger: GlotLID unavailable, rules only:", repr(e)[:120]); _T[lid] = Tagger(None)
    return _T[lid]

def _find(fs, phrases):
    for p in sorted(phrases, key=len, reverse=True):
        if re.search(r"(^|\s)" + re.escape(p) + r"(\s|$)", fs):
            return p
    return None

PCM_MARK = {"dey", "don", "abeg", "wetin", "una", "gree", "wahala", "sabi", "comot", "chop", "oya", "abi", "shey", "wey", "dem", "wan",
            "pikin", "oga", "sef", "jare", "naim", "dis", "dat", "wahala", "kpatakpata", "ehn", "sha", "nor", "abeg", "jor", "ehen", "na"}
PCM_BI = {("no", "gree"), ("no", "be"), ("e", "no"), ("make", "we"), ("how", "far"), ("na", "im"), ("i", "go"), ("you", "go"), ("we", "go"), ("e", "go"),
          ("no", "go"), ("i", "fit"), ("you", "fit"), ("e", "don"), ("i", "don"), ("my", "guy"), ("na", "so"), ("well", "well"), ("no", "dey")}

def pidgin_relabel(words, labs, text):
    """sentence-level Pidgin: a sentence of English-labelled words with Pidgin markers is Pidgin (the word tagger only flips the marker's clause)."""
    if not labs: return labs
    labs = list(labs); f = [fold(w).strip(".,!?;:'\"") for w in words]
    ends = [i for i, w in enumerate(words) if re.search(r"[.!?]$", w)] + [len(words) - 1]
    s0 = 0
    for e in sorted(set(ends)):
        idx = list(range(s0, e + 1)); s0 = e + 1
        if not idx: continue
        hits = sum(1 for i in idx if f[i] in PCM_MARK or f[i] == "omo") + sum(1 for a, b in zip(idx, idx[1:]) if (f[a], f[b]) in PCM_BI and not ((f[a], f[b]) == ("how", "far") and a != idx[0]))
        hits += 1 if f[idx[-1]] == "o" and len(idx) > 2 else 0
        lat = [i for i in idx if labs[i] in ("eng", "pcm")]
        if hits == 0 and len(lat) > 4 and all(labs[i] == "pcm" for i in lat):
            for i in lat: labs[i] = "eng"
            continue
        if lat and (hits >= 2 or (hits >= 1 and len(idx) <= 8)) and sum(labs[i] == "pcm" for i in lat) < len(lat):
            for i in lat: labs[i] = "pcm"
    return labs

def window_relabel(words, labs, t, win=4, thr=0.5):
    """ASR text has no punctuation/diacritics, so the clause-level tagger sees one long 'English' clause. Slide a GlotLID window
    over the words and let confident window votes override the rule labels (rule label keeps a weight-1 vote)."""
    from collections import Counter
    n = len(words)
    if t.lid is None or n < 3: return labs
    votes = [Counter({labs[i]: 1.0}) for i in range(n)]
    for s0 in range(0, max(1, n - win + 1)):
        sc = t._lid_scores(" ".join(words[s0:s0 + win])); l = max(sc, key=sc.get)
        if sc[l] >= thr:
            for i in range(s0, min(n, s0 + win)): votes[i][l] += sc[l]
    return [v.most_common(1)[0][0] for v in votes]

def estimate(text, lid="glotlid", prior=None, prior_w=0.0, asr=False):
    """prior: a previous mix dict (conversation-level style), blended with weight prior_w."""
    t = tagger(lid)
    toks, labs = t.word_langs(text) if text.strip() else ([], [])
    words = [w for _, _, w in toks]
    if asr: labs = window_relabel(words, list(labs), t)
    labs = pidgin_relabel(words, labs, text)
    n = max(1, len(labs))
    mix = {l: labs.count(l) / n for l in LANGS if labs.count(l)}
    if prior and prior_w > 0:
        keys = set(mix) | set(prior)
        mix = {k: (1 - prior_w) * mix.get(k, 0) + prior_w * prior.get(k, 0) for k in keys}
        s = sum(mix.values()) or 1; mix = {k: v / s for k, v in mix.items() if v / s >= 0.03}
    fs = " " + re.sub(r"[^\w\s']", " ", fold(text)) + " "
    fs = re.sub(r"\s+", " ", fs)
    pat = []
    head = " ".join(fs.split()[:5]); tail = " ".join(fs.split()[-5:])
    for l in ("ibo", "yor", "hau", "pcm"):
        g = _find(head, GREET[l]); c = _find(tail, CLOSE[l])
        if g: pat.append(f"greeting:{l}")
        if c: pat.append(f"closing:{l}")
        d = [p for p in DISC[l] if re.search(r"(^|\s)" + re.escape(p) + r"(\s|$)", fs)]
        if d and l in mix: pat.append(f"discourse:{l}")
    # ASR often spells a greeting/closing so that the word tagger calls it English: give the matched words to that language
    toks_f = fs.split()
    for p in pat:
        k, _, l = p.partition(":")
        if k in ("greeting", "closing") and mix.get(l, 0) < 0.1 and labs:
            phr = _find(head if k == "greeting" else tail, (GREET if k == "greeting" else CLOSE)[l]) or ""
            share = min(0.5, len(phr.split()) / max(1, len(labs)))
            m2 = max(mix, key=mix.get)
            mix[m2] = max(0.0, mix[m2] - share); mix[l] = mix.get(l, 0) + share
            mix = {a: b for a, b in mix.items() if b > 0.01}
    # switch style: count switches between labelled runs
    sw = sum(1 for a, b in zip(labs, labs[1:]) if a != b)
    if len(mix) > 1:
        pat.append("intra-sentential" if sw >= 3 else "inter-sentential")
    matrix = max(mix, key=mix.get) if mix else "eng"
    return dict(mix={k: round(v, 2) for k, v in sorted(mix.items(), key=lambda x: -x[1])}, matrix=matrix, pattern=pat,
                words=words, labels=labs, n_words=len(labs), switches=sw)

def describe(est):
    """natural-language mix instruction for the LLM."""
    mix, pat = est["mix"], est["pattern"]
    if not mix:
        return "Reply in English."
    parts = [f"{NAMES[l]} about {round(100 * p / 5) * 5}%" for l, p in mix.items()]
    s = f"Match the user's language mix: {', '.join(parts)} of the words. The main (matrix) language is {NAMES[est['matrix']]}."
    for p in pat:
        k, _, l = p.partition(":")
        if k == "greeting": s += f" The user greeted in {NAMES[l]}: open with a short {NAMES[l]} greeting."
        if k == "closing": s += f" The user closed in {NAMES[l]}: you may close in {NAMES[l]}."
        if k == "discourse": s += f" Sprinkle a {NAMES[l]} discourse word or two the way the user did."
    if "intra-sentential" in pat: s += " The user switches inside sentences; do the same naturally."
    if len(mix) == 1: s += f" Reply only in {NAMES[est['matrix']]}."
    return s

def tvd(a, b):
    k = set(a) | set(b)
    return 0.5 * sum(abs(a.get(x, 0) - b.get(x, 0)) for x in k)

if __name__ == "__main__":
    import sys, json
    for s in sys.argv[1:] or ["Ndewo! How are you doing today? I wan ask you something.", "Ẹ káàárọ̀, ṣé dáadáa ni? I need help with my account."]:
        e = estimate(s); print(json.dumps({k: e[k] for k in ("mix", "matrix", "pattern")}, ensure_ascii=False), "|", describe(e))