from __future__ import annotations

import math

try:
    import numpy as np
except Exception:
    np = None  # type: ignore

try:
    import torch
    import torch.nn.functional as F
except Exception:
    torch = None  # type: ignore
    F = None  # type: ignore


# NumPy 版本
def np_relu(x):
    if np is None:
        raise RuntimeError("NumPy unavailable")
    return np.maximum(0, x)


def np_leaky_relu(x, negative_slope: float = 0.01):
    if np is None:
        raise RuntimeError("NumPy unavailable")
    return np.where(x >= 0, x, negative_slope * x)


def np_sigmoid(x):  # 数值稳定分段
    if np is None:
        raise RuntimeError("NumPy unavailable")
    out = np.empty_like(x, dtype=float)
    pos = x >= 0
    neg = ~pos
    out[pos] = 1.0 / (1.0 + np.exp(-x[pos]))
    ex = np.exp(x[neg])
    out[neg] = ex / (1.0 + ex)
    return out


def np_tanh(x):
    if np is None:
        raise RuntimeError("NumPy unavailable")
    return np.tanh(x)


def np_softmax(x, axis: int = -1):  # 减去 max 提升稳定性
    if np is None:
        raise RuntimeError("NumPy unavailable")
    x = np.asarray(x, dtype=float)
    m = np.max(x, axis=axis, keepdims=True)
    e = np.exp(x - m)
    return e / np.sum(e, axis=axis, keepdims=True)


# PyTorch 版本
def th_relu(x):
    if torch is None:
        raise RuntimeError("PyTorch unavailable")
    return F.relu(x)


def th_leaky_relu(x, negative_slope: float = 0.01):
    if torch is None:
        raise RuntimeError("PyTorch unavailable")
    return F.leaky_relu(x, negative_slope=negative_slope)


def th_sigmoid(x):
    if torch is None:
        raise RuntimeError("PyTorch unavailable")
    return torch.sigmoid(x)


def th_tanh(x):
    if torch is None:
        raise RuntimeError("PyTorch unavailable")
    return torch.tanh(x)


def th_softmax(x, dim: int = -1):
    if torch is None:
        raise RuntimeError("PyTorch unavailable")
    return F.softmax(x, dim=dim)


if __name__ == "__main__":
    print("[Activation] quick check")
    if np is not None:
        a = np.array([-2.0, 0.0, 0.5, 10.0])
        print("np_relu:", np_relu(a))
    if torch is not None:
        t = torch.tensor([-2.0, 0.0, 0.5, 10.0])
        print("th_relu:", th_relu(t))