"""The fit corpus: 16 short passages (8 roleplay, 8 novel prose), split 4 + 4 of each into a fit split and a
held-out split.

Each item becomes one sequence: a chat-template user turn asking the model to continue a passage (128 tokens
of it), then 128 reference tokens as the model's turn. The fit averages the fine-tune's edit over every
position of the fit split; the score compares next-token distributions at the reference positions.
The item texts and reference token ids are stored in data/fit_corpus.json, so the prompts are identical
wherever this runs.
"""
from __future__ import annotations

import json
from pathlib import Path

import numpy as np

DATA = Path(__file__).parent / "data" / "fit_corpus.json"
SPLITS = ("fit", "held_out")


class Corpus:
    """Token arrays and position masks for both splits.

    ids[s]      (8, T) token ids, right-padded
    valid[s]    per sequence, the number of real tokens T_valid (the sequences are cut to the shortest)
    starts[s]   per sequence, where the reference turn starts
    fit_mask[s] (8, T) 1 at the positions a vector is fitted over: 1 <= t < T_valid (BOS and padding excluded)
    rows[s]     (sequence, first, last): the positions whose next-token prediction is scored
    spans[s]    per sequence, its [start, end) in the concatenated scored rows
    """

    def __init__(self, tokenizer, path: Path = DATA, turns: dict[str, list[int]] | None = None):
        data = json.loads(Path(path).read_text())
        self.items = data["items"]
        n_ref = data["reference_tokens"]
        pad = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
        self.ids, self.valid, self.starts, self.fit_mask, self.rows, self.spans = {}, {}, {}, {}, {}, {}
        for s in SPLITS:
            seqs, starts, valid = [], [], []
            for reg in ("roleplay", "narrative"):
                group = []
                for it in (i for i in self.items if i["split"] == s and i["register"] == reg):
                    user = it["framing"].replace("{{TEXT}}", it["context"])
                    prompt = tokenizer.apply_chat_template([{"role": "user", "content": user}],
                                                           add_generation_prompt=True, return_dict=True)["input_ids"]
                    ref = (turns[it["id"]] if turns else it["reference_ids"])[:n_ref]
                    group.append(list(prompt) + list(ref))
                    starts.append(len(prompt))
                # each register's sequences are cut to that register's shortest
                T = min(len(x) for x in group)
                seqs += [x[:T] for x in group]
                valid += [T] * len(group)
            self.valid[s] = valid
            self.starts[s] = starts
            width = max(self.valid[s])
            self.ids[s] = np.array([x + [pad] * (width - len(x)) for x in seqs])
            self.fit_mask[s] = np.array([[1.0 if 1 <= t < v else 0.0 for t in range(width)] for v in self.valid[s]],
                                        np.float32)
            self.rows[s] = [(i, starts[i] - 1, self.valid[s][i] - 1) for i in range(len(seqs))]
            spans, o = [], 0
            for _, a, b in self.rows[s]:
                spans.append((o, o + b - a))
                o += b - a
            self.spans[s] = spans