"""BhashaCore Track B candidate generator: tiny character-level Transformer (Roman -> Devanagari).

Pre-LN encoder/decoder, learned positions, tied decoder input/output embeddings. Inference = batched beam
search over all words of a sentence at once, with an incremental (cached) decoder, optionally constrained to a
Devanagari lexicon trie. Weights are stored as INT8 (per-output-row symmetric scales) + fp16 for small vectors
and de-quantised to fp32 at load time (compute is fp32; storage is INT8).
"""
import io, gzip, json, math, os
import numpy as np
import torch, torch.nn as nn, torch.nn.functional as F

ROMAN = "abcdefghijklmnopqrstuvwxyz"
PAD, BOS, EOS, UNK = 0, 1, 2, 3
DEVA_CHARS = [chr(c) for c in range(0x0900, 0x0980)]  # full block (128), most unused but cheap
SRC_VOCAB = ["<pad>", "<s>", "</s>", "<unk>"] + list(ROMAN)
TGT_VOCAB = ["<pad>", "<s>", "</s>", "<unk>"] + DEVA_CHARS
S2I = {c: i for i, c in enumerate(SRC_VOCAB)}; T2I = {c: i for i, c in enumerate(TGT_VOCAB)}
MAXLEN = 32


def enc_src(w): return [S2I.get(c, UNK) for c in w.lower()[:MAXLEN - 1]] + [EOS]
def enc_tgt(w): return [T2I.get(c, UNK) for c in w[:MAXLEN - 2]] + [EOS]
def dec_tgt(ids): return "".join(TGT_VOCAB[i] for i in ids if i > UNK)


class MHA(nn.Module):
    def __init__(s, d, h):
        super().__init__(); s.h = h; s.dk = d // h
        s.q = nn.Linear(d, d); s.kv = nn.Linear(d, 2 * d); s.o = nn.Linear(d, d)

    def forward(s, x, mem, mask=None, cache=None):
        B, T, D = x.shape
        q = s.q(x).view(B, T, s.h, s.dk).transpose(1, 2)
        if cache is not None and "k" in cache and cache.get("static"):
            k, v = cache["k"], cache["v"]
        else:
            k, v = s.kv(mem).view(B, -1, 2, s.h, s.dk).permute(2, 0, 3, 1, 4)
            if cache is not None:
                if "k" in cache:  # self-attn incremental
                    k = torch.cat([cache["k"], k], 2); v = torch.cat([cache["v"], v], 2)
                cache["k"], cache["v"] = k, v
        if mask is not None and mask.dtype != q.dtype: mask = mask.to(q.dtype)
        a = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
        return s.o(a.transpose(1, 2).reshape(B, T, D))


class EncLayer(nn.Module):
    def __init__(s, d, h, ff, p):
        super().__init__(); s.n1 = nn.LayerNorm(d); s.att = MHA(d, h); s.n2 = nn.LayerNorm(d)
        s.ff = nn.Sequential(nn.Linear(d, ff), nn.GELU(), nn.Linear(ff, d)); s.dr = nn.Dropout(p)

    def forward(s, x, mask):
        y = s.n1(x); x = x + s.dr(s.att(y, y, mask)); return x + s.dr(s.ff(s.n2(x)))


class DecLayer(nn.Module):
    def __init__(s, d, h, ff, p):
        super().__init__(); s.n1 = nn.LayerNorm(d); s.sa = MHA(d, h); s.n2 = nn.LayerNorm(d); s.ca = MHA(d, h)
        s.n3 = nn.LayerNorm(d); s.ff = nn.Sequential(nn.Linear(d, ff), nn.GELU(), nn.Linear(ff, d)); s.dr = nn.Dropout(p)

    def forward(s, x, mem, smask, mmask, cache=None):
        y = s.n1(x); x = x + s.dr(s.sa(y, y, smask, None if cache is None else cache["sa"]))
        x = x + s.dr(s.ca(s.n2(x), mem, mmask, None if cache is None else cache["ca"]))
        return x + s.dr(s.ff(s.n3(x)))


class XlitTransformer(nn.Module):
    def __init__(s, d=128, h=4, enc=3, dec=2, ff=384, p=0.1):
        super().__init__(); s.cfg = dict(d=d, h=h, enc=enc, dec=dec, ff=ff, p=p)
        s.semb = nn.Embedding(len(SRC_VOCAB), d, padding_idx=PAD); s.temb = nn.Embedding(len(TGT_VOCAB), d, padding_idx=PAD)
        s.spos = nn.Embedding(MAXLEN, d); s.tpos = nn.Embedding(MAXLEN, d)
        s.enc = nn.ModuleList(EncLayer(d, h, ff, p) for _ in range(enc)); s.dec = nn.ModuleList(DecLayer(d, h, ff, p) for _ in range(dec))
        s.en = nn.LayerNorm(d); s.dn = nn.LayerNorm(d); s.scale = math.sqrt(d)
        for e in (s.semb, s.temb): nn.init.normal_(e.weight, 0, d ** -0.5)
        nn.init.normal_(s.spos.weight, 0, 0.02); nn.init.normal_(s.tpos.weight, 0, 0.02)

    def encode(s, src):
        B, S = src.shape; pad = src == PAD
        x = s.semb(src) * s.scale + s.spos(torch.arange(S, device=src.device))[None]
        m = torch.zeros(B, 1, 1, S, device=src.device).masked_fill(pad[:, None, None, :], float("-inf"))
        for l in s.enc: x = l(x, m)
        return s.en(x), m

    def logits(s, x): return x @ s.temb.weight.t()

    def forward(s, src, tgt_in):
        mem, mm = s.encode(src); B, T = tgt_in.shape
        x = s.temb(tgt_in) * s.scale + s.tpos(torch.arange(T, device=tgt_in.device))[None]
        causal = torch.full((T, T), float("-inf"), device=tgt_in.device).triu(1)
        for l in s.dec: x = l(x, mem, causal, mm)
        return s.logits(s.dn(x))

    @torch.no_grad()
    def step(s, tok, pos, mem, mm, caches):
        x = s.temb(tok) * s.scale + s.tpos.weight[pos][None, None]
        for l, c in zip(s.dec, caches): x = l(x, mem, None, mm, c)
        return F.log_softmax(s.logits(s.dn(x))[:, -1], -1)

    def n_params(s): return sum(p.numel() for p in s.parameters())


# ------------------------------------------------------------------ lexicon trie (for constrained decoding)
class Trie:
    """Character trie over a word list, stored as flat arrays; node 0 = root. child(node, tok) lookup via dict."""
    def __init__(s, words):
        s.kids = [{}]; s.final = [False]
        for w in words:
            n = 0
            for c in w:
                t = T2I.get(c)
                if t is None: break
                nx = s.kids[n].get(t)
                if nx is None:
                    nx = len(s.kids); s.kids[n][t] = nx; s.kids.append({}); s.final.append(False)
                n = nx
            else:
                s.final[n] = True


# ------------------------------------------------------------------ beam search
@torch.no_grad()
def beam_search(model, words, beam=8, nbest=8, trie=None, maxlen=None, constrain=None):
    """Return per word a list of (devanagari, logprob) best first. All words decoded in one batch.
    trie: constrain outputs to lexicon words (EOS allowed only at final nodes); constrain[i] selects which rows
    are constrained (default: all rows when a trie is given)."""
    if not words: return []
    src = [enc_src(w) for w in words]; S = max(map(len, src))
    src = torch.tensor([x + [PAD] * (S - len(x)) for x in src])
    mem, mm = model.encode(src); W = len(words); K = beam
    mem = mem.repeat_interleave(K, 0); mm = mm.repeat_interleave(K, 0)
    caches = [{"sa": {}, "ca": {"static": False}} for _ in model.dec]
    # first step: only beam 0 alive
    score = torch.full((W, K), float("-inf")); score[:, 0] = 0
    tok = torch.full((W * K, 1), BOS); seqs = [[[] for _ in range(K)] for _ in range(W)]
    if trie and constrain is None: constrain = [True] * len(words)
    nodes = [[0 if constrain[i] else None] * K for i in range(W)] if trie else None
    done = [[] for _ in range(W)]; V = len(TGT_VOCAB)
    L = maxlen or min(MAXLEN - 1, S + 8)
    for t in range(L):
        lp = model.step(tok, t, mem, mm, caches)
        if t == 0:
            for c in caches: c["ca"]["static"] = True
        lp = lp.view(W, K, V)
        if trie:
            mask = torch.zeros((W, K, V))
            for i in range(W):
                if not constrain[i]: continue
                mask[i] = float("-inf")
                for k in range(K):
                    n = nodes[i][k]
                    if n is None: continue
                    ks = list(trie.kids[n].keys())
                    if ks: mask[i, k, ks] = 0
                    if trie.final[n]: mask[i, k, EOS] = 0
            lp = lp + mask
        lp[:, :, PAD] = float("-inf"); lp[:, :, BOS] = float("-inf"); lp[:, :, UNK] = float("-inf")
        tot = (score[:, :, None] + lp).view(W, -1)
        top, idx = tot.topk(2 * K, -1)
        new_score = torch.full((W, K), float("-inf")); reorder = torch.zeros(W, K, dtype=torch.long); new_tok = torch.full((W, K), EOS)
        new_seqs = [[[] for _ in range(K)] for _ in range(W)]; new_nodes = [[None] * K for _ in range(W)] if trie else None
        alive_any = False
        for i in range(W):
            j = 0
            for sc, ix in zip(top[i].tolist(), idx[i].tolist()):
                if sc == float("-inf"): break
                b, v = divmod(ix, V)
                if v == EOS:
                    if len(done[i]) < nbest: done[i].append((dec_tgt(seqs[i][b]), sc))
                    continue
                if j < K:
                    new_score[i, j] = sc; reorder[i, j] = i * K + b; new_tok[i, j] = v
                    new_seqs[i][j] = seqs[i][b] + [v]
                    if trie: new_nodes[i][j] = trie.kids[nodes[i][b]][v] if constrain[i] else None
                    j += 1
            # stop this word if enough finished and best alive can't beat worst finished
            if len(done[i]) >= nbest and (j == 0 or new_score[i, 0].item() < min(x[1] for x in done[i])):
                new_score[i] = float("-inf")
            elif j: alive_any = True
        if not alive_any: break
        flat = reorder.view(-1)
        for c in caches:
            c["sa"]["k"] = c["sa"]["k"][flat]; c["sa"]["v"] = c["sa"]["v"][flat]
        score = new_score; seqs = new_seqs; nodes = new_nodes
        tok = torch.cat([tok[flat], new_tok.view(-1, 1)], 1)[:, -1:]
    out = []
    for i in range(W):
        d = {}
        for s_, sc in done[i]:
            if s_ not in d or sc > d[s_]: d[s_] = sc
        out.append(sorted(d.items(), key=lambda x: -x[1])[:nbest])
    return out


# ------------------------------------------------------------------ INT8 export / load
def export_int8(model, path):
    """Linear/embedding matrices -> int8 with per-row fp16 scale; vectors (bias, LN) -> fp16. Single .npz-like blob."""
    sd = model.state_dict(); meta = {"cfg": model.cfg, "tensors": []}; buf = io.BytesIO()
    for k, v in sd.items():
        v = v.float().numpy()
        if v.ndim == 2:
            sc = np.abs(v).max(1, keepdims=True) / 127.0; sc[sc == 0] = 1
            q = np.clip(np.round(v / sc), -127, 127).astype(np.int8)
            meta["tensors"].append([k, list(v.shape), "i8", buf.tell()]); buf.write(q.tobytes()); buf.write(sc.astype(np.float16).tobytes())
        else:
            meta["tensors"].append([k, list(v.shape), "f16", buf.tell()]); buf.write(v.astype(np.float16).tobytes())
    hdr = json.dumps(meta, separators=(",", ":")).encode()
    with open(path, "wb") as f:
        f.write(len(hdr).to_bytes(4, "little")); f.write(hdr); f.write(buf.getvalue())


def load_int8(path):
    b = open(path, "rb").read(); n = int.from_bytes(b[:4], "little"); meta = json.loads(b[4:4 + n]); data = b[4 + n:]
    m = XlitTransformer(**meta["cfg"]); sd = {}
    for k, shape, kind, off in meta["tensors"]:
        cnt = int(np.prod(shape))
        if kind == "i8":
            q = np.frombuffer(data, np.int8, cnt, off).reshape(shape).astype(np.float32)
            sc = np.frombuffer(data, np.float16, shape[0], off + cnt).astype(np.float32)[:, None]
            sd[k] = torch.from_numpy(q * sc)
        else:
            sd[k] = torch.from_numpy(np.frombuffer(data, np.float16, cnt, off).astype(np.float32).reshape(shape))
    m.load_state_dict(sd); m.eval(); return m


@torch.no_grad()
def score_pairs(model, srcs, tgts):
    """Teacher-forced log P(tgt | src) for each pair (batched)."""
    if not srcs: return []
    s = [enc_src(w) for w in srcs]; t = [enc_tgt(w) for w in tgts]; S = max(map(len, s)); T = max(map(len, t))
    src = torch.tensor([x + [PAD] * (S - len(x)) for x in s]); tin = torch.tensor([[BOS] + x[:-1] + [PAD] * (T - len(x)) for x in t])
    tout = torch.tensor([x + [PAD] * (T - len(x)) for x in t])
    lp = F.log_softmax(model(src, tin), -1).gather(-1, tout[..., None])[..., 0]
    return (lp * (tout != PAD)).sum(-1).tolist()


@torch.no_grad()
def beam_search_fast(model, words, beam=8, nbest=8, trie=None, maxlen=None, constrain=None):
    """Same search as beam_search (identical results), with the per-candidate bookkeeping vectorised:
    hypotheses kept as a token-history tensor; only EOS hits and trie-node updates touch Python."""
    if not words: return []
    src = [enc_src(w) for w in words]; S = max(map(len, src))
    src = torch.tensor([x + [PAD] * (S - len(x)) for x in src])
    mem, mm = model.encode(src); W = len(words); K = beam; V = len(TGT_VOCAB)
    mem = mem.repeat_interleave(K, 0); mm = mm.repeat_interleave(K, 0)
    caches = [{"sa": {}, "ca": {"static": False}} for _ in model.dec]
    score = torch.full((W, K), float("-inf")); score[:, 0] = 0
    tok = torch.full((W * K, 1), BOS); hist = torch.zeros((W * K, 0), dtype=torch.long)
    if trie and constrain is None: constrain = [True] * W
    crow = [i for i in range(W) if trie and constrain[i]]
    nodes = {i: [0] * K for i in crow}
    if trie and not hasattr(trie, "_allow"): trie._allow = {}
    done = [[] for _ in range(W)]; alive = [True] * W
    L = maxlen or min(MAXLEN - 1, S + 8); NEG = float("-inf")
    for t in range(L):
        lp = model.step(tok, t, mem, mm, caches)
        if t == 0:
            for c in caches: c["ca"]["static"] = True
        lp = lp.view(W, K, V)
        if crow:
            mask = torch.full((len(crow), K, V), NEG)
            for j, i in enumerate(crow):
                for k in range(K):
                    n = nodes[i][k]
                    if n is None: continue
                    a = trie._allow.get(n)
                    if a is None:
                        a = list(trie.kids[n].keys()) + ([EOS] if trie.final[n] else []); a = torch.tensor(a, dtype=torch.long); trie._allow[n] = a
                    if len(a): mask[j, k, a] = 0
            lp[crow] = lp[crow] + mask
        lp[:, :, PAD] = NEG; lp[:, :, BOS] = NEG; lp[:, :, UNK] = NEG
        tot = (score[:, :, None] + lp).view(W, -1)
        top, idx = tot.topk(2 * K, -1)
        b = idx // V; v = idx % V; fin = torch.isfinite(top)
        eos = (v == EOS) & fin
        if eos.any():
            for r, c in eos.nonzero().tolist():
                if alive[r] and len(done[r]) < nbest:
                    done[r].append((dec_tgt(hist[r * K + b[r, c]].tolist()), top[r, c].item()))
        ne = (v != EOS) & fin; rank = ne.long().cumsum(1) - 1; keep = ne & (rank < K)
        rr, cc = keep.nonzero(as_tuple=True); kk = rank[rr, cc]
        new_score = torch.full((W, K), NEG); new_score[rr, kk] = top[rr, cc]
        reorder = (torch.arange(W)[:, None] * K).repeat(1, K); reorder[rr, kk] = rr * K + b[rr, cc]
        new_tok = torch.full((W, K), EOS); new_tok[rr, kk] = v[rr, cc]
        if crow:
            nn_ = {i: [None] * K for i in crow}
            bl, vl, kl, rl = b[rr, cc].tolist(), v[rr, cc].tolist(), kk.tolist(), rr.tolist()
            for r, bb, vv, k in zip(rl, bl, vl, kl):
                if r in nn_: nn_[r][k] = trie.kids[nodes[r][bb]][vv]
            nodes = nn_
        nkeep = keep.sum(1).tolist(); best0 = new_score[:, 0].tolist(); any_alive = False
        for i in range(W):
            if not alive[i]: new_score[i] = NEG; continue
            if len(done[i]) >= nbest and (nkeep[i] == 0 or best0[i] < min(x[1] for x in done[i])): new_score[i] = NEG; alive[i] = False
            elif nkeep[i]: any_alive = True
            else: alive[i] = False
        if not any_alive: break
        flat = reorder.view(-1)
        for c in caches:
            c["sa"]["k"] = c["sa"]["k"][flat]; c["sa"]["v"] = c["sa"]["v"][flat]
        hist = torch.cat([hist[flat], new_tok.view(-1, 1)], 1)
        score = new_score; tok = new_tok.view(-1, 1)
    out = []
    for i in range(W):
        d = {}
        for s_, sc in done[i]:
            if s_ not in d or sc > d[s_]: d[s_] = sc
        out.append(sorted(d.items(), key=lambda x: -x[1])[:nbest])
    return out