"""Utilities for computing log mean and entropy regularization.

Inspired from: https://github.com/RobertCsordas/moeut

"""

import math

import torch


def log_mean(x: torch.Tensor, dim: int = 0) -> torch.Tensor:
    """Compute the log mean of a tensor along a specified dimension.

    Parameters
    ----------
    x : torch.Tensor
        The input tensor.
    dim : int, optional
        The dimension along which to compute the log mean. Default is 0.

    Returns
    -------
    torch.Tensor
        The log mean of the input tensor along the specified dimension.

    """
    return x.logsumexp(dim) - math.log(x.shape[dim])


def entropy_l(input_tensor: torch.Tensor) -> torch.Tensor:
    """Compute the entropy of a tensor using the log softmax function.

    Parameters
    ----------
    input_tensor : torch.Tensor
        The input tensor for which to compute the entropy.

    Returns
    -------
    torch.Tensor
        The entropy of the input tensor.

    """
    return -(input_tensor * input_tensor.exp()).sum(-1)


def entropy_reg(sel: torch.Tensor, dim: int) -> torch.Tensor:
    """Compute the entropy regularization loss for a selection tensor.

    Parameters
    ----------
    sel : torch.Tensor
        The selection tensor, typically the output of a softmax operation.
    dim : int
        The dimension along which to compute the entropy.

    Returns
    -------
    torch.Tensor
        The entropy regularization loss.

    """
    sel = torch.nn.functional.log_softmax(sel, dim=-1)
    sel = log_mean(sel, dim)
    return -entropy_l(sel).mean()
