"""Utilities for logits-derived token scoring."""

from __future__ import annotations

import torch
import torch.nn.functional as F
from torch import Tensor


def apply_temperature(logits: Tensor, temperature: float | None = None) -> Tensor:
    """Scale logits by temperature for calibration-style scoring."""

    if temperature is None:
        return logits
    if temperature <= 0:
        raise ValueError("Temperature must be positive.")
    if temperature == 1.0:
        return logits
    return logits / temperature


def logits_to_log_probs(logits: Tensor, temperature: float | None = None) -> Tensor:
    """Convert logits into log-probabilities in a numerically stable way."""

    return F.log_softmax(apply_temperature(logits, temperature=temperature), dim=-1)


def logits_to_probs(logits: Tensor, temperature: float | None = None) -> Tensor:
    """Convert logits into probabilities."""

    return logits_to_log_probs(logits, temperature=temperature).exp()


def select_final_token(tensor: Tensor, attention_mask: Tensor | None = None) -> Tensor:
    """Select the final valid token position from a sequence tensor."""

    if tensor.ndim < 2:
        raise ValueError("Expected a tensor with shape [batch, seq, ...].")

    batch_size, seq_len = tensor.shape[:2]
    if attention_mask is None:
        positions = torch.full(
            (batch_size,),
            seq_len - 1,
            dtype=torch.long,
            device=tensor.device,
        )
    else:
        if attention_mask.shape != (batch_size, seq_len):
            raise ValueError("attention_mask must match the leading [batch, seq] tensor shape.")
        mask = attention_mask.to(device=tensor.device, dtype=torch.long)
        token_positions = torch.arange(seq_len, device=tensor.device, dtype=torch.long).unsqueeze(0)
        positions = (mask * token_positions).amax(dim=-1)

    gather_index = positions.view(batch_size, 1, *([1] * (tensor.ndim - 2)))
    gather_index = gather_index.expand(batch_size, 1, *tensor.shape[2:])
    return tensor.gather(dim=1, index=gather_index).squeeze(1)


def gather_token_scores(scores: Tensor, token_ids: Tensor) -> Tensor:
    """Gather per-token scores from the vocabulary dimension."""

    if scores.ndim < 2:
        raise ValueError("Expected a score tensor with a vocabulary dimension.")
    if token_ids.ndim != scores.ndim - 1:
        raise ValueError("token_ids must match the score tensor without the vocabulary dimension.")
    if scores.shape[:-1] != token_ids.shape:
        raise ValueError("token_ids must match the leading score tensor dimensions.")

    token_ids = token_ids.to(device=scores.device, dtype=torch.long).unsqueeze(-1)
    return scores.gather(dim=-1, index=token_ids).squeeze(-1)
