"""Held-out val sets for docs/MIX_DATA_LONG_MEMORY.md (same protocol as eval_windows.py:
contiguous 512-transition windows, state reset per window, eval/no_grad, normal + field-alone).

* val_varied : 6 windows, one per held-out book slice of data/heldout/val_varied.txt (never in train)
* val_dialog : 4 windows of data/heldout/val_dialogue.txt (held-out dialogue blocks, never in train
               verbatim; templated lines can recur in train)
* val2048    : the old fixed val (inside the old corpus, which is also ~25 % of the mix stream)
"""
import sys
from pathlib import Path
import numpy as np
import torch
HERE = Path(__file__).resolve().parent
ROOT = HERE.parents[2]
sys.path.insert(0, str(HERE.parent / "decfix_bptt"))
import dfb_common as C  # noqa: E402  (val2048, probe_chats, load_model)
from src.io.atomizer import Atomizer  # noqa: E402
EW = C.EW
N = 512


def _win(text_bytes, off):
    t = text_bytes[off: off + 4 * N].decode("utf-8", errors="ignore")
    pk = Atomizer(max_span_bytes=1).encode(t, reset=True)[: N + 1]
    return list(zip(pk[:-1], pk[1:]))


def varied_windows():
    raw = (ROOT / "data/heldout/val_varied.txt").read_bytes()
    out = []
    starts = [i for i in range(len(raw)) if raw.startswith(b"### ", i) and (i == 0 or raw[i - 1:i] == b"\n")]
    for s in starts:
        body = raw.index(b"\n", s) + 1
        out.append({"src": raw[s + 4: body - 1].decode(), "pairs": _win(raw, body + 2000)})
    return out


def dialog_windows():
    raw = (ROOT / "data/heldout/val_dialogue.txt").read_bytes()
    out = []
    for f in (0.1, 0.35, 0.6, 0.85):
        o = raw.index(b"\n\nUtilisateur:", int(len(raw) * f)) + 2
        out.append({"src": f"val_dialogue@{o}", "pairs": _win(raw, o)})
    return out


_CACHE = {}


def eval_set(m, name):
    if name not in _CACHE:
        _CACHE[name] = {"varied": varied_windows, "dialog": dialog_windows, "val2048": EW.windows}[name]()
    normal, alone = [], []
    for w in _CACHE[name]:
        normal += EW.run(m, w["pairs"], False); alone += EW.run(m, w["pairs"], True)
    return float(np.mean(normal)), float(np.mean(alone))


def all_vals(m):
    out = {}
    for n in ("dialog", "varied", "val2048"):
        a, b = eval_set(m, n)
        out[f"{n}_normal"] = a; out[f"{n}_field_alone"] = b
    return out