"""The layout layer (Q-023, N-001): which spellings change at the start and end of a line, learned on
half E of the bifolium split, and a normaliser that writes each line-edge word in its inside-line form.

Rules.  A rule is a minimal alternation at a line position P (first word of a line that does not
open a paragraph, first word of a paragraph, last word of a line): at the start, one glyph unit a
written before a stem in place of nothing or of one other unit b, or a unit b left out; at the end,
one glyph unit a written after a stem in place of an ending b of one to three units (so -m for -r,
-l, -in, -iin).
For every stem s seen with both s+a and s+b, the a-form's share at P is compared with its share
inside lines.  A rule is kept when, on half E, the share at P is at least three times the inside
share, at least 15 stems carry it, and a within-page shuffle of positions gives z >= 5 (the bar is
high because some two thousand alternations are screened).

Normalising.  A word at P that matches a kept rule (s+a) is rewritten s+b when s+b occurs inside a
line anywhere in the book and is commoner there than the word itself; among several b the commonest
inside form wins.  Nothing else is touched.
A second, model-free treatment drops the line-edge words altogether (interior()), so a result that
holds under both cannot be a layout effect.  The same code runs on the null texts (the Latin control should
lose almost nothing) and on a planted control (Latin with three invented line-edge rules, which the
procedure must find and undo).

Usage: python3 layout_layer.py explore     -> out/layout_rules_E.csv, out/layout_candidates_E.csv
       python3 layout_layer.py control     -> out/layout_control.csv (Latin, planted Latin)
Library: rules(), normalise_pages(pages, rules)
"""
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, json, random
import numpy as np
import pandas as pd
HERE = os.path.dirname(os.path.abspath(__file__)); sys.path.insert(0, HERE); sys.path.insert(0, os.path.join(HERE, ".."))
from corpus import tokenize_glyphs

RULES_CSV = os.path.join(HERE, "out", "layout_rules_E.csv")
POS = ("first", "pfirst", "last")
MIN_RATIO, MIN_STEMS, MIN_Z, MAX_AFFIX = 3.0, 15, 5.0, 3


def positioned(pages):
    """(page, position, word) for every word; positions first / pfirst / last / inside.
    pages carry lines and, per line, a par_start flag (texts.voynich adds it; default False)."""
    out = []
    for p in pages:
        flags = p.get("par_start") or [False] * len(p["lines"])
        for l, ps in zip(p["lines"], flags):
            n = len(l)
            for i, w in enumerate(l):
                if n == 1: pos = "single"
                elif i == 0: pos = "pfirst" if ps else "first"
                elif i == n - 1: pos = "last"
                else: pos = "inside"
                out.append((p["page"], pos, w))
    return out


def _units(w, glyphs): return tuple(tokenize_glyphs(w)) if glyphs else tuple(w)


def candidates(pages, glyphs=True, n_perm=300, seed=0):
    occ = positioned(pages); U = {w: _units(w, glyphs) for _, _, w in occ}
    inside = collections.Counter(w for _, pos, w in occ if pos == "inside")
    vocab = set(U); key = {}
    # index by (stem, side) -> words, for affixes of 0..MAX_AFFIX units
    by_stem = collections.defaultdict(set)
    for w in vocab:
        t = U[w]
        for k in range(0, min(MAX_AFFIX, len(t) - 1) + 1):
            by_stem[("end", t[:len(t) - k])].add((t[len(t) - k:], w))
            by_stem[("start", t[k:])].add((t[:k], w))
    pairs = collections.defaultdict(set)  # (side, a, b) -> stems
    for (side, stem), lst in by_stem.items():
        if len(lst) < 2 or not stem: continue
        for a, wa in lst:
            for b, wb in lst:
                if a == b or wa == wb: continue
                if not (len(a) == 1 or (not a and len(b) == 1)): continue  # one unit written, or one unit left out, at the edge
                if side == "end" and not b: continue
                if not a and side == "end": continue
                if side == "start" and len(b) > 1: continue
                pairs[(side, a, b)].add(stem)
    rows = []; rng = random.Random(seed)
    for (side, a, b), stems in pairs.items():
        if len(stems) < MIN_STEMS: continue
        ok = set(); wa_of = {}; wb_of = {}
        for s in stems:
            wa = s + a if side == "end" else a + s; wb = s + b if side == "end" else b + s
            wa_of[wa] = s; wb_of[wb] = s
        for P in (("last",) if side == "end" else ("first", "pfirst")):
            sel = [(pg, pos, U[w] in wa_of) for pg, pos, w in occ if pos in (P, "inside") and (U[w] in wa_of or U[w] in wb_of)]
            if len(sel) < 30: continue
            atP = [x for _, pos, x in sel if pos == P]; ins = [x for _, pos, x in sel if pos == "inside"]
            if len(atP) < 10 or not ins: continue
            sP, sI = float(np.mean(atP)), float(np.mean(ins)); obs = sP - sI
            by = collections.defaultdict(list)
            for k, (pg, _, _) in enumerate(sel): by[pg].append(k)
            null = []
            for _ in range(n_perm):
                lab = [pos for _, pos, _ in sel]
                for ks in by.values():
                    v = [lab[k] for k in ks]; rng.shuffle(v)
                    for k, x in zip(ks, v): lab[k] = x
                e = [sel[k][2] for k in range(len(sel)) if lab[k] == P]; i_ = [sel[k][2] for k in range(len(sel)) if lab[k] == "inside"]
                null.append((np.mean(e) if e else 0) - (np.mean(i_) if i_ else 0))
            z = (obs - np.mean(null)) / (np.std(null) + 1e-12)
            rows.append(dict(position=P, side=side, edge_form="".join(a) or "(none)", inside_form="".join(b) or "(none)", stems=len(stems),
                             tokens_at_P=len(atP), share_at_P=round(sP, 4), share_inside=round(sI, 4),
                             ratio=round(sP / sI, 2) if sI else float("inf"), z=round(float(z), 2)))
    return pd.DataFrame(rows).sort_values("z", ascending=False) if rows else pd.DataFrame()


def keep(cands):
    if cands.empty: return cands
    k = cands[(cands.ratio >= MIN_RATIO) & (cands.stems >= MIN_STEMS) & (cands.z >= MIN_Z)].copy()
    return k.sort_values("z", ascending=False)


def rules(path=RULES_CSV):
    return pd.read_csv(path, keep_default_na=False)


def normaliser(pages, R, glyphs=True):
    """returns f(word, position) -> word, using the inside-line vocabulary of `pages`"""
    occ = positioned(pages); inside = collections.Counter(w for _, pos, w in occ if pos == "inside")
    iu = {}
    for w, c in inside.items(): iu.setdefault(_units(w, glyphs), (w, c))
    R = R.sort_values("z", ascending=False).to_dict("records") if len(R) else []
    def f(w, pos):
        if pos not in POS: return w
        t = _units(w, glyphs); best = None
        for r in R:
            if r["position"] != pos: continue
            a = () if r["edge_form"] == "(none)" else tuple(tokenize_glyphs(r["edge_form"])) if glyphs else tuple(r["edge_form"])
            b = () if r["inside_form"] == "(none)" else (tuple(tokenize_glyphs(r["inside_form"])) if glyphs else tuple(r["inside_form"]))
            if r["side"] == "end" and len(t) > len(a) and t[-len(a):] == a: cand = t[:-len(a)] + b
            elif r["side"] == "start" and len(t) > len(a) and t[:len(a)] == a: cand = b + t[len(a):]
            else: continue
            if cand in iu and (best is None or iu[cand][1] > best[1]): best = iu[cand]
        own = iu.get(t, (w, 0))[1]
        return best[0] if best and best[1] > own else w  # rewrite only to an alternant commoner inside lines than the word itself
    return f


def normalise_pages(pages, R, glyphs=True, ref=None):
    f = normaliser(ref or pages, R, glyphs); out = []; changed = 0; n = 0
    for p in pages:
        flags = p.get("par_start") or [False] * len(p["lines"]); lines = []
        for l, ps in zip(p["lines"], flags):
            m = len(l); nl = []
            for i, w in enumerate(l):
                pos = "single" if m == 1 else ("pfirst" if ps else "first") if i == 0 else "last" if i == m - 1 else "inside"
                v = f(w, pos); changed += v != w; n += 1; nl.append(v)
            lines.append(nl)
        out.append(dict(p, lines=lines))
    return out, changed / max(n, 1)


def interior(pages):
    """every line without its first and last word (lines of one or two words vanish)"""
    return [dict(p, lines=[l[1:-1] for l in p["lines"] if len(l) > 2], par_start=[f for l, f in zip(p["lines"], p.get("par_start") or [False] * len(p["lines"])) if len(l) > 2]) for p in pages]


def layout_mi(pages, glyphs=True):
    """mutual information (bits) between line position (first/inside/last) and the word's first and last
    glyph unit, as in the battery's layout method, over lines of 3+ words"""
    def mi(pairs):
        n = len(pairs); cxy = collections.Counter(pairs); cx = collections.Counter(x for x, _ in pairs); cy = collections.Counter(y for _, y in pairs)
        return sum(c / n * np.log2(c * n / (cx[x] * cy[y])) for (x, y), c in cxy.items())
    f, l = [], []
    for p in pages:
        for ln in p["lines"]:
            if len(ln) < 3: continue
            for i, w in enumerate(ln):
                pos = 0 if i == 0 else 2 if i == len(ln) - 1 else 1; t = _units(w, glyphs)
                f.append((pos, t[0])); l.append((pos, t[-1]))
    return round(float(mi(f)), 4), round(float(mi(l)), 4)


def half(pages, h):
    H = json.load(open(os.path.join(HERE, "out", "split.json")))["half"]
    return [p for p in pages if H[p["page"]] == h]


def plant(pages, seed=0):
    """Latin with three invented line-edge rules, one of each kind the procedure looks for: line-final -um
    written -ũ (80 % of eligible words); a mark z- written before 30 % of line-initial words; and a
    line-initial c- left out of 60 % of the words that have it.  (A first plant, h- left out at the line
    start, was too rare to test anything: 26 tokens in half E, most of them with an h-less twin inside.)"""
    rng = random.Random(seed); out = []
    for p in pages:
        lines = []
        for ln in p["lines"]:
            ln = list(ln)
            if len(ln) > 1 and ln[-1].endswith("um") and len(ln[-1]) > 3 and rng.random() < 0.8: ln[-1] = ln[-1][:-2] + "ũ"
            if len(ln) > 1:
                if ln[0].startswith("c") and len(ln[0]) > 3 and rng.random() < 0.6: ln[0] = ln[0][1:]
                elif rng.random() < 0.3: ln[0] = "z" + ln[0]
            lines.append(ln)
        out.append(dict(p, lines=lines))
    return out


if __name__ == "__main__":
    from texts import voynich, latin
    cmd = sys.argv[1] if len(sys.argv) > 1 else "explore"
    V = voynich()
    if cmd == "explore":
        c = candidates(half(V, "E")); c.to_csv(os.path.join(HERE, "out", "layout_candidates_E.csv"), index=False)
        k = keep(c); k.to_csv(RULES_CSV, index=False)
        print(len(c), "candidates;", len(k), "kept"); print(k.to_string(index=False))
    elif cmd == "control":
        rows = []
        LA = latin(V); PL = plant(LA)
        for name, T, gl in [("latin", LA, False), ("latin+planted", PL, False)]:
            TE = half(T, "E"); TC = half(T, "C")
            c = candidates(TE, glyphs=gl); k = keep(c)
            k.to_csv(os.path.join(HERE, "out", f"layout_rules_{name}_E.csv"), index=False)
            N, share = normalise_pages(TC, k, glyphs=gl, ref=T)
            O = half(LA, "C"); pl = ok = fp = un = 0
            for po, pn, pp in zip(O, N, TC):
                for lo, ln, lp in zip(po["lines"], pn["lines"], pp["lines"]):
                    for wo, wn, wp in zip(lo, ln, lp):
                        if wo != wp: pl += 1; ok += wn == wo
                        else: un += 1; fp += wn != wo
            rows.append(dict(text=name, planted_tokens_C=pl, restored_exactly=round(ok / pl, 3) if pl else None, false_rewrites=round(fp / max(un, 1), 4), kept_rules=len(k), rules="; ".join(f"{r.position}:{r.edge_form}->{r.inside_form}" for r in k.itertuples()) if len(k) else "",
                             changed_share_C=round(share, 4), mi_first_C=layout_mi(TC, gl)[0], mi_last_C=layout_mi(TC, gl)[1],
                             mi_first_C_norm=layout_mi(N, gl)[0], mi_last_C_norm=layout_mi(N, gl)[1]))
            print(rows[-1], flush=True)
        pd.DataFrame(rows).to_csv(os.path.join(HERE, "out", "layout_control.csv"), index=False)