"""Compact Devanagari lexicon + pruned backoff word-bigram LM (BhashaCore Track B/C assets), format v2.

Word ids: words sorted by (unigram bucket = round(-2*ln count) i.e. frequency class, then alphabetical) so
ids are roughly frequency-ordered (small ids for frequent words -> small LM deltas), while the list is still
front-coded within each bucket. Unigram ln-prob is per bucket (implicit; no per-word bytes).
lex.bin : header | per-bucket (n_words varint, ln-prob f16) | front-coded list (1 byte/char = codepoint-0x0900,
          last suffix byte has bit 7 set; 0x7F = empty suffix)
lm.bin  : header | successor count per history (varint; history N = sentence start) | delta-coded successor
          ids (varint) | 4-bit codebook ln p(w|v) (16 centroids, f32) | 4-bit codebook ln backoff for histories
          with successors (16 centroids)
"""
import gzip, math, struct
import numpy as np

HK = 1 << 21   # in-memory bigram key = history * HK + successor (lexicons up to 2M words; file format unchanged)


def varints(vals):
    out = bytearray()
    for v in vals:
        v = int(v)
        while v >= 0x80: out.append((v & 0x7F) | 0x80); v >>= 7
        out.append(v)
    return bytes(out)


def read_varints(b, n, pos=0):
    out = [0] * n; i = 0
    while i < n:
        v = sh = 0
        while True:
            c = b[pos]; pos += 1; v |= (c & 0x7F) << sh; sh += 7
            if c < 0x80: break
        out[i] = v; i += 1
    return out, pos


def kmeans1d(x, k=16, it=30):
    x = np.asarray(x, np.float64); c = np.quantile(x, (np.arange(k) + 0.5) / k)
    for _ in range(it):
        a = np.abs(x[:, None] - c[None]).argmin(1)
        for j in range(k):
            if (a == j).any(): c[j] = x[a == j].mean()
    c = np.sort(c); a = np.abs(x[:, None] - c[None]).argmin(1) if len(x) < 2_000_000 else None
    return c.astype(np.float32), a.astype(np.uint8)


def pack4(a):
    a = np.r_[a, np.zeros(len(a) % 2, np.uint8)].astype(np.uint8); return (a[0::2] * 16 + a[1::2]).astype(np.uint8).tobytes()


def unpack4(b, n):
    a = np.frombuffer(b, np.uint8); return np.stack([a >> 4, a & 15], 1).reshape(-1)[:n]


def build(words_by_freq, counts_by_freq, bigram_pairs, bigram_counts, n_lex, maxbi, D=0.75, cmin=2, K=262144, alpha=False):
    """words_by_freq[i] has count id i+2 in bigram_pairs (0=boundary, 1=unk)."""
    N = n_lex; ws = words_by_freq[:N]; cs = counts_by_freq[:N]
    bk = [-int(round(2 * math.log(c))) for c in cs]
    if alpha:   # format v4: one global alphabetical front-coded list + 16-level (4-bit) uniform log-frequency band per word
        order = sorted(range(N), key=lambda i: ws[i])
    else:
        order = sorted(range(N), key=lambda i: (bk[i], ws[i]))
    nid = np.empty(N + 2, np.int64); nid[0] = N; nid[1] = -1
    for a, i in enumerate(order): nid[i + 2] = a
    tot = float(sum(cs)) * 1.02  # ~2% mass reserved for OOV words
    buckets = []
    for i in order:
        if not buckets or buckets[-1][0] != bk[i]: buckets.append([bk[i], 0, 0.0])
        buckets[-1][1] += 1; buckets[-1][2] += cs[i]
    b_lp = [math.log(s / n / tot) for _, n, s in buckets]
    unk_lp = math.log(0.02 / 50000)
    puni = np.empty(N + 2); puni[0] = math.log(0.04); puni[1] = unk_lp; j = 0
    if alpha:
        lc = np.log(np.array(cs, np.float64) / tot); NB = 16   # 16 uniform log-frequency bands (4 bits/word)
        ucb = np.linspace(lc.min(), lc.max(), NB).astype(np.float32)
        ua = np.abs(lc[:, None] - ucb[None].astype(np.float64)).argmin(1).astype(np.uint8)
        puni[2:] = ucb[ua]
    else:
        for (_, n, _), lp in zip(buckets, b_lp):
            for _ in range(n): puni[[order[j] + 2]] = lp; j += 1
    v = bigram_pairs // K; w = bigram_pairs % K; c = bigram_counts.astype(np.float64)
    inlex = (v - 2 < N) & (w - 2 < N)
    hist_tot = np.bincount(np.where(v - 2 < N, v, 1), weights=c, minlength=N + 2)
    keep = inlex & (v != 1) & (w != 1) & (c >= cmin)
    pk_all = np.log(np.maximum(c - D, 0.1) / np.maximum(hist_tot[np.minimum(v, N + 1)], 1))
    gain = np.where(keep, c * (pk_all - puni[np.minimum(w, N + 1)]), -np.inf)
    if keep.sum() > maxbi:
        thr = np.partition(gain, len(gain) - maxbi)[len(gain) - maxbi]; keep &= gain >= thr
    kv, kw, pk = v[keep], w[keep], pk_all[keep]
    av = nid[kv]; aw = nid[kw]   # history boundary 0 -> N ; successor boundary 0 -> N (= </s>)
    o = np.lexsort((aw, av)); av, aw, pk = av[o], aw[o], pk[o]
    cb, qa = kmeans1d(pk); pkd = cb[qa]
    cnt = np.bincount(av, minlength=N + 1)
    first = np.r_[True, av[1:] != av[:-1]]; delta = np.where(first, aw, aw - np.r_[0, aw[:-1]] - 1)
    # backoff for histories with successors (dequantised probs; unigram of successors in new-id space)
    pu_new = np.empty(N + 1); pu_new[nid[np.arange(2, N + 2)]] = puni[2:]; pu_new[N] = puni[0]
    s_bi = np.bincount(av, weights=np.exp(pkd), minlength=N + 1); s_un = np.bincount(av, weights=np.exp(pu_new[aw]), minlength=N + 1)
    hs = np.nonzero(cnt)[0]
    bow = np.log(np.maximum(1 - s_bi[hs], 1e-4) / np.maximum(1 - s_un[hs], 1e-4))
    cbb, qb = kmeans1d(bow)
    fc = bytearray(); prev = b""
    for i in order:
        bb = bytes(ord(ch) - 0x0900 for ch in ws[i]); k = 0
        while k < min(len(bb), len(prev), 127) and bb[k] == prev[k]: k += 1
        suf = bytearray(bb[k:]) or bytearray([0x7F]); suf[-1] |= 0x80; fc.append(k); fc += suf; prev = bb
    bh = varints([n for _, n, _ in buckets]) + np.array(b_lp, np.float16).tobytes()
    if alpha:   # v4: prefix-length stream and suffix stream stored separately (compresses ~8% better)
        P = bytearray(); S = bytearray(); prev = b""
        for i in order:
            bb = bytes(ord(ch) - 0x0900 for ch in ws[i]); k = 0
            while k < min(len(bb), len(prev), 127) and bb[k] == prev[k]: k += 1
            suf = bytearray(bb[k:]) or bytearray([0x7F]); suf[-1] |= 0x80; P.append(k); S += suf; prev = bb
        lex = b"LX4\0" + struct.pack("<IIfI", N, len(S), unk_lp, NB) + ucb.astype(np.float32).tobytes() + pack4(ua[order]) + bytes(P) + bytes(S)
    else:
        lex = struct.pack("<IIIf", N, len(buckets), len(fc), unk_lp) + bh + bytes(fc)
    b1, b2 = varints(cnt), varints(delta)
    lm = (struct.pack("<IIIII", N, len(av), len(b1), len(b2), len(hs)) + cb.tobytes() + cbb.tobytes() + b1 + b2
          + pack4(qa) + pack4(qb))
    return lex, lm, {"lex_words": N, "bigrams": int(len(av)), "histories": int(len(hs)), "buckets": len(buckets)}


class LexLM:
    def __init__(self, lex_path, lm_path=None):
        if lex_path is None:   # core-only (no dictionary pack): empty lexicon, no LM
            self.words = []; self.N = 0; self.id = {}; self.uni = []; self.unk = 0.0; self.bi = {}; self.bow = {}; return
        b = open(lex_path, "rb").read()
        if b[:4] == b"LX4\0":
            N, nfc, self.unk, NB = struct.unpack("<IIfI", b[4:20]); pos = 20; ucb = np.frombuffer(b, np.float32, NB, pos); pos += 4 * NB
            band = unpack4(b[pos:pos + (N + 1) // 2], N); pos += (N + 1) // 2; uni = ucb[band].tolist(); bn = None
            P = b[pos:pos + N]; S = b[pos + N:pos + N + nfc]; words = []; prev = []; j = 0
            for k in P:
                cur = prev[:k]
                while True:
                    c = S[j]; j += 1; ch = c & 0x7F
                    if ch != 0x7F: cur.append(ch)
                    if c & 0x80: break
                words.append("".join(chr(0x0900 + x) for x in cur)); prev = cur
            pos = end = 0
        else:
            N, nb, nfc, self.unk = struct.unpack("<IIIf", b[:16]); pos = 16
            bn, pos = read_varints(b, nb, pos); blp = np.frombuffer(b, np.float16, nb, pos).astype(np.float32).tolist(); pos += 2 * nb
        if bn is not None: words = []; prev = []; end = pos + nfc
        while pos < end:
            k = b[pos]; pos += 1; cur = prev[:k]
            while True:
                c = b[pos]; pos += 1; ch = c & 0x7F
                if ch != 0x7F: cur.append(ch)
                if c & 0x80: break
            words.append("".join(chr(0x0900 + x) for x in cur)); prev = cur
        self.words = words; self.N = N; self.id = {w: i for i, w in enumerate(words)}
        self.uni = uni if bn is None else [lp for n, lp in zip(bn, blp) for _ in range(n)]
        self.bi = {}; self.bow = {}
        if lm_path:
            b = open(lm_path, "rb").read(); N2, nbg, l1, l2, nh = struct.unpack("<IIIII", b[:20]); pos = 20
            cb = np.frombuffer(b, np.float32, 16, pos); pos += 64; cbb = np.frombuffer(b, np.float32, 16, pos); pos += 64
            cnt, pos = read_varints(b, N + 1, pos); delta, pos = read_varints(b, nbg, pos)
            q = cb[unpack4(b[pos:pos + (nbg + 1) // 2], nbg)].tolist(); pos += (nbg + 1) // 2
            qbw = cbb[unpack4(b[pos:pos + (nh + 1) // 2], nh)].tolist()
            i = 0; hi = 0
            for h, n in enumerate(cnt):
                if not n: continue
                self.bow[h] = qbw[hi]; hi += 1; acc = -1
                for _ in range(n):
                    acc = delta[i] if acc < 0 else acc + delta[i] + 1
                    self.bi[h * HK + acc] = q[i]; i += 1

    def unigram(self, wid):
        return self.unk if wid is None else (math.log(0.04) if wid == self.N else self.uni[wid])

    def logp(self, prev, wid):
        """ln P(wid | prev). prev None = sentence start; wid None = OOV word; wid N = </s>."""
        h = self.N if prev is None else prev
        if wid is not None:
            p = self.bi.get(h * HK + wid)
            if p is not None: return p
        return self.bow.get(h, 0.0) + self.unigram(wid)


def write_assets(out_dir, n_lex, maxbi):
    import os
    root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
    rows = [l.rstrip("\n").split("\t") for l in open(f"{root}/data/processed/trackB/hi_wordfreq.tsv", encoding="utf-8")]
    d = np.load(f"{root}/data/processed/trackB/lm/bigrams.npz")
    lex, lm, st = build([r[0] for r in rows], [int(r[1]) for r in rows], d["pairs"], d["counts"], n_lex, maxbi)
    os.makedirs(out_dir, exist_ok=True)
    open(f"{out_dir}/lex.bin", "wb").write(lex); open(f"{out_dir}/lm.bin", "wb").write(lm)
    st.update(lex_bytes=len(lex), lex_gz=len(gzip.compress(lex, 9)), lm_bytes=len(lm), lm_gz=len(gzip.compress(lm, 9)))
    return st