"""Batches of rendered decisions, and the per-example augmentation for training."""

from __future__ import annotations

import random

import torch

from ..render import Renderer, Rendered
from ..schema import Decision


def collate(items: list[Rendered], pad_id: int, labels: list[int] | None = None) -> dict:
    """Right-padded ids, and candidate positions padded with -1.

    Right padding is safe for a causal backbone: padding sits after every
    position that is read, so nothing read has attended to it.
    """
    L = max(len(r.input_ids) for r in items)
    K = max(len(r.option_positions) for r in items)
    ids = torch.full((len(items), L), pad_id, dtype=torch.long)
    mask = torch.zeros((len(items), L), dtype=torch.long)
    pos = torch.full((len(items), K), -1, dtype=torch.long)
    for i, r in enumerate(items):
        ids[i, :len(r.input_ids)] = torch.tensor(r.input_ids)
        mask[i, :len(r.input_ids)] = 1
        pos[i, :len(r.option_positions)] = torch.tensor(r.option_positions)
    out = {"input_ids": ids, "attention_mask": mask, "option_pos": pos}
    if labels is not None:
        out["labels"] = torch.tensor(labels, dtype=torch.long)
    return out


def training_view(d: Decision, rng: random.Random, max_options: int = 24
                  ) -> tuple[list[str], int]:
    """(candidates, target index) as the model sees this record in one epoch.

    Choice candidates are shuffled, so no position carries the answer, and a
    long label set (CLINC has 151 intents) is subsampled to between 4 and
    `max_options`, gold always included, so the model learns from lists of
    every length. Scales and yes/no keep their order: that order is meaning.
    """
    cands, t = d.candidates(), d.target()
    if d.type != "choice":
        return cands, t
    gold = cands[t]
    others = [c for i, c in enumerate(cands) if i != t]
    if len(cands) > max_options:
        k = rng.randint(min(4, max_options), max_options) - 1
        others = rng.sample(others, k)
    picked = others + [gold]
    rng.shuffle(picked)
    return picked, picked.index(gold)


def render_for_training(r: Renderer, d: Decision, rng: random.Random,
                        max_options: int = 24) -> tuple[Rendered, int]:
    cands, t = training_view(d, rng, max_options)
    return r.render(d, cands), t