import torch
from typing import Tuple


def sigmoid_soft_top_k(
    logits: torch.Tensor,
    k: int,
    tau: float,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Calculate soft top k mask using a sigmoid function.

    Args:
        logits: Parameters used for weighting.
        k: Number of elements that should be selected.
        tau: Temperature that is used to vary steepness of sigmoid function.

    Returns:
        mask: Soft top k mask that can be used to select k elements.
        x0: Exactly k (normalized) logits are larger than x0. (Useful to gain insight in logits.)
    """
    if k == len(logits):
        # Not doing anything
        return torch.ones_like(logits), torch.zeros(
            1,
            device=logits.device,
        )

    std, mean = torch.std_mean(logits)
    logits = (logits - mean) / torch.maximum(std, torch.tensor(1e-3))

    # Find splitter value: k values are >=q
    q = 1 - k / len(logits)
    x0 = torch.quantile(logits, q).detach()

    # Center and stretch depending on temperature tau
    logits = (logits - x0) / tau

    return torch.sigmoid(logits), x0.detach()