from typing import Optional
import numpy as np

from cleanlab.internal.constants import EPSILON


def softmax(
    x: np.ndarray, temperature: float = 1.0, axis: Optional[int] = None, shift: bool = False
) -> np.ndarray:
    """Softmax function.

    Parameters
    ----------
    x : np.ndarray
        Input array.

    temperature : float
        Temperature of the softmax function.

    axis : Optional[int]
        Axis to apply the softmax function. If None, the softmax function is
        applied to all elements of the input array.

    shift : bool
        Whether to shift the input array before applying the softmax function.
        This is useful to avoid numerical issues when the input array contains
        large values, that could result in overflows when applying the exponential
        function.

    Returns
    -------
    np.ndarray
        Softmax function applied to the input array.
    """
    x = x / max(temperature, EPSILON)
    if shift:
        x = x - np.max(x, axis=axis, keepdims=True)
    exp_x = np.exp(x)
    return exp_x / np.sum(exp_x, axis=axis, keepdims=True)