"""AksharaLP v3 runtime (Python reference): v2 candidates + reranker (BhcTranslit2), then sentence-level rescoring of the
candidate lattice with (a) the v2 bigram LM, (b) a tiny neural sentence LM over Devanagari subwords (bhc/nlm.py) and
(c) agreement/grammar consistency features (bhc/agree.py). HARD guardrail: the search only ever picks a word from the
candidate list produced for that typed roman token (lattice), never anything else.
Lattice slot = list of (cand, channel_logsoftmax, wid); fixed slots (punctuation, digits, Devanagari) pass through."""
import math, os, re, json
from bhc.translit2 import BhcTranslit2, LATIN
from bhc.text import TOKEN_RE

DEV_RE = re.compile(r"[\u0900-\u0963\u0971-\u097F]+$")
NUM_RE = re.compile(r"[0-9\u0966-\u096F]+$")
V3DEF = dict(lat_k=12, beam_s=16, w_nlm=0.0, w_agr=0.0, w_agn=0.0, nlm_file="nlm.q8", nlm_margin=99.0, nlm_k=12, rec_n=3)


class BhcTranslit3(BhcTranslit2):
    def __init__(self, asset_dir, threads=1, use_dict=True, nlm=None, **over):
        super().__init__(asset_dir, threads=threads, use_dict=use_dict, **{k: v for k, v in over.items() if k not in V3DEF})
        for k, v in V3DEF.items(): self.p.setdefault(k, v)
        self.p.update({k: v for k, v in over.items() if k in V3DEF})
        self.nlm = nlm
        if self.nlm is None:
            d = self.dict_dir or asset_dir; f = f"{d}/{self.p['nlm_file']}"
            if os.path.exists(f):
                from bhc.nlm import NLM
                self.nlm = NLM.load(f)
        self.name = "bhc-xlit3"

    def build_lattice(self, slots, k=None):
        k = k or self.p["lat_k"]; xs = [t for kind, t in slots if kind == "x"]
        cl = iter(self.cand_feats(xs)) if xs else iter(()); lat = []
        for kind, t in slots:
            if kind == "x":
                rows = next(cl); sc = [(s, self.score(fv), wid) for s, fv, wid in rows]
                if sc:
                    mx = max(x[1] for x in sc); lse = mx + math.log(sum(math.exp(x[1] - mx) for x in sc))
                    sc = sorted(((s, v - lse, wid) for s, v, wid in sc), key=lambda x: -x[1])[:k]
                lat.append(("x", sc or [(t, 0.0, None)]))
            else:
                lat.append(("fix", [(t, 0.0, self.lm.id.get(t, -1) if DEV_RE.match(t) else -2)]))
        return lat

    def search(self, lat):
        """Beam search over the lattice. score = sum w_ch_s*channel + w_lm*bigram(-oov) + w_nlm*nlm_logp + w_agr*agree."""
        from bhc.agree import agree_feats
        p = self.p; lm = self.lm; use_bi = p["use_lm"] and lm.bi; nlm = self.nlm if p["w_nlm"] else None
        # beam item: (score, prev_wid, words(list), nlm_state)
        beams = [(0.0, None, [], nlm.init_state() if nlm else None)]
        for kind, opts in lat:
            if kind == "fix":
                t, _, wid = opts[0]
                if wid == -2:          # punctuation / digits / latin: no LM context change except <n> for numbers in the NLM
                    if nlm and NUM_RE.match(t):
                        st = nlm.advance_many([b[3] for b in beams], ["<n>"])
                        beams = [(s, h, ch + [t], st[i][0]) for i, (s, h, ch, _) in enumerate(beams)]
                    else:
                        beams = [(s, h, ch + [t], ns) for s, h, ch, ns in beams]
                    continue
                opts = [(t, 0.0, wid)]
            opts = opts[:p["lat_k"]]
            if nlm:
                cand = [o for o in opts if o[1] >= opts[0][1] - p["nlm_margin"]][:p["nlm_k"]]
                lp, nst = nlm.score_words([b[3] for b in beams], [c[0] for c in cand])   # lp[i][j], nst[i][j]
            else:
                cand = opts
            best = {}
            for i, (s, h, ch, ns) in enumerate(beams):
                for j, (c, csc, wid) in enumerate(cand):
                    w = None if wid in (None, -1) else wid
                    tot = s + p["w_ch_s"] * csc
                    if use_bi: tot += p["w_lm"] * (lm.logp(h, w) - (p["oov_pen_s"] if w is None else 0.0))
                    if nlm: tot += p["w_nlm"] * lp[i][j]
                    words = ch + [c]
                    if p["w_agr"] or p["w_agn"]:
                        ap_, an_ = agree_feats(words, len(words) - 1); tot += p["w_agr"] * ap_ - p["w_agn"] * an_
                    key = c if not (nlm or p["w_agr"] or p["w_agn"]) else tuple(words[-p["rec_n"]:])   # v2-exact recombination for the bigram-only case
                    if key not in best or tot > best[key][0]: best[key] = (tot, w, words, nst[i][j] if nlm else None)
            beams = sorted(best.values(), key=lambda x: -x[0])[:p["beam_s"]]
        def final(b):
            f = b[0]
            if use_bi: f += p["w_lm"] * lm.logp(b[1], lm.N)
            if nlm: f += p["w_nlm"] * nlm.eos_logp(b[3])
            return f
        return max(beams, key=final)[2]

    def decode_slots(self, slots):
        return self.search(self.build_lattice(slots))