"""Dynamic padding + next-token label construction."""
from __future__ import annotations

import numpy as np
import torch

from ..tokenizers.multimodal import Sample


def collate(samples: list[Sample], pad_id: int, multiple: int = 8, *, target_length: int | None = None,
            batch_size: int | None = None) -> tuple[torch.Tensor, torch.Tensor]:
    """Right-pad to the batch max (rounded up to `multiple`) and build shifted labels.

    labels[j] = tokens[j+1] where j+1 >= target_start, else -100 (prompt and padding are not
    trained on).  Right padding is safe under causal attention: real tokens never attend to it.
    """
    L = max(len(s.tokens) for s in samples) - 1
    if target_length is not None:
        if target_length < L:
            raise ValueError("fixed padding cannot truncate a sample")
        L = target_length
    else:
        L = -(-L // multiple) * multiple
    B = len(samples) if batch_size is None else batch_size
    if B < len(samples):
        raise ValueError("fixed batch size cannot discard samples")
    inp = np.full((B, L), pad_id, dtype=np.int64)
    lab = np.full((B, L), -100, dtype=np.int64)
    for i, s in enumerate(samples):
        t = s.tokens.astype(np.int64)
        n = len(t) - 1
        inp[i, :n] = t[:-1]
        lab[i, :n] = t[1:]
        lab[i, : max(0, s.target_start - 1)] = -100
    return torch.from_numpy(inp), torch.from_numpy(lab)