"""常用函数。softmax / log_softmax / gelu 写成融合算子（手写反向），更快也更数值稳定。"""
import numpy as np

from .tensor import Tensor


def relu(x):
    return x.relu()


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


def sigmoid_np(a):
    """数值稳定：只对 -|a| 求 exp，永远不会溢出。"""
    e = np.exp(-np.abs(a))
    return np.where(a >= 0, 1 / (1 + e), e / (1 + e))


def sigmoid(x):
    out = sigmoid_np(x.data)
    return Tensor._op(out, (x,), lambda g: (g * out * (1 - out),))


def gelu(x):
    """GELU（tanh 近似，GPT-2 同款）。"""
    c = np.sqrt(2 / np.pi)
    a = x.data
    inner = c * (a + 0.044715 * a ** 3)
    t = np.tanh(inner)
    out = 0.5 * a * (1 + t)

    def backward(g):
        dinner = c * (1 + 3 * 0.044715 * a ** 2)
        return (g * (0.5 * (1 + t) + 0.5 * a * (1 - t ** 2) * dinner),)

    return Tensor._op(out.astype(a.dtype), (x,), backward)


def softmax(x, axis=-1):
    e = np.exp(x.data - x.data.max(axis=axis, keepdims=True))
    s = e / e.sum(axis=axis, keepdims=True)
    return Tensor._op(s, (x,), lambda g: (s * (g - (g * s).sum(axis=axis, keepdims=True)),))


def log_softmax(x, axis=-1):
    shifted = x.data - x.data.max(axis=axis, keepdims=True)
    out = shifted - np.log(np.exp(shifted).sum(axis=axis, keepdims=True))
    s = np.exp(out)
    return Tensor._op(out, (x,), lambda g: (g - s * g.sum(axis=axis, keepdims=True),))


def cross_entropy(logits, targets):
    """logits: (N, C)；targets: (N,) 整数类别。返回平均负对数似然。"""
    targets = targets.data if isinstance(targets, Tensor) else np.asarray(targets)
    logp = log_softmax(logits, axis=-1)
    return -logp[np.arange(len(targets)), targets].mean()


def mse_loss(pred, target):
    return ((pred - target) ** 2).mean()


def dropout(x, p, training=True):
    if not training or p == 0:
        return x
    mask = (np.random.rand(*x.shape) >= p).astype(x.dtype) / (1 - p)
    return x * Tensor(mask)