"""Memory-light lexicon trie for constrained beam search over large (0.5-2M word) lexicons.
A node is a (lo, hi, depth) range of the sorted word list; children are found with bisect and cached lazily, so
only nodes actually visited by the beam are ever materialised. Same interface as xlit_model.Trie (kids[n], final[n])."""
import bisect
from bhc.xlit_model import T2I, TGT_VOCAB


class _Kids:
    def __init__(s, t): s.t = t
    def __getitem__(s, n): return s.t._children(n)


class _Final:
    def __init__(s, t): s.t = t
    def __getitem__(s, n):
        lo, hi, d = s.t.nodes[n]; return len(s.t.w[lo]) == d


class LazyTrie:
    def __init__(s, words):
        s.w = sorted(w for w in set(words) if w and all(c in T2I for c in w))
        s.nodes = [(0, len(s.w), 0)]; s.ids = {s.nodes[0]: 0}; s.cache = {}
        s.kids = _Kids(s); s.final = _Final(s)

    def _children(s, n):
        c = s.cache.get(n)
        if c is not None: return c
        lo, hi, d = s.nodes[n]; w = s.w; out = {}
        pos = lo
        if pos < hi and len(w[pos]) == d: pos += 1          # the word equal to the prefix sorts first
        while pos < hi:
            ch = w[pos][d]; pre = w[pos][:d + 1]
            end = bisect.bisect_left(w, pre[:-1] + chr(ord(ch) + 1), pos, hi)
            key = (pos, end, d + 1); nid = s.ids.get(key)
            if nid is None: nid = len(s.nodes); s.nodes.append(key); s.ids[key] = nid
            out[T2I[ch]] = nid; pos = end
        if len(s.cache) > 400000: s.cache.clear()
        s.cache[n] = out; return out