"""Loop outputs and symbolic readout semantics."""

from dataclasses import dataclass

import torch
from torch import Tensor


@dataclass
class RecurrentOutput:
    """Tuple index t-1 represents loop t; h_0 is not a decoded start symbol.

    Logits are [batch, vocab] in answer mode or [batch, sequence, vocab]
    in all-token mode. Optional hidden states are [batch, sequence, hidden].
    Margins are [batch, loops]. Tensors retain autograd unless the caller
    disables it; observations never feed back into the recurrent trajectory.
    """

    loop_logits: tuple[Tensor, ...]
    hidden_states: tuple[Tensor, ...] | None = None
    initial_hidden_state: Tensor | None = None
    margins: Tensor | None = None
    state_logits: Tensor | None = None  # Direct R readout [B,T,26], never fed back.
    stop_logits: Tensor | None = None  # [batch, executed_loops], when completion is enabled.

    @property
    def logits(self) -> Tensor:
        """The final loop's logits."""
        return self.loop_logits[-1]


def answer_margin(logits: Tensor, labels: Tensor, allowed_token_ids: Tensor) -> Tensor:
    """Correct raw logit minus the best *other allowed* logit, per example.

    Shapes: logits [batch, vocab], labels [batch] (token IDs), allowed IDs
    [answers]. Labels must belong to a unique answer set of at least two IDs.
    A tie has margin zero, so it is not margin-based correctness.
    """
    if logits.ndim != 2 or labels.shape != logits.shape[:1]:
        raise ValueError("Expected logits [batch, vocab] and labels [batch]")
    if labels.dtype != torch.long or allowed_token_ids.dtype != torch.long:
        raise ValueError("Labels and allowed_token_ids must be torch.long token IDs")
    if labels.device != logits.device or allowed_token_ids.device != logits.device:
        raise ValueError("Logits, labels, and allowed_token_ids must share a device")
    if allowed_token_ids.ndim != 1 or allowed_token_ids.numel() < 2:
        raise ValueError("Provide at least two allowed token IDs in a 1D tensor")
    if allowed_token_ids.unique().numel() != allowed_token_ids.numel():
        raise ValueError("Allowed token IDs must be unique")
    if (allowed_token_ids < 0).any() or (allowed_token_ids >= logits.shape[-1]).any():
        raise ValueError("Allowed token IDs must be within the model vocabulary")
    correct = labels[:, None] == allowed_token_ids[None, :]
    if not correct.any(dim=-1).all():
        raise ValueError("Every label must be an allowed token ID")
    scores = logits.index_select(-1, allowed_token_ids)
    wrong = scores.masked_fill(correct, -torch.inf).max(dim=-1).values
    return logits.gather(-1, labels[:, None]).squeeze(-1) - wrong