"""T20/T21, final form.

T20 sign-level collapse at matched vocabulary.  Units: glyph units for the manuscript and the mixed-key
simulators, letters for Latin; alternations = single-unit substitutions in otherwise identical words
(position-free classes, the same procedure for every text); merged in order of the context similarity
of their minimal pairs, or at random.  Because merging any alternation shrinks the vocabulary and so
raises repetition, the comparison is made at equal numbers of word types: for each text, the context
curve is compared with the random curve interpolated to the same type count ("excess").  The
simulators give the expectation under the hypothesis, with the oracle (true plaintext words) as the
ceiling; Latin gives the expectation for a language.

T21b position-blindness of variant choice.  For the commonest alternation classes, the mutual
information between which variant is written and (a) the word's position in its line (first, last,
inside) and (b) the last glyph of the preceding word, against 200 within-page shuffles of the
variants.  Homophones chosen by hand, at random or by cycling are blind to position; allographs of a
script, or a choice tied to the surrounding text, are not.
Writes out/collapse3.csv, out/position_blind.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 collections, csv, math, random, sys
import numpy as np
from corpus import tokenize_glyphs
from harness import structured, split, flatten, languages, log
import battery as B, generators as G, mixedkey as MK
from t_collapse2 import alternations, canonicalizer, remap

STATS = ["order_gain", "heldout_bigram_gain", "rep4"]
parse_g = lambda w: tuple(tokenize_glyphs(w))
parse_l = lambda w: tuple(w)


def curve(name, pages, parse, ks=(5, 10, 20, 40, 80), n_rand=4, oracle=None):
    words = B.flat(pages); alts = alternations(words, parse, positional=False)
    rows = [dict(text=name, method="none", k=0, types=len(set(words)), **B.all_stats(pages, STATS))]
    for k in ks:
        if k > len(alts): break
        f = canonicalizer([a for a, s, n in alts[:k]], words, parse, positional=False); cp = remap(pages, f)
        rows.append(dict(text=name, method="context", k=k, types=len(set(B.flat(cp))), **B.all_stats(cp, STATS)))
        for d in range(n_rand):
            pick = random.Random(100 + d).sample(alts, k); f2 = canonicalizer([a for a, s, n in pick], words, parse, positional=False); c2 = remap(pages, f2)
            rows.append(dict(text=name, method="random", k=k, types=len(set(B.flat(c2))), **B.all_stats(c2, STATS)))
    if oracle is not None: rows.append(dict(text=name, method="oracle", k=None, types=len(set(B.flat(oracle))), **B.all_stats(oracle, STATS)))
    ctx = [r for r in rows if r["method"] == "context"]; rnd = sorted([r for r in rows if r["method"] in ("random", "none")], key=lambda r: r["types"])
    for r in ctx:
        for s in STATS:
            xs = [x["types"] for x in rnd]; ys = [float(x[s]) for x in rnd]
            r[f"excess_{s}"] = round(float(r[s]) - float(np.interp(r["types"], xs, ys)), 4) if min(xs) <= r["types"] <= max(xs) else None
    print(name, "alternations", len(alts), "top", [(a[1:], round(s, 2)) for a, s, n in alts[:5]], flush=True)
    for r in ctx: print(f"  k={r['k']} types {r['types']} " + " ".join(f"{s}={float(r[s]):.3f} (excess {r['excess_' + s]})" for s in STATS), flush=True)
    if oracle is not None: print("  oracle", {s: round(float(rows[-1][s]), 3) for s in STATS}, flush=True)
    return rows, alts


def position_mi(pages, alts, parse, top=12, n_shuf=200, seed=0):
    """MI(variant; line position) and MI(variant; previous word's last glyph) for the top alternation classes"""
    rng = random.Random(seed); out = []
    for (pos, a, b), s, n in alts[:top]:
        occ = []  # (page, variant, line_pos, prev_last)
        for p in pages:
            for ln in p["lines"]:
                for i, w in enumerate(ln):
                    t = parse(w)
                    if not t: continue
                    has_a, has_b = a in t, b in t
                    if has_a == has_b: continue
                    lp = "first" if i == 0 else ("last" if i == len(ln) - 1 else "inside")
                    prev = parse(ln[i - 1])[-1] if i > 0 and parse(ln[i - 1]) else "^"
                    occ.append((p["page"], a if has_a else b, lp, prev))
        if len(occ) < 50: continue
        def mi(pairs):
            c = collections.Counter(pairs); cx = collections.Counter(x for x, y in pairs); cy = collections.Counter(y for x, y in pairs); N = len(pairs)
            return sum(v / N * math.log2((v / N) / ((cx[x] / N) * (cy[y] / N))) for (x, y), v in c.items())
        obs_l = mi([(v, lp) for _, v, lp, pr in occ]); obs_p = mi([(v, pr) for _, v, lp, pr in occ])
        by_page = collections.defaultdict(list)
        for k, o in enumerate(occ): by_page[o[0]].append(k)
        nl, npv = [], []
        for _ in range(n_shuf):
            var = [o[1] for o in occ]
            for ks in by_page.values():
                vs = [var[k] for k in ks]; rng.shuffle(vs)
                for k, v in zip(ks, vs): var[k] = v
            nl.append(mi([(v, o[2]) for v, o in zip(var, occ)])); npv.append(mi([(v, o[3]) for v, o in zip(var, occ)]))
        z = lambda x, xs: (x - np.mean(xs)) / (np.std(xs) + 1e-12)
        out.append(dict(alternation=f"{a}/{b}", n=len(occ), context_score=round(s, 3), mi_line=round(obs_l, 4), mi_line_null=round(float(np.mean(nl)), 4), z_line=round(float(z(obs_l, nl)), 1),
                        mi_prev=round(obs_p, 4), mi_prev_null=round(float(np.mean(npv)), 4), z_prev=round(float(z(obs_p, npv)), 1)))
    return out


if __name__ == "__main__":
    pages = structured(); tr, te = split(pages); lay = G.layout_of(pages); trw = flatten(tr)
    rows = []; pos_rows = []
    r, alts_v = curve("Voynich", pages, parse_g); rows += r
    for x in position_mi(pages, alts_v, parse_g): pos_rows.append(dict(text="Voynich", **x))
    lat = G.fill(lay, languages()["Latin"][:34217])
    r, alts_l = curve("Latin", lat, parse_l); rows += r
    for x in position_mi(lat, alts_l, parse_l): pos_rows.append(dict(text="Latin", **x))
    pw = lambda t: "".join(u.strip("<>") for s, u in t)
    for sp in ("word", "split"):
        sim, truths, K = MK.gen_tranchedino_sectioned(trw, lay, seed=7, spacing=sp, n_syl=60, n_tri=0, n_nom=80, k_v=3, k_c=2, p_syl=0.7, return_truth=True)
        it = iter(truths); oracle = [dict(p, lines=[[pw(next(it)) for w in ln] for ln in p["lines"]]) for p in sim]
        r, alts_s = curve(f"mixed key, {sp} spacing", sim, parse_g, oracle=oracle); rows += r
        for x in position_mi(sim, alts_s, parse_g): pos_rows.append(dict(text=f"mixed key, {sp} spacing", **x))
    keys = sorted({k for r in rows for k in r}, key=lambda k: (k not in ("text", "method", "k", "types"), k))
    with open("out/collapse3.csv", "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=keys); w.writeheader(); w.writerows(rows)
    with open("out/position_blind.csv", "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(pos_rows[0].keys())); w.writeheader(); w.writerows(pos_rows)
    for r in pos_rows: print(r)
    def ex(text, s="rep4"):
        v = [r.get(f"excess_{s}") for r in rows if r["text"] == text and r["method"] == "context" and r.get(f"excess_{s}") is not None]
        return round(float(np.mean(v)), 3) if v else None
    log("H-CIPHER-MIXED", "T20-collapse (matched vocabulary)", "mean excess of context-chosen over random merges at equal type count: repeated 4-grams / held-out bigram gain",
        f"Voynich {ex('Voynich')} / {ex('Voynich', 'heldout_bigram_gain')}", f"Latin {ex('Latin')} / {ex('Latin', 'heldout_bigram_gain')}; mixed key word {ex('mixed key, word spacing')} / {ex('mixed key, word spacing', 'heldout_bigram_gain')}; split {ex('mixed key, split spacing')} / {ex('mixed key, split spacing', 'heldout_bigram_gain')}", "see note", "out/collapse3.csv")
    vz = [r for r in pos_rows if r["text"] == "Voynich"]
    log("H-CIPHER-MIXED", "T21b-position-blind", "variant choice vs line position and previous glyph, top alternation classes: median z against within-page shuffles",
        f"line {np.median([r['z_line'] for r in vz]):.1f}; previous glyph {np.median([r['z_prev'] for r in vz]):.1f}", "simulators and Latin in out/position_blind.csv", "see note", "out/position_blind.csv")