import json
import torch
import torch.nn.functional as F
from tokenizer import CharTokenizer, PAD, BOS, EOS
from model import TranslitModel
from config import CFG

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


class Trie:
    def __init__(self):
        self.root = {}

    def insert(self, word, payload):
        node = self.root
        for c in word:
            node = node.setdefault(c, {})
        node.setdefault("$", []).append(payload)

    def prefix(self, pre, limit=400):
        node = self.root
        for c in pre:
            if c not in node:
                return []
            node = node[c]
        out, stack = [], [node]
        while stack and len(out) < limit:
            n = stack.pop()
            if "$" in n:
                out.extend(n["$"])
            for k, v in n.items():
                if k != "$":
                    stack.append(v)
        return out


class IME:
    def __init__(self, ckpt_dir="checkpoints_ime"):
        # ---------- lexicon trie ----------
        with open(f"{ckpt_dir}/lexicon.json", encoding="utf-8") as f:
            lex = json.load(f)
        self.trie = Trie()
        for g, targets in lex.items():
            for m, n in targets.items():
                self.trie.insert(g, (g, m, n))

        # ---------- neural fallback model ----------
        self.tok = CharTokenizer.load(f"{ckpt_dir}/vocab.json")
        self.model = TranslitModel(
            len(self.tok), CFG.d_model, CFG.nhead,
            CFG.num_layers, CFG.dim_ff,
            dropout=0.0, max_len=64,
        ).to(DEVICE)
        sd = torch.load(f"{ckpt_dir}/best.pt", map_location=DEVICE)["model"]
        sd = {k.replace("_orig_mod.", ""): v for k, v in sd.items()}
        self.model.load_state_dict(sd)
        self.model.eval()

    # ---------- dictionary path (trie prefix lookup + noise filtering) ----------
    def _dict_suggest(self, q, k):
        hits = self.trie.prefix(q, limit=400)
        exact = [h for h in hits if h[0] == q]
        pref  = [h for h in hits if h[0] != q]

        # exact-match noise filter: if a dominant exact form exists,
        # drop rare "alternatives" (misaligned one-off pairs)
        if exact:
            top = max(h[2] for h in exact)
            if top >= 5:
                exact = [h for h in exact if h[2] >= max(2, top * 0.02)]
        exact.sort(key=lambda h: -h[2])

        # prefix completions: Aksharantar has millions of n=1 glued
        # compounds (e.g. "entesuhruthaya"). Keep a completion only if it
        # was seen more than once OR is not much longer than the query,
        # then rank by frequency and shortness.
        pref = [h for h in pref if h[2] >= 2 or len(h[0]) <= len(q) + 6]
        pref.sort(key=lambda h: (-h[2], len(h[0])))

        out, seen = [], set()
        for g, m, n in exact + pref:
            if m not in seen:
                seen.add(m)
                out.append(m)
            if len(out) == k:
                break
        return out

    # ---------- neural path (beam search) ----------
    @torch.no_grad()
    def _model_suggest(self, q, k=3, beam=8, max_len=48):
        src = torch.tensor([self.tok.encode(q)], device=DEVICE)
        beams = [(torch.tensor([[BOS]], device=DEVICE), 0.0, False)]
        for _ in range(max_len):
            if all(f for _, _, f in beams):
                break
            cands = []
            for ids, lp, fin in beams:
                if fin:
                    cands.append((ids, lp, True))
                    continue
                logits = self.model(src, ids)[0, -1]
                logp = F.log_softmax(logits.float(), -1)
                tl, ti = logp.topk(beam)
                for l, ix in zip(tl.tolist(), ti.tolist()):
                    nids = torch.cat(
                        [ids, torch.tensor([[ix]], device=DEVICE)], 1)
                    cands.append((nids, lp + l, ix == EOS))
            cands.sort(key=lambda c: c[1] / (c[0].size(1) ** 0.7),
                       reverse=True)
            beams = cands[:beam]

        out, seen = [], set()
        for ids, _, _ in beams:
            s = self.tok.decode(ids[0].tolist())
            if s and s not in seen:
                seen.add(s)
                out.append(s)
            if len(out) == k:
                break
        return out

    # ---------- public API ----------
# ---------- public API ----------
    def suggest(self, manglish_word, k=5):
        q = manglish_word.strip()
        if not q:
            return []
        if self._is_passthrough(q):
            return self._number_suggestions(q)
    
        ql = q.lower()
        dict_hits = self._dict_suggest(ql, k)
    
        # is there an exact-length dictionary match? (romanization == input)
        has_exact = any(h[0] == ql for h in self.trie.prefix(ql, limit=50))
    
        if not has_exact:
            # no bare-form entry -> the model's direct transliteration is
            # usually the clean word the user wants (e.g. dhoni -> ധോണി)
            model_hits = self._model_suggest(ql, k=2)
            merged = []
            for s in model_hits + dict_hits:
                if s not in merged:
                    merged.append(s)
            return merged[:k]
    
        return dict_hits[:k]

    @staticmethod
    def _is_passthrough(q):
        # token is all digits / punctuation / has no latin letters to transliterate
        return not any(c.isalpha() and ord(c) < 128 for c in q)

    # Western -> Malayalam digit map
    _ML_DIGITS = str.maketrans("0123456789", "൦൧൨൩൪൫൬൭൮൯")

    def _number_suggestions(self, q):
        out = [q]                                  # keep as-is (10)
        if any(c.isdigit() for c in q):
            ml = q.translate(self._ML_DIGITS)      # malayalam numerals (൧൦)
            if ml != q:
                out.append(ml)
        return out


if __name__ == "__main__":
    ime = IME()
    tests = [
        "thi", "thila",                          # prefix -> dict completions
        "amma", "ente", "adukkala", "veedu",     # common words
        "krithyamaayi",
        "njan", "nammal", "keralam",             # now covered by aksharantar
        "thiruvananthapuram",
        "blockchain", "kunjava",                 # OOV / rare -> model fallback
    ]
    for w in tests:
        print(f"{w:22s} -> {ime.suggest(w)}")