"""
kiri/nn/activations.py
"""

import numpy as np
from kiri.backend.detect import BACKEND
from kiri.nn.layers import Module
from kiri.tensor import _CPUTensor

if BACKEND == "mlx":
    import mlx.core as mx
    import mlx.nn as mlx_nn


class ReLU(Module):
    def forward(self, x):
        if BACKEND == "mlx":
            return mx.maximum(x, 0)
        return x.relu()


class LeakyReLU(Module):
    def __init__(self, negative_slope=0.01):
        super().__init__()
        self.negative_slope = negative_slope

    def forward(self, x):
        if BACKEND == "mlx":
            return mx.where(x >= 0, x, self.negative_slope * x)
        out = _CPUTensor(
            np.where(x.data >= 0, x.data, self.negative_slope * x.data),
            (x,), "leaky_relu"
        )
        def _backward():
            x.grad += np.where(x.data >= 0, 1.0, self.negative_slope) * out.grad
        out._backward = _backward
        return out


class Sigmoid(Module):
    def forward(self, x):
        if BACKEND == "mlx":
            return mx.sigmoid(x)
        return x.sigmoid()


class Tanh(Module):
    def forward(self, x):
        if BACKEND == "mlx":
            return mx.tanh(x)
        return x.tanh()


class Softmax(Module):
    def __init__(self, dim=-1):
        super().__init__()
        self.dim = dim

    def forward(self, x):
        if BACKEND == "mlx":
            return mx.softmax(x, axis=self.dim)
        # numerically stable softmax
        data = x.data
        shifted = data - data.max(axis=self.dim, keepdims=True)
        exp_x = np.exp(shifted)
        s = exp_x / exp_x.sum(axis=self.dim, keepdims=True)
        out = _CPUTensor(s, (x,), "softmax")
        def _backward():
            dx = s * (out.grad - (out.grad * s).sum(axis=self.dim, keepdims=True))
            x.grad += dx
        out._backward = _backward
        return out


class GELU(Module):
    def forward(self, x):
        if BACKEND == "mlx":
            return mlx_nn.gelu(x)
        import math
        data = x.data
        val = 0.5 * data * (1 + np.tanh(np.sqrt(2/math.pi) * (data + 0.044715 * data**3)))
        out = _CPUTensor(val, (x,), "gelu")
        def _backward():
            c  = np.sqrt(2/math.pi)
            t  = np.tanh(c * (data + 0.044715 * data**3))
            dt = c * (1 + 3 * 0.044715 * data**2) * (1 - t**2)
            x.grad += out.grad * 0.5 * ((1 + t) + data * dt)
        out._backward = _backward
        return out
