import logging

import triton
import triton.language as tl

from flag_gems.utils import tl_extra_shim

from ..utils.pointwise_dynamic import pointwise_dynamic

logger = logging.getLogger("flag_gems").getChild(__name__.lstrip("."))
div_rn = tl_extra_shim.div_rn


@pointwise_dynamic(is_tensor=[True, False], promotion_methods=[(0, "DEFAULT")])
@triton.jit
def silu_forward(x, inplace):
    x_fp32 = x.to(tl.float32)
    y = 1.0 / (1.0 + tl.exp(-x_fp32)) * x_fp32
    return y


@pointwise_dynamic(promotion_methods=[(0, "DEFAULT")])
@triton.jit
def silu_backward_kernel(x, dy):
    dy_fp32 = dy.to(tl.float32)
    x_fp32 = x.to(tl.float32)
    sigma = 1.0 / (1.0 + tl.exp(-x_fp32))
    dx = dy_fp32 * sigma * (1.0 + x_fp32 * (1.0 - sigma))
    return dx


def silu(self):
    logger.debug("GEMS_CAMBRICON SILU FORWARD")
    output = silu_forward(self, False)
    return output


def silu_backward(grad_output, self):
    logger.debug("GEMS_CAMBRICON SILU BACKWARD")
    grad_input = silu_backward_kernel(self, grad_output)
    return grad_input


def silu_(A):
    logger.debug("GEMS_CAMBRICON SILU_ FORWARD")
    out = silu_forward(A, True, out0=A)
    return out
