"""Native right-padded batches. Supervision is returned separately."""
from __future__ import annotations
from rook.render import encode

MODEL_FIELDS = {"input_ids", "attention_mask", "candidate_positions", "candidate_mask", "query_positions"}


def collate_lists(encoded, pad_id: int, max_length: int):
    if not encoded or type(pad_id) is not int or pad_id < 0:
        raise ValueError("empty batch or missing padding token")
    length = ((max(len(e.ids) for e in encoded) + 7) // 8) * 8
    if length > max_length:
        raise ValueError("padded sequence exceeds explicit limit; no truncation")
    width = max(len(e.keys) for e in encoded)
    result = {key: [] for key in MODEL_FIELDS}
    for e in encoded:
        n, c = len(e.ids), len(e.keys)
        if not 2 <= c <= 255 or len(e.candidate_positions) != c or len(set(e.candidate_positions)) != c or not all(0 <= p < e.query_position == n-1 for p in e.candidate_positions):
            raise ValueError("invalid native endpoints/query")
        result["input_ids"].append(list(e.ids) + [pad_id] * (length-n))
        result["attention_mask"].append([1] * n + [0] * (length-n))
        result["candidate_positions"].append(list(e.candidate_positions) + [0] * (width-c))
        result["candidate_mask"].append([True] * c + [False] * (width-c))
        result["query_positions"].append(e.query_position)
    return result


def tensor_batch(inputs, tokenizer, device, max_length=1024):
    import torch
    encoded = [encode(item, tokenizer, max_length) for item in inputs]
    pad = tokenizer.pad_token_id
    if pad is None:
        pad = tokenizer.eos_token_id
    values = collate_lists(encoded, pad, max_length)
    return {k: torch.tensor(v, dtype=torch.bool if k == "candidate_mask" else torch.long, device=device) for k, v in values.items()}, encoded


def masked_ce(logits, mask, labels):
    import torch
    import torch.nn.functional as F
    if logits.dtype != torch.float32 or mask.dtype != torch.bool or labels.dtype != torch.long or logits.shape != mask.shape or labels.shape != (logits.shape[0],):
        raise ValueError("invalid candidate CE tensor contract")
    if not mask.any(1).all() or (labels < 0).any() or (labels >= logits.shape[1]).any() or not mask.gather(1, labels[:, None]).all():
        raise ValueError("label selects an invalid candidate")
    if not torch.isfinite(logits[mask]).all():
        raise FloatingPointError("nonfinite valid candidate logits")
    loss = F.cross_entropy(logits.masked_fill(~mask, -float("inf")), labels)
    if not torch.isfinite(loss):
        raise FloatingPointError("nonfinite candidate CE")
    return loss