"""Shared dataset primitives: examples, vocabularies, split containers and hashing.

Every example is a token sequence ``<bos> <prompt tokens> = <a_1 ... a_k>`` where the ``k`` answer
tokens form a *scratchpad* (one intermediate result per execution step; k = 1 for single-answer
tasks). The primary metric is exact match of the final answer token under greedy decoding; the
trace metric is exact match of all ``k`` answer tokens. ``cluster`` identifies the
generating template (e.g. the operator program) and is the unit of resampling in the cluster
bootstrap: items sharing a cluster are not treated as independent.
"""
from __future__ import annotations

import hashlib
import json
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Iterable

import numpy as np

PAD, BOS, EQ = "<pad>", "<bos>", "="


@dataclass(frozen=True)
class Example:
    tokens: tuple[str, ...]  # full sequence including the answer tokens
    n_answer: int  # number of answer (scratchpad) tokens at the end of ``tokens``
    cluster: str  # template / program identifier (bootstrap unit)
    meta: dict = field(default_factory=dict)

    @property
    def prompt(self) -> tuple[str, ...]:
        return self.tokens[: -self.n_answer]

    @property
    def answers(self) -> tuple[str, ...]:
        return self.tokens[-self.n_answer :]

    @property
    def answer(self) -> str:
        return self.tokens[-1]

    @property
    def item_hash(self) -> str:
        return hashlib.sha256(" ".join(self.tokens).encode()).hexdigest()[:16]


class Vocab:
    def __init__(self, tokens: Iterable[str]) -> None:
        base = [PAD, BOS, EQ]
        extra = [t for t in dict.fromkeys(tokens) if t not in base]
        self.itos: list[str] = base + extra
        self.stoi: dict[str, int] = {t: i for i, t in enumerate(self.itos)}

    def __len__(self) -> int:
        return len(self.itos)

    @property
    def pad_id(self) -> int:
        return self.stoi[PAD]

    def encode(self, tokens: Iterable[str]) -> list[int]:
        return [self.stoi[t] for t in tokens]

    def decode(self, ids: Iterable[int]) -> list[str]:
        return [self.itos[i] for i in ids]

    def to_json(self) -> str:
        return json.dumps(self.itos)

    @classmethod
    def from_json(cls, text: str) -> "Vocab":
        return cls(json.loads(text))


@dataclass
class Split:
    name: str
    examples: list[Example]
    generator: str
    generator_version: str
    seed: int
    config: dict

    def __len__(self) -> int:
        return len(self.examples)

    @property
    def hash(self) -> str:
        return split_hash(self.examples)

    def item_hashes(self) -> set[str]:
        return {e.item_hash for e in self.examples}

    def answer_distribution(self) -> dict[str, float]:
        vals, counts = np.unique([e.answer for e in self.examples], return_counts=True)
        return {str(v): float(c) / len(self.examples) for v, c in zip(vals, counts)}

    def save(self, path: Path) -> None:
        path.parent.mkdir(parents=True, exist_ok=True)
        payload = {
            "name": self.name,
            "generator": self.generator,
            "generator_version": self.generator_version,
            "seed": self.seed,
            "config": self.config,
            "hash": self.hash,
            "n": len(self.examples),
            "examples": [asdict(e) for e in self.examples],
        }
        path.write_text(json.dumps(payload, sort_keys=True))

    @classmethod
    def load(cls, path: Path) -> "Split":
        d = json.loads(path.read_text())
        ex = [Example(tuple(e["tokens"]), int(e["n_answer"]), e["cluster"], e.get("meta", {})) for e in d["examples"]]
        s = cls(d["name"], ex, d["generator"], d["generator_version"], d["seed"], d["config"])
        if s.hash != d["hash"]:
            raise ValueError(f"split hash mismatch for {path}: stored {d['hash']} computed {s.hash}")
        return s


def split_hash(examples: Iterable[Example]) -> str:
    h = hashlib.sha256()
    for e in examples:
        h.update((" ".join(e.tokens) + "\n").encode())
    return h.hexdigest()[:16]


def assert_disjoint(*splits: Split) -> None:
    seen: dict[str, str] = {}
    for s in splits:
        for e in s.examples:
            if e.item_hash in seen:
                raise AssertionError(f"item leakage: {e.tokens} present in {seen[e.item_hash]} and {s.name}")
            seen[e.item_hash] = s.name


def encode_split(split: Split, vocab: Vocab, max_len: int) -> dict[str, np.ndarray]:
    """Encode a split for teacher-forced training and greedy evaluation.

    Returns a dict of arrays:
      ``ids``        [N,T] full sequences (prompt + answers), right-padded with ``<pad>``;
      ``loss_mask``  [N,T] 1 at positions whose *next* token is an answer token;
      ``prompt_len`` [N]   length of the prompt (index of the first answer token);
      ``n_answer``   [N]   number of answer tokens;
      ``answers``    [N,K] answer token ids right-padded with ``<pad>`` (K = max n_answer).
    """
    n = len(split.examples)
    K = max(e.n_answer for e in split.examples)
    ids = np.full((n, max_len), vocab.pad_id, dtype=np.int64)
    loss_mask = np.zeros((n, max_len), dtype=np.float32)
    prompt_len = np.zeros(n, dtype=np.int64)
    n_answer = np.zeros(n, dtype=np.int64)
    answers = np.full((n, K), vocab.pad_id, dtype=np.int64)
    for i, e in enumerate(split.examples):
        seq = vocab.encode(e.tokens)
        if len(seq) > max_len:
            raise ValueError(f"sequence longer than max_len={max_len}: {e.tokens}")
        ids[i, : len(seq)] = seq
        pl = len(e.prompt)
        assert e.prompt[-1] == EQ
        prompt_len[i] = pl
        n_answer[i] = e.n_answer
        loss_mask[i, pl - 1 : pl - 1 + e.n_answer] = 1.0
        answers[i, : e.n_answer] = vocab.encode(e.answers)
    return {"ids": ids, "loss_mask": loss_mask, "prompt_len": prompt_len, "n_answer": n_answer, "answers": answers}