"""Chancery mixed keys of the Tranchedino type, for simulation (T19) and planted controls (T22).

A key has four parts, each sign a string of Voynich glyph units drawn from the manuscript's own
inventory of word-initial / word-final strings and slot fillers:
  letters      each plaintext letter -> k_v signs (vowels) or k_c signs (consonants)
  syllables    the n_syl commonest within-word letter pairs of the plaintext -> s_forms signs each
  nomenclator  the n_nom commonest plaintext words -> one whole-word sign each (a slot-grammar word)
  nulls        n_null signs that encode nothing, inserted at rate p_null per sign
Encoding a plaintext word: if it is in the nomenclator, its word sign (with probability p_nom);
otherwise left to right, a syllable sign where one applies (with probability p_syl), else a letter
sign.  Choice among a unit's forms: 'random' (uniform), 'cycle' (each unit rotates through its forms),
'sticky' (repeat the unit's previous form with probability 0.7).
Spacing practice:
  word    one written word per plaintext word
  split   each plaintext word's signs broken into written words of one or two signs
  stream  plaintext spaces dropped; signs regrouped into written words of 1-3 signs
Every written word comes with its ground truth: the list of (sign, unit) pairs it contains.
"""
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, sys
from corpus import tokenize_glyphs
from generators import SlotModel, fill

VOWELS = set("aeiouy")


class MixedKey:
    def __init__(self, voy_words, plain_words_for_tables, seed=0, k_v=3, k_c=2, n_syl=60, s_forms=1, n_nom=80, n_null=4,
                 max_sign_units=3, positional=False, n_tri=0, unique_signs=False):
        """positional: each unit has one set of signs for the first position of a plaintext word (drawn
        from the manuscript's word-initial strings) and another for later positions (word-final strings),
        as the Naibbe tables separate prefixes from suffixes"""
        rng = random.Random(seed); self.rng = rng; self.positional = positional
        M = SlotModel(voy_words, seed); pre, suf = M.affixes(max_units=max_sign_units)
        pool = collections.Counter(); pool.update(pre); pool.update(suf)
        # slot-position pools: what the manuscript writes at the start, middle and end of a word
        from units import G1
        ini, mid, fin = collections.Counter(), collections.Counter(), collections.Counter()
        for w in voy_words:
            m = G1.match(w)
            if not m: continue
            g = m.groups()
            if g[0] + g[1]: ini[g[0] + g[1]] += 1
            if g[2] + g[3]: mid[g[2] + g[3]] += 1
            if g[4]: fin[g[4]] += 1
        pools = {"any": pool, "pre": pre, "suf": suf, "ini": ini, "mid": mid, "fin": fin}
        lists = {k: [s for s, c in p.most_common() if c >= 2] for k, p in pools.items()}
        used_by = {k: set() for k in pools}  # signs in different position tables may share strings
        # unique_signs: no string is used in two tables, so every sign has one reading (as in the
        # historical keys); the fingerprint grid was run without it (a string could carry one reading at
        # the start of a word and another at the end, which a reader who knows the rule can still decode)
        used = set()  # every sign, for the nomenclator's uniqueness check

        def draw(which="any"):
            P = pools[which]; strings = lists[which]; weights = [P[s] for s in strings]; U = used_by[which]
            for _ in range(5000):
                s = rng.choices(strings, weights)[0]
                if s not in U and not (unique_signs and s in used): U.add(s); used.add(s); return s
            for s in strings:
                if s not in U and not (unique_signs and s in used): U.add(s); used.add(s); return s
            for _ in range(5000):  # inventory exhausted: compose two strings
                s = rng.choices(strings, weights)[0] + rng.choices(strings, weights)[0]
                if s not in U and not (unique_signs and s in used): U.add(s); used.add(s); return s
            raise RuntimeError("sign inventory exhausted")

        text = [w for w in plain_words_for_tables if w]
        letters = sorted(set("".join(text)))
        # tables by position in the plaintext word: positional=False one table; True initial/later;
        # "slots" initial / medial / final, drawn from the manuscript's own slot fillers
        classes = {False: {"any": "any"}, True: {"ini": "pre", "any": "suf"}, "slots": {"ini": "ini", "any": "mid", "fin": "fin"}}[positional]
        self.tables = {c: {} for c in classes}
        units_all = list(letters)
        pairs = collections.Counter(w[i:i + 2] for w in text for i in range(len(w) - 1))
        self.syllables = [p for p, c in pairs.most_common(n_syl)]
        tris = collections.Counter(w[i:i + 3] for w in text for i in range(len(w) - 2))
        self.syllables += [t for t, c in tris.most_common(n_tri)]
        for u in units_all + self.syllables:
            k = s_forms if len(u) > 1 else (k_v if u in VOWELS else k_c)
            for c, pool_name in classes.items(): self.tables[c][u] = [draw(pool_name) for _ in range(k)]
        self.unit_signs = self.tables["any"]; self.unit_signs_init = self.tables.get("ini", {})
        self.nulls = [draw() for _ in range(n_null)]
        wc = collections.Counter(text); self.nomen = {}
        nomen_used = set()
        for w, c in wc.most_common(n_nom):
            for _ in range(1000):
                s = M.word()
                if s not in used and s not in nomen_used and len(tokenize_glyphs(s)) >= 2: break
            nomen_used.add(s); used.add(s); self.nomen[w] = s
        self.sign_to_unit = {}
        for c, T in self.tables.items():
            for u, ss in T.items():
                for s in ss: self.sign_to_unit.setdefault(s, u)  # a string shared by two position tables keeps its first reading
        for s in self.nulls: self.sign_to_unit[s] = ""
        for w, s in self.nomen.items(): self.sign_to_unit[s] = "<" + w + ">"
        self.last = {}

    def choose(self, unit, mode, cls="any"):
        T = self.tables.get(cls, self.tables["any"]); forms = T[unit]; key = (unit, cls)
        if len(forms) == 1 or mode == "random": return self.rng.choice(forms)
        if mode == "cycle":
            i = (self.last.get(key, -1) + 1) % len(forms); self.last[key] = i; return forms[i]
        if mode == "sticky":
            if key in self.last and self.rng.random() < 0.7: return forms[self.last[key]]
            i = self.rng.randrange(len(forms)); self.last[key] = i; return forms[i]
        raise ValueError(mode)

    def encode_word(self, w, p_syl=0.7, p_nom=1.0, p_null=0.05, choice="random"):
        """list of (sign, unit) pairs for one plaintext word"""
        out = []
        if w in self.nomen and self.rng.random() < p_nom:
            out.append((self.nomen[w], "<" + w + ">"))
        else:
            i = 0; units = []
            while i < len(w):
                if i + 2 < len(w) and w[i:i + 3] in self.unit_signs and self.rng.random() < p_syl:
                    units.append(w[i:i + 3]); i += 3
                elif i + 1 < len(w) and w[i:i + 2] in self.unit_signs and self.rng.random() < p_syl:
                    units.append(w[i:i + 2]); i += 2
                elif w[i] in self.unit_signs:
                    units.append(w[i]); i += 1
                else: i += 1
            for j, u in enumerate(units):
                cls = "ini" if j == 0 else ("fin" if j == len(units) - 1 and len(units) > 1 else "any")
                out.append((self.choose(u, choice, cls), u))
        if p_null > 0:
            res = []
            for pr in out:
                if self.rng.random() < p_null: res.append((self.rng.choice(self.nulls), ""))
                res.append(pr)
            out = res
        return out

    def encipher(self, plain_words, spacing="word", choice="random", p_syl=0.7, p_nom=1.0, p_null=0.05, n_out=None):
        """returns written words and, for each, its ground truth [(sign, unit), ...]"""
        written, truth = [], []
        if spacing == "stream":
            stream = []
            for w in plain_words:
                stream += self.encode_word(w, p_syl, p_nom, p_null, choice)
                if n_out and len(stream) > n_out * 3: break
            i = 0
            while i < len(stream):
                k = self.rng.choices([1, 2, 3], [0.35, 0.45, 0.20])[0]
                chunk = stream[i:i + k]; i += k
                written.append("".join(s for s, u in chunk)); truth.append(chunk)
                if n_out and len(written) >= n_out: break
            return written, truth
        for w in plain_words:
            enc = self.encode_word(w, p_syl, p_nom, p_null, choice)
            if not enc: continue
            if spacing == "word":
                written.append("".join(s for s, u in enc)); truth.append(enc)
            elif spacing == "split":
                i = 0
                while i < len(enc):
                    k = 1 if len(enc) - i == 1 else self.rng.choice([1, 2, 2])
                    chunk = enc[i:i + k]; i += k
                    written.append("".join(s for s, u in chunk)); truth.append(chunk)
            else: raise ValueError(spacing)
            if n_out and len(written) >= n_out: break
        return written, truth


def gen_tranchedino(train_words, layout, seed=0, language="Latin", spacing="word", choice="random", k_v=3, k_c=2, n_syl=60, n_nom=80,
                    n_null=4, p_syl=0.7, p_null=0.05, max_sign_units=3, positional=False, **kw):
    from generators import plaintext
    plain = plaintext(language, n=80000, seed=seed)
    K = MixedKey(train_words, plain, seed=seed, k_v=k_v, k_c=k_c, n_syl=n_syl, n_nom=n_nom, n_null=n_null, max_sign_units=max_sign_units, positional=positional)
    n = sum(sum(p["lines"]) for p in layout)
    written, truth = K.encipher(plain, spacing=spacing, choice=choice, p_syl=p_syl, p_null=p_null, n_out=n)
    return fill(layout, written)


def sectioned_plaintext(layout, seed=0, src_dir=os.path.join(REF, "sections")):
    """plaintext words in manuscript order, each page drawing from a Latin source of its section's
    genre (herbal: Isidore 17, Columella; astronomical and zodiac: Hyginus; cosmological and
    balneological: Pliny NH 2, Isidore 13-14; pharmaceutical: Isidore 4, Apicius; recipes: Apicius,
    Cato), read on consecutively within each source from a random start"""
    rng = random.Random(seed); src = {}; pos = {}
    out = []
    for p in layout:
        sec = p.get("I") or "T"
        if sec not in src:
            src[sec] = open(f"{src_dir}/{sec}.txt").read().split(); pos[sec] = rng.randrange(len(src[sec]))
        n = int(sum(p["lines"]) * 1.3) + 5  # plaintext words per page: more than written words, since nomenclator and splits vary
        S = src[sec]
        for _ in range(n): out.append((S[pos[sec] % len(S)], p["page"])); pos[sec] += 1
    return out


def gen_tranchedino_sectioned(train_words, layout, seed=0, spacing="word", choice="random", k_v=3, k_c=2, n_syl=60, n_nom=80,
                              n_null=4, p_syl=0.7, p_null=0.05, max_sign_units=3, positional="slots", n_tri=0, return_truth=False, **kw):
    """the mixed key applied page by page to a plaintext whose sections match the manuscript's"""
    plain = sectioned_plaintext(layout, seed)
    words = [w for w, pg in plain]
    K = MixedKey(train_words, words, seed=seed, k_v=k_v, k_c=k_c, n_syl=n_syl, n_nom=n_nom, n_null=n_null, max_sign_units=max_sign_units, positional=positional, n_tri=n_tri)
    by_page = collections.defaultdict(list)
    for w, pg in plain: by_page[pg].append(w)
    out = []; truths = []
    for p in layout:
        need = sum(p["lines"]); written, truth = K.encipher(by_page[p["page"]], spacing=spacing, choice=choice, p_syl=p_syl, p_null=p_null, n_out=need)
        it = iter(zip(written, truth)); lines = []
        for n in p["lines"]:
            ln = [x for _, x in zip(range(n), it)]
            if ln: lines.append([w for w, t in ln]); truths += [t for w, t in 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, truths, K) if return_truth else out