"""Turning an SFT jsonl into batches, and counting what the model actually saw.

Three things here are decisions rather than plumbing.

**One example per sequence.** Pretraining packs documents end to end into 2048
tokens, which is right when the objective is "predict the next token of a
stream". It is wrong here: packed example B would attend to example A, and the
loss on B would depend on which example happened to precede it. For a diagnostic
whose whole output is "what does N examples buy you", that is a confound bought
for a throughput gain we do not need -- the largest arm is 50,000 examples
against a billion pretraining tokens.

**Length bucketing, not sorted batching.** Padding every batch to the longest
example in a random draw wastes most of the tensor when lengths are skewed.
Sorting the whole corpus by length would fix that and introduce a worse problem:
every batch would then be homogeneous in length, and length correlates with
family here -- tool examples carry schemas and are long, chitchat is short. The
model would see batches of nearly pure tools followed by batches of nearly pure
chitchat. So examples are shuffled, cut into large windows, sorted *within* a
window, and the resulting micro-batches shuffled again. Padding stays low and no
batch is a single family.

**Counting.** The experiment asks for capability per example, per SFT token and
per GPU-second, so all three are tracked, and "SFT tokens" means **supervised**
tokens -- the ones that carried loss -- not the length of the padded tensor.
Reporting the padded figure would inflate the denominator by roughly the padding
ratio and flatter every arm equally, which is worse than useless because it
looks fine.
"""

from __future__ import annotations

import json
import random
from collections.abc import Iterator
from dataclasses import dataclass, field
from pathlib import Path

import torch

from .chat import IGNORE_INDEX, ChatTemplate, Message, RenderedChat

__all__ = ["SFTExample", "SFTDataset", "SFTBatch", "load_jsonl",
           "step_weights"]


@dataclass
class SFTExample:
    example_id: str
    family: str
    subfamily: str
    lang: str
    messages: list[dict]
    source: dict = field(default_factory=dict)

    def to_messages(self) -> list[Message]:
        return [Message(role=m["role"], content=m["content"]) for m in self.messages]


def load_jsonl(path: str | Path) -> list[SFTExample]:
    out = []
    with open(path, encoding="utf-8") as fh:
        for line in fh:
            if not line.strip():
                continue
            row = json.loads(line)
            out.append(SFTExample(
                example_id=row.get("example_id", ""),
                family=row["family"],
                subfamily=row.get("subfamily", row["family"]),
                lang=row.get("lang", "en"),
                messages=row["messages"],
                source=row.get("source", {}),
            ))
    return out


@dataclass
class SFTBatch:
    input_ids: torch.Tensor      # [B, L]
    labels: torch.Tensor         # [B, L], -100 where masked
    #: Supervised tokens in this batch. The denominator for "loss per token"
    #: and the thing that accumulates into `sft_tokens`.
    n_supervised: int
    n_examples: int
    #: Padding tokens, tracked so the bucketing can be shown to work rather
    #: than assumed to.
    n_padding: int

    def to(self, device) -> SFTBatch:
        return SFTBatch(
            input_ids=self.input_ids.to(device, non_blocking=True),
            labels=self.labels.to(device, non_blocking=True),
            n_supervised=self.n_supervised,
            n_examples=self.n_examples,
            n_padding=self.n_padding,
        )


class SFTDataset:
    """Rendered examples, batched deterministically from a seed."""

    def __init__(
        self,
        examples: list[SFTExample],
        template: ChatTemplate,
        *,
        micro_batch_size: int = 8,
        bucket_windows: int = 32,
        seed: int = 1234,
    ) -> None:
        self.template = template
        self.micro_batch_size = int(micro_batch_size)
        self.bucket_size = self.micro_batch_size * int(bucket_windows)
        self.seed = int(seed)

        self.rendered: list[RenderedChat] = []
        self.kept: list[SFTExample] = []
        self.dropped: list[tuple[str, str]] = []

        for ex in examples:
            try:
                r = template.render(ex.to_messages())
            except ValueError as exc:
                # A malformed conversation is a data bug, and one that would
                # otherwise show up as "SFT did not work". Recorded with its
                # reason rather than skipped.
                self.dropped.append((ex.example_id, f"malformed: {exc}"))
                continue
            if r is None:
                self.dropped.append((ex.example_id, "too long"))
                continue
            if r.n_supervised == 0:
                self.dropped.append((ex.example_id, "nothing supervised"))
                continue
            self.rendered.append(r)
            self.kept.append(ex)

    # ------------------------------------------------------------------ #

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

    @property
    def total_supervised_tokens(self) -> int:
        return sum(r.n_supervised for r in self.rendered)

    @property
    def total_tokens(self) -> int:
        """Every token the model reads, supervised or not. The prompt is most
        of it, and it is real compute even though it carries no loss."""
        return sum(len(r) for r in self.rendered)

    def stats(self) -> dict:
        lengths = [len(r) for r in self.rendered]
        sup = [r.n_supervised for r in self.rendered]
        by_family: dict[str, int] = {}
        for ex in self.kept:
            by_family[ex.family] = by_family.get(ex.family, 0) + 1
        return {
            "examples": len(self.kept),
            "dropped": len(self.dropped),
            "drop_reasons": _tally(reason for _, reason in self.dropped),
            "by_family": by_family,
            "by_lang": _tally(ex.lang for ex in self.kept),
            "tokens_total": sum(lengths),
            "tokens_supervised": sum(sup),
            "supervised_frac": (sum(sup) / sum(lengths)) if lengths else 0.0,
            "len_mean": (sum(lengths) / len(lengths)) if lengths else 0.0,
            "len_p50": _pct(lengths, 50),
            "len_p95": _pct(lengths, 95),
            "len_max": max(lengths) if lengths else 0,
        }

    # ------------------------------------------------------------------ #

    def epoch_batches(self, epoch: int) -> list[list[int]]:
        """Index batches for one epoch. Deterministic in `(seed, epoch)`.

        Deterministic because two branches that differ only in dataset size
        must not also differ in the accident of their shuffle, and because a
        resumed run has to replay the order it was on.
        """
        # An explicit formula, not `hash((seed, epoch))`. Python's tuple hash
        # happens to be stable for ints, but it is stable by accident of the
        # implementation rather than by promise, and the order of a training
        # run is not a thing to leave resting on that.
        rng = random.Random(self.seed * 1_000_003 + epoch)
        order = list(range(len(self.rendered)))
        rng.shuffle(order)

        batches: list[list[int]] = []
        for start in range(0, len(order), self.bucket_size):
            window = order[start:start + self.bucket_size]
            window.sort(key=lambda i: len(self.rendered[i]))
            for j in range(0, len(window), self.micro_batch_size):
                batches.append(window[j:j + self.micro_batch_size])
        rng.shuffle(batches)
        return batches

    def collate(self, indices: list[int]) -> SFTBatch:
        chosen = [self.rendered[i] for i in indices]
        ids, labels = self.template.pad_batch(chosen)
        input_ids = torch.tensor(ids, dtype=torch.long)
        label_t = torch.tensor(labels, dtype=torch.long)
        return SFTBatch(
            input_ids=input_ids,
            labels=label_t,
            n_supervised=int((label_t != IGNORE_INDEX).sum()),
            n_examples=len(chosen),
            n_padding=int(input_ids.numel() - sum(len(r) for r in chosen)),
        )

    def iter_epoch(self, epoch: int) -> Iterator[SFTBatch]:
        for batch in self.epoch_batches(epoch):
            yield self.collate(batch)

    def steps_per_epoch(self) -> int:
        return len(self.epoch_batches(0))


def step_weights(batches: list[SFTBatch]) -> list[float]:
    """Each microbatch's share of its optimizer step's supervised tokens.

    The one arithmetic fact that makes an accumulated step equal to a single
    backward pass over the whole step:

        F.cross_entropy(..., ignore_index=-100)

    returns the mean over the *valid* tokens of the tensor it was given. So
    scaling microbatch k by `n_k / sum(n)` and summing yields

        sum_k sum_i CE_ki / sum_k n_k

    which is the token-weighted mean. Scaling by `1/K` instead yields a mean of
    means, and the two agree only when every microbatch carries the same amount
    of supervision.

    With assistant-only loss they never do. Measured on this corpus, supervised
    tokens are 7.4% of an abstention sequence and 57% of an instruction one, and
    length bucketing groups similar examples together -- so the imbalance is
    systematic, not noise. Lives here, beside the batches it weights, so the
    trainer and `tests/test_sft_grad_accum.py` are provably running the same
    arithmetic rather than two copies of it.

    A microbatch with no supervision gets weight 0: its loss is NaN (a mean over
    zero terms) and it must not enter the sum.
    """
    total = sum(b.n_supervised for b in batches)
    if total == 0:
        return [0.0] * len(batches)
    return [b.n_supervised / total for b in batches]


def _tally(values) -> dict:
    out: dict = {}
    for v in values:
        out[v] = out.get(v, 0) + 1
    return dict(sorted(out.items(), key=lambda kv: -kv[1]))


def _pct(values: list[int], p: int) -> int:
    if not values:
        return 0
    s = sorted(values)
    return s[min(len(s) - 1, int(len(s) * p / 100))]