"""(1a) Byte n-gram with interpolated Witten-Bell backoff, n=1..8.
Val = the exact ATOM 2048-transition val (4x512, context reset at each window start).
Two training sets: 'full' (whole corpus, same as ATOM which saw val ~4.5x) and
'heldout' (corpus with every occurrence-range of the 4 val windows cut out, +-4KB).
Also samples the 3 chat prompts (dialogue wrap, T=0.8, top-k 12, <=100 bytes, stop at \n\n).
"""
import json, math, random, sys, time
from collections import defaultdict
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent))
from common import OFFSETS, PROMPTS, ROOT, N, val_windows_bytes

OUT = Path(__file__).resolve().parent


def heldout_corpus():
    parts = []
    for f in ("data/corpus_fr_medium.txt", "data/corpus_fr_medium_s21.txt"):
        raw = (ROOT / f).read_bytes()
        cuts = sorted((o - 4096, o + 4 * N + 4096) for p, o in OFFSETS if p == f)
        pos = 0
        for a, b in cuts:
            parts.append(raw[pos:max(a, pos)]); pos = b
        parts.append(raw[pos:])
    return b"".join(parts)


class WB:
    def __init__(self, data: bytes, n: int):
        self.n = n
        self.c = [defaultdict(lambda: defaultdict(int)) for _ in range(n)]
        for i in range(len(data)):
            for k in range(n):
                if i - k < 0: break
                self.c[k][data[i - k:i]][data[i]] += 1
        self.tot = [{h: sum(d.values()) for h, d in ck.items()} for ck in self.c]

    def dist(self, hist: bytes):
        p = [1.0 / 256] * 256
        for k in range(self.n):
            if k > len(hist): break
            h = hist[len(hist) - k:] if k else b""
            d = self.c[k].get(h)
            if not d: break
            t = self.tot[k][h]; u = len(d)
            lam = t / (t + u)
            q = [(1 - lam) * x for x in p]
            for b, cnt in d.items():
                q[b] += lam * cnt / t
            p = q
        return p

    def ce(self, windows):
        tot, n = 0.0, 0
        for w in windows:
            for i in range(len(w) - 1):
                p = self.dist(w[max(0, i + 1 - (self.n - 1)): i + 1])
                tot += -math.log(max(p[w[i + 1]], 1e-12)); n += 1
        return tot / n

    def sample(self, prompt: str, rng, T=0.8, k=12, max_len=100):
        ctx = f"Utilisateur: {prompt}\nAssistant: ".encode()
        out = bytearray()
        while len(out) < max_len:
            p = self.dist((ctx + out)[-(self.n - 1):] if self.n > 1 else b"")
            logits = [math.log(max(x, 1e-12)) / T for x in p]
            top = sorted(range(256), key=lambda b: -logits[b])[:k]
            m = max(logits[b] for b in top)
            ws = [math.exp(logits[b] - m) for b in top]
            out.append(rng.choices(top, ws)[0])
            if out.endswith(b"\n\n"): break
        return out.decode("utf-8", "replace")


def main():
    wins = val_windows_bytes()
    full = (ROOT / "data/corpus_fr_medium.txt").read_bytes() + (ROOT / "data/corpus_fr_medium_s21.txt").read_bytes()
    held = heldout_corpus()
    res = {"val": "ATOM 2048-transition val (eval_windows.py), nats/byte", "full_bytes": len(full), "heldout_bytes": len(held), "rows": []}
    for name, data in (("full", full), ("heldout", held)):
        for n in (1, 2, 3, 4, 5, 6, 8):
            t0 = time.time(); m = WB(data, n); ce = m.ce(wins)
            row = {"train": name, "n": n, "val_ce": round(ce, 4), "sec": round(time.time() - t0, 1)}
            if n in (3, 5, 8) and name == "heldout":
                rng = random.Random(0)
                row["samples"] = {p: [m.sample(p, rng) for _ in range(3)] for p in PROMPTS}
            res["rows"].append(row); print(json.dumps(row, ensure_ascii=False), flush=True)
    (OUT / "ngram_results.json").write_text(json.dumps(res, indent=1, ensure_ascii=False))


if __name__ == "__main__":
    main()