"""T15. Simulators of mechanical text-production procedures, each writing into the manuscript's real
layout (pages, lines, words per line), so that the output can be scored with the same battery.

Every generator takes the training pages (to fit its tables) and a layout (page/line word counts)
and returns pages in harness format.  Plaintext for cipher generators: the Latin Bible (full text
minus the screening sample), or its vowel-less form.

G1 slot    : words drawn independently from a slot grammar fitted on the training text (prefix,
             gallows, middle, x, suffix as independent slots per word template).  No memory.
G2 autocopy: Timm-Schinner-type self-citation.  Each new word copies a word from the previous
             `window` words (within the page) and modifies it with probability p_mod (substitute
             a glyph by a confusable one, insert, delete, or swap a suffix); with probability
             p_new a fresh slot word is drawn instead.
G3 grille  : Rugg-type table-and-grille.  A table of R rows x 3 columns (prefix, middle, suffix
             syllables sampled from the training inventory); a grille with three holes at fixed
             column offsets is slid down the table one row per word; the table is re-filled every
             `refill` words.
G4 naibbe  : Greshko-type affix-table cipher of Latin.  Plaintext letters (and, with probability
             p_bigram, letter bigrams) are each enciphered as prefix + suffix drawn from per-unit
             lists of Voynich-style affix strings (k homophones each); plaintext spaces dropped.
G5 homonull: Tranchedino-type homophonic substitution of Latin with nulls: each letter -> one of
             up to 4 glyph units, nulls inserted at rate p_null, plaintext spaces dropped and the
             stream re-spaced into words with the manuscript's word-length distribution.
G6 wordletter: Trithemius-type word-for-letter table: each plaintext letter -> one of k whole
             Voynich-style words from a fixed table (k homophones per letter), spaces dropped.
"""
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, random, re, sys
from corpus import tokenize_glyphs
from harness import full_language
from units import G1 as SLOT1, G2 as SLOT2

CONFUSE = {"ch": ["sh"], "sh": ["ch"], "k": ["t"], "t": ["k"], "a": ["o"], "o": ["a"], "e": ["ee"], "ee": ["e"], "r": ["l", "s"], "l": ["r"], "s": ["r", "d"], "d": ["s"],
           "in": ["iin"], "iin": ["in"], "y": ["dy"], "dy": ["y"], "p": ["f"], "f": ["p"], "cth": ["ckh"], "ckh": ["cth"]}


def layout_of(pages):
    return [dict(page=p["page"], quire=p.get("quire"), hand=p.get("hand"), lang=p.get("lang"), I=p.get("I"), lines=[len(ln) for ln in p["lines"]]) for p in pages]


def fill(layout, words):
    """pour a word sequence into the layout; returns pages in harness format"""
    it = iter(words); out = []
    for p in layout:
        lines = []
        for n in p["lines"]:
            ln = []
            for _ in range(n):
                try: ln.append(next(it))
                except StopIteration: break
            if ln: lines.append(ln)
        if lines: out.append(dict(page=p["page"], quire=p["quire"], hand=p["hand"], lang=p["lang"], I=p["I"], lines=lines))
    return out


# ---------------- slot grammar fitted on training text
class SlotModel:
    def __init__(self, words, seed=0):
        self.rng = random.Random(seed); self.words = words
        self.templates = collections.Counter()  # tuple of slot indices used
        self.fillers = collections.defaultdict(collections.Counter)  # slot index -> filler counts
        self.n_two = 0; self.n = 0
        for w in words:
            m = SLOT1.match(w); two = False
            if not m:
                m = SLOT2.match(w); two = bool(m)
            if not m: continue
            self.n += 1; self.n_two += two
            groups = m.groups()
            tpl = tuple(i for i, g in enumerate(groups) if g)
            self.templates[tpl] += 1
            for i, g in enumerate(groups):
                if g: self.fillers[i][g] += 1
        self.tpl_list = list(self.templates); self.tpl_w = [self.templates[t] for t in self.tpl_list]
        self.fill_list = {i: list(c) for i, c in self.fillers.items()}; self.fill_w = {i: [c[x] for x in self.fill_list[i]] for i, c in self.fillers.items()}

    def word(self):
        tpl = self.rng.choices(self.tpl_list, self.tpl_w)[0]
        return "".join(self.rng.choices(self.fill_list[i], self.fill_w[i])[0] for i in tpl)

    def affixes(self, words=None, max_units=3):
        """prefix-side and suffix-side string inventories for the table ciphers: every word-initial
        and word-final string of 1..max_units glyph units in the training words, with counts"""
        pre = collections.Counter(); suf = collections.Counter()
        for w in (words or self.words):
            u = tokenize_glyphs(w)
            for L in range(1, min(max_units, len(u) - 1) + 1):
                pre["".join(u[:L])] += 1; suf["".join(u[-L:])] += 1
        return pre, suf


def gen_slot(train_words, layout, seed=0, **kw):
    M = SlotModel(train_words, seed); n = sum(sum(p["lines"]) for p in layout)
    return fill(layout, [M.word() for _ in range(n)])


def modify(w, rng, p_sub=0.5):
    u = tokenize_glyphs(w)
    r = rng.random()
    if r < p_sub:
        idx = [i for i, g in enumerate(u) if g in CONFUSE]
        if idx:
            i = rng.choice(idx); u[i] = rng.choice(CONFUSE[u[i]]); return "".join(u)
    r = rng.random()
    if r < 0.4 and len(u) > 2: del u[rng.randrange(len(u))]; return "".join(u)
    if r < 0.7 and len(u) < 8: u.insert(rng.randrange(len(u) + 1), rng.choice(["e", "o", "a", "i", "l", "r", "d", "ch"])); return "".join(u)
    return "".join(u[:-1]) + rng.choice(["y", "dy", "aiin", "ain", "ar", "al", "or", "ol", "am", "ey", "edy"])


def gen_autocopy(train_words, layout, seed=0, window=12, p_mod=0.8, p_new=0.1, **kw):
    rng = random.Random(seed); M = SlotModel(train_words, seed); out_pages = []
    for p in layout:
        hist = []; lines = []
        for n in p["lines"]:
            ln = []
            for _ in range(n):
                if not hist or rng.random() < p_new: w = M.word()
                else:
                    w = rng.choice(hist[-window:])
                    if rng.random() < p_mod: w = modify(w, rng)
                ln.append(w); hist.append(w)
            lines.append(ln)
        out_pages.append(dict(page=p["page"], quire=p["quire"], hand=p["hand"], lang=p["lang"], I=p["I"], lines=lines))
    return out_pages


def gen_grille(train_words, layout, seed=0, rows=40, refill=60, p_skip=0.15, **kw):
    rng = random.Random(seed); M = SlotModel(train_words, seed); pre, suf = M.affixes()
    mid = M.fillers.get(2, collections.Counter()); mid.update(M.fillers.get(7, collections.Counter()))
    def table():
        return [(rng.choices(list(pre), list(pre.values()))[0], rng.choices(list(mid), list(mid.values()))[0] if rng.random() < 0.7 else "", rng.choices(list(suf), list(suf.values()))[0]) for _ in range(rows)]
    T = table(); r = 0; k = 0; words = []; n = sum(sum(p["lines"]) for p in layout)
    while len(words) < n:
        a, b, c = T[r % rows]
        w = (a if rng.random() > p_skip else "") + b + (c if rng.random() > p_skip else "")
        if w: words.append(w)
        r += 1 + (rng.random() < 0.3); k += 1
        if k % refill == 0: T = table()
    return fill(layout, words)


def plaintext(name="Latin", n=60000, seed=0):
    ws = full_language(name); rng = random.Random(seed); start = rng.randrange(0, max(1, len(ws) - n))
    return ws[start:start + n]


def gen_naibbe(train_words, layout, seed=0, k=8, p_bigram=0.5, n_bigrams=60, language="Latin", **kw):
    """each plaintext unit has k prefix strings and k suffix strings, drawn independently per use
    (as the card draws of the Naibbe cipher index the two affix tables independently)"""
    rng = random.Random(seed); M = SlotModel(train_words, seed); pre, suf = M.affixes()
    pre_l, pre_w = list(pre), list(pre.values()); suf_l, suf_w = list(suf), list(suf.values())
    n = sum(sum(p["lines"]) for p in layout)
    text = "".join(plaintext(language, seed=seed))
    letters = sorted(set(text)); bigrams = [a + b for (a, b), c in collections.Counter(zip(text, text[1:])).most_common(n_bigrams)]
    table = {}
    for u in letters + bigrams:
        table[u] = ([rng.choices(pre_l, pre_w)[0] for _ in range(k)], [rng.choices(suf_l, suf_w)[0] for _ in range(k)])
    words = []; i = 0
    while len(words) < n and i < len(text):
        if rng.random() < p_bigram and i + 1 < len(text) and text[i:i + 2] in table:
            u = text[i:i + 2]; i += 2
        else: u = text[i]; i += 1
        if u not in table: continue
        P, S = table[u]; words.append(rng.choice(P) + rng.choice(S))
    return fill(layout, words)


def gen_homonull(train_words, layout, seed=0, k=4, p_null=0.1, language="Latin", **kw):
    rng = random.Random(seed); n = sum(sum(p["lines"]) for p in layout)
    units = [u for u, c in collections.Counter(g for w in train_words for g in tokenize_glyphs(w)).most_common(29)]
    text = "".join(plaintext(language, seed=seed)); letters = sorted(set(text))
    pool = units * 4; rng.shuffle(pool); table = {}
    for l in letters: table[l] = [pool.pop() for _ in range(min(k, len(pool)))] or [rng.choice(units)]
    nulls = rng.sample(units, 3)
    lens = [len(tokenize_glyphs(w)) for w in train_words]
    stream = []
    for ch in text:
        if ch not in table: continue
        stream.append(rng.choice(table[ch]))
        if rng.random() < p_null: stream.append(rng.choice(nulls))
        if len(stream) > n * 7: break
    words = []; i = 0
    while len(words) < n and i < len(stream):
        L = rng.choice(lens); words.append("".join(stream[i:i + L])); i += L
    return fill(layout, words)


def gen_wordletter(train_words, layout, seed=0, k=4, language="Latin", **kw):
    rng = random.Random(seed); n = sum(sum(p["lines"]) for p in layout); M = SlotModel(train_words, seed)
    text = "".join(plaintext(language, seed=seed)); letters = sorted(set(text))
    vocab = list({M.word() for _ in range(60 * k * len(letters))}); rng.shuffle(vocab)
    table = {}
    for l in letters: table[l] = [vocab.pop() for _ in range(k)]
    words = [rng.choice(table[ch]) for ch in text if ch in table][:n]
    return fill(layout, words)


GENERATORS = {"slot": gen_slot, "autocopy": gen_autocopy, "grille": gen_grille, "naibbe": gen_naibbe, "naibbe_novowel": lambda t, l, seed=0, **kw: gen_naibbe(t, l, seed, language="Latin.novowel", **kw),
              "homonull": gen_homonull, "wordletter": gen_wordletter}