"""T2. One-to-one substitution screen against 103 languages × {plain, vowel-less}.
For each: best key by bigram log-likelihood (hill-climbing, validated on a planted-key
control), gap to the language's own text, dictionary hit rate under the best key vs
random keys, and a held-out check (key fitted on training quires, scored on held-out).
Decision rule (PROTOCOL §4): a language is a candidate only if the held-out gap is within
the control's recovery tolerance AND dictionary hits exceed random by more than the
control's margin. Writes out/subst_screen.csv.
"""
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 sys, csv, random, collections, math, time
import numpy as np
from harness import *

random.seed(11); np.random.seed(11)
pages = structured(); tr, te = split(pages)
V_tr, V_te, V_all = flatten(tr), flatten(te), flatten(pages)


def stream(words, units=True):
    return char_stream(words, units)


def bigram_logp(seq, alphabet, k=0.5):
    idx = {a: i for i, a in enumerate(alphabet)}
    M = np.full((len(alphabet), len(alphabet)), k)
    for a, b in zip(seq, seq[1:]):
        if a in idx and b in idx: M[idx[a], idx[b]] += 1
    return np.log(M / M.sum(1, keepdims=True)), idx


def counts(seq, alphabet):
    idx = {a: i for i, a in enumerate(alphabet)}
    M = np.zeros((len(alphabet), len(alphabet)))
    for a, b in zip(seq, seq[1:]):
        if a in idx and b in idx: M[idx[a], idx[b]] += 1
    return M


def solve(Cm, logP, n_src, n_tgt, restarts=8, iters=5000):
    best, bm = -1e18, None
    for r in range(restarts):
        m = np.append(np.random.permutation(n_tgt)[:n_src], n_tgt)
        sc = lambda m: (Cm * logP[np.ix_(m, m)]).sum()
        s = sc(m)
        for it in range(iters):
            m2 = m.copy(); i = random.randrange(n_src)
            if random.random() < 0.6:
                j = random.randrange(n_src); m2[i], m2[j] = m2[j], m2[i]
            else:
                t = random.randrange(n_tgt)
                if t in m2[:n_src]: continue
                m2[i] = t
            s2 = sc(m2)
            if s2 >= s: m, s = m2, s2
        if s > best: best, bm = s, m
    return best, bm


vs_tr, vs_te, vs_all = stream(V_tr), stream(V_te), stream(V_all)
valpha_all = [a for a, c in collections.Counter(vs_all).most_common() if c >= 100 and a != " "]
langs = languages()
rows = []
t0 = time.time()
for li, (name, words) in enumerate(sorted(langs.items())):
    words = words[:len(V_all)]
    rs = stream(words, units=False)
    ref_alpha = [a for a, c in collections.Counter(rs).most_common() if c >= 20 and a != " "]
    if len(ref_alpha) > 60: ref_alpha = ref_alpha[:60]   # syllabaries/logographic: cap
    ref_alpha = ref_alpha + [" "]; n_tgt = len(ref_alpha) - 1
    voy_alpha = valpha_all[:n_tgt] + [" "]; n_src = len(voy_alpha) - 1
    logP, idx = bigram_logp(rs, ref_alpha)
    Cr = counts(rs, ref_alpha); own = (Cr * logP).sum() / Cr.sum()
    # control: planted key on the language's own text (train half), solve, score
    perm = np.append(np.random.permutation(n_tgt), n_tgt)
    enc = [ref_alpha[perm[idx[a]]] if a in idx else a for a in rs]
    Ce = counts(enc, ref_alpha); sc_, _ = solve(Ce, logP, n_tgt, n_tgt, restarts=4, iters=4000); ctrl = sc_ / Ce.sum()
    # Voynich: fit on training quires, score on held-out
    Ctr = counts(vs_tr, voy_alpha); sfit, m = solve(Ctr, logP, n_src, n_tgt)
    fit_tr = sfit / Ctr.sum()
    Cte = counts(vs_te, voy_alpha); fit_te = (Cte * logP[np.ix_(m, m)]).sum() / Cte.sum()
    mr = np.append(np.random.permutation(n_tgt)[:n_src], n_tgt); rnd = (Cte * logP[np.ix_(mr, mr)]).sum() / Cte.sum()
    mapping = {voy_alpha[i]: ref_alpha[m[i]] for i in range(n_src)}
    vocab = set(words)
    def hits(mp):
        h = t = 0
        for w in V_te:
            g = tokenize_glyphs(w)
            if any(x not in mp for x in g): continue
            t += 1; h += "".join(mp[x] for x in g) in vocab
        return h / max(t, 1)
    hb = hits(mapping)
    keys = list(mapping.keys()); vals = list(mapping.values()); hr = []
    for _ in range(3):
        random.shuffle(vals); hr.append(hits(dict(zip(keys, vals))))
    hr = sum(hr) / len(hr)
    rows.append(dict(language=name, ref_alphabet=n_tgt, own=own, control=ctrl, fit_train=fit_tr, fit_heldout=fit_te, random=rnd,
                     gap_heldout=own - fit_te, hits_best=hb, hits_random=hr, hits_margin=hb - hr,
                     sample=" ".join("".join(mapping.get(x, "?") for x in tokenize_glyphs(w)) for w in V_te[:12])))
    print(f"{li:3d} {name:28s} gap {own - fit_te:6.3f}  hits {hb:.3f} vs {hr:.3f}   [{time.time() - t0:.0f}s]", flush=True)
    with open("out/subst_screen.csv", "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=rows[0].keys()); w.writeheader(); w.writerows(rows)
print("done")