"""T22. Mixed-key solver over an ELEMENT vocabulary (after Lasry, "Deciphering Historical Syllabic
Ciphers", HistoCrypt 2024): the plaintext is a sequence of elements - letters, syllables, nomenclator
words, a word boundary, and a null - and each cipher unit encodes exactly one element.  The language
model is an element trigram trained on the reference text decomposed by the same scheme the key is
assumed to use (nomenclator words whole; otherwise syllables with probability p_syl, longest first;
nulls inserted at rate p_null, so that the model expects them).  Because every unit emits one
element, the length of the decoded text is fixed by the ciphertext, and the score cannot be raised by
writing longer or shorter plaintext - the failure of the two character-level objectives tried first
(salad of common Latin strings at 0-16 % key accuracy on planted keys; recorded in mixed_solver.py).

Search: simulated annealing over unit -> element (reassign 0.6, swap 0.4), score = sum log P3 over
windows touched by the moved units + lam * T * KL(element frequencies || reference), then coordinate
polishing.  With learned segmentation, alternated with beam re-segmentation of the written words.
"""
import os, sys
_here = os.path.dirname(os.path.abspath(__file__))  # data/analysis, or a subfolder of it
sys.path.insert(0, _here if os.path.exists(os.path.join(_here, "paths.py")) else os.path.dirname(_here))
from paths import *  # noqa: repository locations
import collections, math, random, sys
import numpy as np
from corpus import tokenize_glyphs

NULL, SPACE = "∅", " "


def decompose(words, syllables, nomen, p_syl=0.7, p_null=0.05, spaces=True, seed=0):
    """reference words -> element sequence"""
    rng = random.Random(seed); syl3 = {s for s in syllables if len(s) == 3}; syl2 = {s for s in syllables if len(s) == 2}
    out = []
    for w in words:
        els = []
        if w in nomen: els = [w]
        else:
            i = 0
            while i < len(w):
                if w[i:i + 3] in syl3 and rng.random() < p_syl: els.append(w[i:i + 3]); i += 3
                elif w[i:i + 2] in syl2 and rng.random() < p_syl: els.append(w[i:i + 2]); i += 2
                else: els.append(w[i]); i += 1
        for e in els:
            if rng.random() < p_null: out.append(NULL)
            out.append(e)
        if spaces: out.append(SPACE)
    return out


class ElemLM:
    def __init__(self, ref_words, syllables, nomen, p_syl=0.7, p_null=0.05, spaces=True, k=(0.5, 0.5), seed=0):
        seq = decompose(ref_words, syllables, nomen, p_syl, p_null, spaces, seed)
        cnt = collections.Counter(seq)
        self.alpha = [e for e, c in cnt.most_common() if c >= 3]
        if NULL not in self.alpha: self.alpha.append(NULL)
        self.idx = {e: i for i, e in enumerate(self.alpha)}; self.L = L = len(self.alpha)
        self.sp = self.idx.get(SPACE) if spaces else None; self.null = self.idx[NULL]
        s = np.array([self.idx[e] for e in seq if e in self.idx], dtype=np.int64); self.seq = s
        u = np.bincount(s, minlength=L).astype(float); P1 = (u + 1) / (u.sum() + L)
        c2 = np.zeros((L, L)); np.add.at(c2, (s[:-1], s[1:]), 1)
        P2 = (c2 + k[0] * P1[None, :]) / (c2.sum(1, keepdims=True) + k[0])
        c3 = np.zeros((L, L, L)); np.add.at(c3, (s[:-2], s[1:-1], s[2:]), 1)
        P3 = (c3 + k[1] * P2[None, :, :]) / (c3.sum(2, keepdims=True) + k[1])
        self.logP3 = np.log(P3).reshape(-1); self.uni = u / u.sum()
        self.kind = ["null" if e == NULL else "space" if e == SPACE else "word" if e in nomen else "syl" if len(e) > 1 else "letter" for e in self.alpha]

    def score_seq(self, s):
        if len(s) < 3: return 0.0
        L = self.L; return float(self.logP3[(s[:-2] * L + s[1:-1]) * L + s[2:]].mean())

    def own_score(self): return self.score_seq(self.seq)


class Units:
    """cipher text as unit ids; written-word boundaries are fixed to the SPACE element (if spaces)"""
    def __init__(self, segmented, lm, min_count=2, spaces=True):
        cnt = collections.Counter(s for ss in segmented for s in ss)
        self.alpha = [u for u, c in cnt.most_common() if c >= min_count]; self.idx = {u: i for i, u in enumerate(self.alpha)}
        self.nU = len(self.alpha); seq = []
        for ss in segmented:
            seq += [self.idx[s] for s in ss if s in self.idx]
            if spaces: seq.append(-1)
        self.seq = np.array(seq, dtype=np.int64); self.freq = np.array([cnt[u] for u in self.alpha], float)
        self.pos = [np.flatnonzero(self.seq == i) for i in range(self.nU)]


def solve(U, lm, lam=1.0, restarts=2, iters=150000, seed=0, verbose=False, polish_sweeps=4, allowed=None, init=None):
    """allowed: optional bool array (nU x L) of permitted unit -> element assignments; init: starting key for restart 0"""
    rng = random.Random(seed); L = lm.L
    elems = [i for i in range(L) if i != lm.sp]
    uni = lm.uni[elems] / lm.uni[elems].sum(); eidx = {e: k for k, e in enumerate(elems)}
    freq = U.freq / U.freq.sum(); seq = U.seq; T = len(seq); nU = U.nU
    starts = []
    for i in range(nU):
        p = U.pos[i]; w = np.concatenate([p - 2, p - 1, p]); starts.append(np.unique(w[(w >= 0) & (w <= T - 3)]))
    logP3 = lm.logP3; spv = lm.sp if lm.sp is not None else 0
    def wsc(P, st): return logP3[(P[st] * L + P[st + 1]) * L + P[st + 2]].sum()
    ne = len(elems); Z = 1.0 + ne * 1e-9; lu = np.log(uni + 1e-9)
    ek = np.full(L, -1); ek[elems] = np.arange(ne)   # element id -> index in elems
    def term(k, F):  # contribution of element index k (array or int) with mass F to the KL
        f = (F + 1e-9) / Z; return f * (np.log(f) - lu[k])
    def kl_full(F): return float(term(np.arange(ne), F).sum())
    def mass(m):
        F = np.zeros(ne); np.add.at(F, ek[m], freq); return F
    def dkl_move(F, a, b, x):  # move mass x from element index a to b
        if a == b: return 0.0
        return float(term(a, F[a] - x) - term(a, F[a]) + term(b, F[b] + x) - term(b, F[b]))
    ok = (lambda i, e: True) if allowed is None else (lambda i, e: bool(allowed[i, e]))
    best = (-1e18, None, None, None); T0 = 0.02 * T
    for r in range(restarts):
        order = np.argsort(-U.freq); lord = [elems[k] for k in np.argsort(-uni)]
        m = np.zeros(nU, int); slot = 0
        for ui in order:
            for _ in range(len(lord)):
                e = lord[slot % len(lord)]; slot += 1
                if ok(ui, e): m[ui] = e; break
        if init is not None:
            m0 = np.array(init, int); m = np.where(m0 >= 0, m0, m)
        if r > 0:
            for i in range(nU):
                if rng.random() < (0.5 if init is None else 0.1):
                    e = rng.choice(elems)
                    if ok(i, e): m[i] = e
        P = np.where(seq >= 0, m[np.maximum(seq, 0)], spv).astype(np.int64)
        total = float(logP3[(P[:-2] * L + P[1:-1]) * L + P[2:]].sum()); F = mass(m); kl = kl_full(F); score = total - lam * T * kl
        for it in range(iters):
            temp = T0 * (1 - it / iters) + 1e-9; i = rng.randrange(nU)
            if rng.random() < 0.6:
                new = rng.choice(elems)
                if new == m[i] or not ok(i, new): continue
                st = starts[i]; o = wsc(P, st); P[U.pos[i]] = new; n = wsc(P, st)
                a, b = ek[m[i]], ek[new]; d = dkl_move(F, a, b, freq[i]); sc2 = score + (n - o) - lam * T * d
                if sc2 >= score or rng.random() < math.exp((sc2 - score) / temp):
                    F[a] -= freq[i]; F[b] += freq[i]; m[i] = new; kl += d; score = sc2
                else: P[U.pos[i]] = m[i]
            else:
                j = rng.randrange(nU)
                if m[i] == m[j] or not ok(i, m[j]) or not ok(j, m[i]): continue
                st = np.union1d(starts[i], starts[j]); o = wsc(P, st)
                P[U.pos[i]] = m[j]; P[U.pos[j]] = m[i]; n = wsc(P, st)
                a, b = ek[m[i]], ek[m[j]]; F2 = F.copy(); F2[a] += freq[j] - freq[i]; F2[b] += freq[i] - freq[j]
                d = float(term(a, F2[a]) + term(b, F2[b]) - term(a, F[a]) - term(b, F[b])); sc2 = score + (n - o) - lam * T * d
                if sc2 >= score or rng.random() < math.exp((sc2 - score) / temp):
                    F = F2; m[i], m[j] = m[j], m[i]; kl += d; score = sc2
                else: P[U.pos[i]] = m[i]; P[U.pos[j]] = m[j]
        for sw in range(polish_sweeps):
            ch = 0
            for i in np.argsort(-U.freq):
                st = starts[i]; cur = m[i]; base = wsc(P, st); bl, bg, bd = cur, 0.0, 0.0
                a = ek[cur]; x = freq[i]; ks = np.arange(ne)
                dk = term(a, F[a] - x) - term(a, F[a]) + term(ks, F + x) - term(ks, F)   # KL change for moving to each element
                for e in elems:
                    if e == cur or not ok(i, e): continue
                    P[U.pos[i]] = e; g = wsc(P, st) - base - lam * T * float(dk[ek[e]])
                    if g > bg: bl, bg, bd = e, g, float(dk[ek[e]])
                P[U.pos[i]] = bl
                if bl != cur: F[a] -= x; F[ek[bl]] += x; m[i] = bl; kl += bd; score += bg; ch += 1
            if not ch: break
        kl = kl_full(mass(m))  # re-sync against drift
        if verbose: print(f"  restart {r}: score/elem {score / T:.4f} kl {kl:.3f}", flush=True)
        if score > best[0]: best = (score, m.copy(), score / T, kl)
    return best


def carry_key(prev, U, lm):
    """initial key for a new unit inventory from the previous round's {segment: element id};
    units not seen before get -1 (left to the frequency initialisation)"""
    return np.array([prev.get(u, -1) for u in U.alpha], int)


def key_dict(U, m): return {u: int(m[i]) for i, u in enumerate(U.alpha)}


def decode_words(segmented, U, lm, m):
    out = []
    def el(s):
        if s not in U.idx: return "?"
        e = lm.alpha[m[U.idx[s]]]
        return "" if e == NULL else e
    for ss in segmented: out.append("".join(el(s) for s in ss))
    return out


def heldout_score(segmented, U, lm, m, spaces=True):
    s = []
    for ss in segmented:
        s += [m[U.idx[x]] for x in ss if x in U.idx]
        if spaces and lm.sp is not None: s.append(lm.sp)
    return lm.score_seq(np.array(s, dtype=np.int64))


def resegment(words, lex, U, lm, m, beam=5, spaces=True):
    """re-cut each written word into lexicon segments maximising the element-trigram score under the key"""
    L = lm.L; ctx = (lm.sp or 0, lm.sp or 0); out = []
    for w in words:
        u = tokenize_glyphs(w); Bm = [(0.0, 0, ctx, [])]; done = None
        while Bm:
            nb = []
            for sc, pos, c2, segs in Bm:
                if pos == len(u):
                    if done is None or sc > done[0]: done = (sc, c2, segs)
                    continue
                for Ln, s in lex.candidates(u, pos):
                    if s in U.idx:
                        e = m[U.idx[s]]; nsc = sc + lm.logP3[(c2[0] * L + c2[1]) * L + e]; nc = (c2[1], e)
                    else: nsc, nc = sc - 12.0, c2
                    nb.append((nsc, pos + Ln, nc, segs + [s]))
            nb.sort(key=lambda x: -x[0]); Bm = nb[:beam]
        if done is None: out.append(list(u)); continue
        out.append(done[2]); ctx = done[1]
        if spaces and lm.sp is not None: ctx = (ctx[1], lm.sp)
    return out