"""Utility functions for mt_dao supported models."""

import torch
from torch import nn
from torch.distributed.tensor._api import DTensor  # noqa: PLC2701
from transformers.modeling_outputs import (
    CausalLMOutputWithPast,
)

CROSS_ENTROPY_IGNORE_INDEX = -100


def dtensor_safe_check_numel(tensor: torch.Tensor | DTensor) -> int:
    """Safely checks the number of elements in a tensor, handling DTensor types.

    Parameters
    ----------
    tensor : torch.Tensor | DTensor
        The tensor to check.

    Returns
    -------
    int
        The number of elements in the tensor.

    """
    if isinstance(tensor, DTensor):
        tensor = tensor._local_tensor  # noqa: SLF001
    return tensor.numel()


def get_targets(labels: torch.Tensor) -> torch.Tensor:
    """Shift labels for causal language modeling.

    Parameters
    ----------
    labels : torch.Tensor
        The labels tensor of shape (batch_size, sequence_length).

    Returns
    -------
    torch.Tensor
        The shifted labels tensor of shape (batch_size, sequence_length), where the last
        token in each sequence is set to `CROSS_ENTROPY_IGNORE_INDEX`.

    """
    targets = torch.roll(labels, shifts=-1)
    targets[:, -1] = CROSS_ENTROPY_IGNORE_INDEX
    return targets


def compute_loss_from_logits(
    outputs: CausalLMOutputWithPast,
    labels: torch.Tensor,
    loss_fn: nn.Module,
    *,
    shift_labels: bool,
) -> torch.Tensor:
    """Compute the loss from model outputs and labels.

    Parameters
    ----------
    outputs : CausalLMOutputWithPast
        The model outputs containing logits.
    labels : torch.Tensor
        The labels tensor of shape (batch_size, sequence_length).
    loss_fn : nn.Module
        The loss function to compute the loss from logits.
    shift_labels : bool
        Whether to shift the labels for causal language modeling.

    Returns
    -------
    torch.Tensor
        The computed loss value.

    """
    targets = get_targets(labels) if shift_labels else labels

    logits = outputs.logits
    assert logits is not None, "outputs.logits must not be None"

    losses = loss_fn(
        logits.view(-1, logits.size(-1)),
        targets.view(-1),
    )
    ignore_index = getattr(loss_fn, "ignore_index", CROSS_ENTROPY_IGNORE_INDEX)

    if torch.all(targets == ignore_index):
        loss = losses.sum()
    else:
        loss = losses.sum() / (targets != ignore_index).sum()

    return loss
