"""T13. Context-length scaling of held-out predictability.
Interpolated Kneser-Ney n-gram models of order 1..8 over glyph units (with word spaces) and, in a
second run, over words; trained on the training quires, scored in bits per unit on the held-out
quires.  The same for surrogates (chars shuffled within words; words shuffled within lines; lines
shuffled across the book) and for reference languages at the same length (80/20 split).
Hypothesis H-STRUCT-CONTEXT: predictive information reaches beyond the word (gains continue past
order ~6 in glyph units and past order 1 in words, as in languages).  Writes out/context_scaling.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, math, sys
pass
from corpus import tokenize_glyphs
from harness import structured, split, SURROGATES, languages, log

ORDERS = range(1, 9)


class KN:
    """interpolated Kneser-Ney with absolute discount D=0.75, continuation counts for lower orders."""
    def __init__(self, seq, n):
        self.n = n; self.D = 0.75
        self.cnt = [collections.Counter() for _ in range(n + 1)]  # cnt[k][(ctx..., w)] with len ctx = k-1
        self.ctx = [collections.Counter() for _ in range(n + 1)]
        self.types = [collections.Counter() for _ in range(n + 1)]  # number of distinct continuations of ctx
        for k in range(1, n + 1):
            for i in range(len(seq) - k + 1):
                g = tuple(seq[i:i + k]); self.cnt[k][g] += 1
        # continuation counts for orders < n
        self.cont = [collections.Counter() for _ in range(n + 1)]
        for k in range(1, n):
            for g in self.cnt[k + 1]: self.cont[k][g[1:]] += 1
        for k in range(1, n + 1):
            src = self.cnt[k] if k == n else self.cont[k]
            for g, c in src.items():
                self.ctx[k][g[:-1]] += c; self.types[k][g[:-1]] += 1
        self.V = len(self.cnt[1]) + 1

    def prob(self, hist, w):
        p = 1.0 / self.V
        for k in range(1, self.n + 1):
            h = tuple(hist[-(k - 1):]) if k > 1 else ()
            src = self.cnt[k] if k == self.n else self.cont[k]
            tot = self.ctx[k].get(h, 0)
            if tot == 0: continue
            c = src.get(h + (w,), 0)
            p = max(c - self.D, 0) / tot + self.D * self.types[k][h] / tot * p
        return p

    def bits(self, seq):
        s = 0.0
        for i, w in enumerate(seq):
            s -= math.log2(self.prob(seq[max(0, i - self.n + 1):i], w))
        return s / len(seq)


def glyph_stream(pages):
    out = []
    for p in pages:
        for ln in p["lines"]:
            for w in ln: out += tokenize_glyphs(w) + [" "]
    return out


def word_stream(pages):
    return [w for p in pages for ln in p["lines"] for w in ln]


rows = []
pages = structured(); tr, te = split(pages)


def run(name, trs, tes, unit):
    r = dict(text=name, unit=unit, n_train=len(trs), n_test=len(tes))
    for n in ORDERS:
        m = KN(trs, n); r[f"k{n}"] = round(m.bits(tes), 4)
    print(r, flush=True); rows.append(r)


for unit, fn in (("glyph", glyph_stream), ("word", word_stream)):
    run("Voynich", fn(tr), fn(te), unit)
    for sname in ("chars-in-words", "words-in-lines", "lines-global"):
        sp = SURROGATES[sname](pages, seed=1); a, b = split(sp); run(f"Voynich {sname}", fn(a), fn(b), unit)
    L = languages()
    for lang in ("Latin", "Italian", "German", "Hebrew", "Arabic", "Malagasy", "Latin.novowel", "Italian.novowel"):
        ws = L[lang][:34217]; cut = int(len(ws) * 0.8)
        if unit == "glyph":
            a = [c for w in ws[:cut] for c in list(w) + [" "]]; b = [c for w in ws[cut:] for c in list(w) + [" "]]
        else: a, b = ws[:cut], ws[cut:]
        run(lang, a, b, unit)

import csv
with open("out/context_scaling.csv", "w", newline="") as f:
    w = csv.DictWriter(f, fieldnames=rows[0].keys()); w.writeheader(); w.writerows(rows)
v = [r for r in rows if r["text"] == "Voynich" and r["unit"] == "glyph"][0]; s = [r for r in rows if r["text"] == "Voynich lines-global" and r["unit"] == "glyph"][0]
log("H-STRUCT-CONTEXT", "T13-context-scaling", "held-out bits/glyph at order 1,2,4,6,8", f"{v['k1']},{v['k2']},{v['k4']},{v['k6']},{v['k8']}", f"lines shuffled: {s['k1']},{s['k2']},{s['k4']},{s['k6']},{s['k8']}", "see note", "out/context_scaling.csv")