"""
Activations functions for fundl.
"""
import jax.numpy as np


def sigmoid(x, version="tanh"):
    if version not in ["tanh", "exp"]:
        raise ValueError("version must be one of ['tanh' or 'exp']")

    sigmoids = {
        "tanh": lambda x: 0.5 * np.tanh(x) + 0.5,
        "exp": lambda x: 1 / (1 + np.exp(-x)),
    }
    return sigmoids[version](x)


def relu(x, alpha=0.1):
    return x * (x > 0) + alpha * x * (x <= 0)


def tanh(x):
    return np.tanh(x)


def leaky_relu(x, a):
    """Leaky ReLu"""
    pass


def identity(x):
    return x


def softmax(x, axis: int = -1):
    return np.exp(x) / np.sum(np.exp(x), axis=axis).reshape(-1, 1)
