import flarejax as fj
import jax
from jaxtyping import Array, Float

import jax.nn as jnn

__all__ = [
    "ELU",
    "GLU",
    "CeLU",
    "HardSigmoid",
    "HardSiLU",
    "HardTanh",
    "LeakyReLU",
    "LogSigmoid",
    "LogSoftmax",
    "LogSumExp",
    "OneHot",
    "ReLU",
    "ReLU6",
    "SeLU",
    "Sigmoid",
    "SiLU",
    "Softmax",
    "SoftPlus",
    "SoftSign",
    "SparsePlus",
    "SquarePlus",
    "Standardize",
]


class CeLU(fj.Module):
    __module_name = "flarenet.CeLU"

    @fj.typecheck
    def __call__(self, x: Float[Array, "*s"]) -> Float[Array, "*s"]:
        return jnn.relu(x)


class ELU(fj.Module):
    __module_name = "flarenet.ELU"

    @fj.typecheck
    def __call__(self, x: Float[Array, "*s"]) -> Float[Array, "*s"]:
        return jnn.elu(x)


class GELU(fj.Module):
    __module_name = "flarenet.GELU"

    @fj.typecheck
    def __call__(self, x: Float[Array, "*s"]) -> Float[Array, "*s"]:
        return jnn.gelu(x)


class GLU(fj.Module):
    __module_name = "flarenet.GLU"

    axis: int = -1

    @fj.typecheck
    def __call__(self, x: Float[Array, "..."]) -> Float[Array, "..."]:
        return jnn.glu(x, self.axis)


class HardSigmoid(fj.Module):
    __module_name = "flarenet.HardSigmoid"

    @fj.typecheck
    def __call__(self, x: Float[Array, "*s"]) -> Float[Array, "*s"]:
        return jnn.hard_sigmoid(x)


class HardSiLU(fj.Module):
    __module_name = "flarenet.HardSiLU"

    @fj.typecheck
    def __call__(self, x: Float[Array, "*s"]) -> Float[Array, "*s"]:
        return jnn.hard_silu(x)


class HardTanh(fj.Module):
    __module_name = "flarenet.HardTanh"

    @fj.typecheck
    def __call__(self, x: Float[Array, "*s"]) -> Float[Array, "*s"]:
        return jnn.hard_tanh(x)


class LeakyReLU(fj.Module):
    __module_name = "flarenet.LeakyReLU"

    negative_slope: float = 1e-2

    @fj.typecheck
    def __call__(self, x: Float[Array, "*s"]) -> Float[Array, "*s"]:
        return jnn.leaky_relu(x, self.negative_slope)


class LogSigmoid(fj.Module):
    __module_name = "flarenet.LogSigmoid"

    @fj.typecheck
    def __call__(self, x: Float[Array, "*s"]) -> Float[Array, "*s"]:
        return jnn.log_sigmoid(x)


class LogSoftmax(fj.Module):
    __module_name = "flarenet.LogSoftmax"

    @fj.typecheck
    def __call__(self, x: Float[Array, "*s"]) -> Float[Array, "*s"]:
        return jnn.log_softmax(x, axis=-1)


class LogSumExp(fj.Module):
    __module_name = "flarenet.LogSumExp"

    @fj.typecheck
    def __call__(self, x: Float[Array, "..."]) -> Float[Array, "..."]:
        return jnn.logsumexp(x, axis=-1)


class Standardize(fj.Module):
    __module_name = "flarenet.Standardize"

    axis: int = -1

    @fj.typecheck
    def __call__(self, x: Float[Array, "*s"]) -> Float[Array, "*s"]:
        return jnn.standardize(x, axis=self.axis)


class OneHot(fj.Module):
    __module_name = "flarenet.OneHot"

    num_classes: int
    axis: int = -1

    @fj.typecheck
    def __call__(
        self,
        x: Float[Array, "*b"],
    ) -> Float[Array, "*b {self.dim}"]:
        return jnn.one_hot(x, self.num_classes, axis=self.axis)


class ReLU(fj.Module):
    __module_name = "flarenet.ReLU"

    @fj.typecheck
    def __call__(self, x: Float[Array, "*s"]) -> Float[Array, "*s"]:
        return jnn.relu(x)


class ReLU6(fj.Module):
    __module_name = "flarenet.ReLU6"

    @fj.typecheck
    def __call__(self, x: Float[Array, "*s"]) -> Float[Array, "*s"]:
        return jnn.relu6(x)


class SeLU(fj.Module):
    __module_name = "flarenet.SeLU"

    @fj.typecheck
    def __call__(self, x: Float[Array, "*s"]) -> Float[Array, "*s"]:
        return jnn.selu(x)


class Sigmoid(fj.Module):
    __module_name = "flarenet.Sigmoid"

    @fj.typecheck
    def __call__(self, x: Float[Array, "*s"]) -> Float[Array, "*s"]:
        return jnn.sigmoid(x)


class SoftSign(fj.Module):
    __module_name = "flarenet.SoftSign"

    @fj.typecheck
    def __call__(self, x: Float[Array, "*s"]) -> Float[Array, "*s"]:
        return jnn.soft_sign(x)


class Softmax(fj.Module):
    __module_name = "flarenet.Softmax"

    axis: int = -1

    @fj.typecheck
    def __call__(self, x: Float[Array, "*s"]) -> Float[Array, "*s"]:
        return jnn.softmax(x, axis=self.axis)


class SoftPlus(fj.Module):
    __module_name = "flarenet.SoftPlus"

    @fj.typecheck
    def __call__(self, x: Float[Array, "*s"]) -> Float[Array, "*s"]:
        return jnn.softplus(x)


class SparsePlus(fj.Module):
    __module_name = "flarenet.SparsePlus"

    @fj.typecheck
    def __call__(self, x: Float[Array, "*s"]) -> Float[Array, "*s"]:
        return jnn.sparse_plus(x)


class SiLU(fj.Module):
    __module_name = "flarenet.SiLU"

    @fj.typecheck
    def __call__(self, x: Float[Array, "*s"]) -> Float[Array, "*s"]:
        return jnn.silu(x)


class SquarePlus(fj.Module):
    __module_name = "flarenet.SquarePlus"

    @fj.typecheck
    def __call__(self, x: Float[Array, "*s"]) -> Float[Array, "*s"]:
        return jnn.squareplus(x)
