"""广义熵度量算法

特征列表：
- Renyi 熵 (Renyi Entropy)
- Tsallis 熵 (Tsallis Entropy)
- 碰撞熵 (Collision Entropy)
- 香农熵 (Shannon Entropy)
"""

import numpy as np
from typing import cast

from zyplib.utils.ensure import ensure_ndarray

from ..utils import adapt_1d_func


# =============================================================================
# Renyi 熵 (Renyi Entropy)
# =============================================================================


def _renyi_entropy_1d(x: np.ndarray, alpha: float, bins: int, normalize: bool) -> float:
    """计算一维信号的 Renyi 熵 (基于直方图)"""
    if len(x) < 10:
        return np.nan

    hist, _ = np.histogram(x, bins=bins, density=True)

    bin_width = (np.max(x) - np.min(x)) / bins
    if bin_width == 0:
        return 0.0

    probs = hist * bin_width
    probs = probs[probs > 0]

    if len(probs) == 0:
        return 0.0

    # Renyi 熵公式
    if np.abs(alpha - 1.0) < 1e-10:
        # α → 1 时退化为香农熵
        entropy = -np.sum(probs * np.log2(probs))
    else:
        entropy = (1.0 / (1.0 - alpha)) * np.log2(np.sum(probs**alpha))

    if normalize:
        max_entropy = np.log2(bins)
        if max_entropy > 0:
            entropy /= max_entropy

    return entropy


def renyi_entropy(
    signal: np.ndarray,
    alpha: float = 2.0,
    bins: int = 100,
    normalize: bool = False,
    axis: int = -1,
) -> np.ndarray:
    """计算信号的 Renyi 熵

    Renyi 熵是香农熵的推广，由参数 α 控制对概率分布不同部分的敏感度。

    公式: $H_\\alpha = \\frac{1}{1-\\alpha} \\log_2 \\sum_i p_i^\\alpha$

    复杂度: O(N + bins)

    Parameters
    ----------
    - `signal`: `np.ndarray`, 输入信号
    - `alpha`: `float`, Renyi 参数, by default 2.0
        - α = 0: Hartley 熵（只与支撑集大小有关）
        - α → 1: 香农熵
        - α = 2: 碰撞熵（与 Kolmogorov-Sinai 熵相关）
        - α → ∞: 最小熵（只与最大概率有关）
    - `bins`: `int`, 直方图的 bin 数, by default 100
    - `normalize`: `bool`, 是否归一化, by default False
    - `axis`: `int`, 计算轴, by default -1

    Returns
    ----------
    - `np.ndarray`, Renyi 熵值
    """
    signal = cast(np.ndarray, ensure_ndarray(signal))
    return adapt_1d_func(signal, _renyi_entropy_1d, axis, alpha, bins, normalize)


def collision_entropy(
    signal: np.ndarray, bins: int = 100, normalize: bool = False, axis: int = -1
) -> np.ndarray:
    """碰撞熵 (Collision Entropy)，即 α=2 的 Renyi 熵

    碰撞熵与 Kolmogorov-Sinai 熵相关，是 Renyi 熵最常用的特例之一。

    Parameters
    ----------
    - `signal`: `np.ndarray`, 输入信号
    - `bins`: `int`, 直方图 bin 数, by default 100
    - `normalize`: `bool`, 是否归一化, by default False
    - `axis`: `int`, 计算轴, by default -1

    Returns
    ----------
    - `np.ndarray`, 碰撞熵值
    """
    return renyi_entropy(signal, alpha=2.0, bins=bins, normalize=normalize, axis=axis)


# =============================================================================
# Tsallis 熵
# =============================================================================


def _tsallis_entropy_1d(x: np.ndarray, q: float, bins: int) -> float:
    """计算一维信号的 Tsallis 熵"""
    if len(x) < 10:
        return np.nan

    hist, _ = np.histogram(x, bins=bins, density=True)
    bin_width = (np.max(x) - np.min(x)) / bins
    if bin_width == 0:
        return 0.0

    probs = hist * bin_width
    probs = probs[probs > 0]

    if len(probs) == 0:
        return 0.0

    # Tsallis 熵公式
    if np.abs(q - 1.0) < 1e-10:
        # q → 1 时退化为香农熵
        return -np.sum(probs * np.log(probs))
    else:
        return (1.0 / (q - 1.0)) * (1.0 - np.sum(probs**q))


def tsallis_entropy(
    signal: np.ndarray, q: float = 1.5, bins: int = 100, axis: int = -1
) -> np.ndarray:
    """计算信号的 Tsallis 熵

    Tsallis 熵是非广延统计力学的基础，适用于具有长程相关性的系统。

    公式: $S_q = \\frac{1}{q-1} \\left(1 - \\sum_i p_i^q \\right)$

    复杂度: O(N + bins)

    Parameters
    ----------
    - `signal`: `np.ndarray`, 输入信号
    - `q`: `float`, 非广延性参数, by default 1.5
        - q < 1: 增强稀有事件的权重
        - q = 1: 退化为香农熵
        - q > 1: 增强常见事件的权重
    - `bins`: `int`, 直方图 bin 数, by default 100
    - `axis`: `int`, 计算轴, by default -1

    Returns
    ----------
    - `np.ndarray`, Tsallis 熵值
    """
    signal = cast(np.ndarray, ensure_ndarray(signal))
    return adapt_1d_func(signal, _tsallis_entropy_1d, axis, q, bins)


# =============================================================================
# 离散香农熵 (Shannon Entropy)
# =============================================================================


def _shannon_entropy_discrete_1d(x: np.ndarray, bins: int) -> float:
    """计算一维信号的离散香农熵"""
    if len(x) < 2:
        return np.nan

    hist, _ = np.histogram(x, bins=bins)
    probs = hist / np.sum(hist)
    probs = probs[probs > 0]

    if len(probs) == 0:
        return 0.0

    return -np.sum(probs * np.log2(probs))


def shannon_entropy(signal: np.ndarray, bins: int = 100, axis: int = -1) -> np.ndarray:
    """计算信号的离散香农熵

    通过直方图将连续信号离散化后计算香农熵。

    复杂度: O(N + bins)

    Parameters
    ----------
    - `signal`: `np.ndarray`, 输入信号
    - `bins`: `int`, 直方图 bin 数, by default 100
    - `axis`: `int`, 计算轴, by default -1

    Returns
    ----------
    - `np.ndarray`, 香农熵值（以 bit 为单位）
    """
    signal = cast(np.ndarray, ensure_ndarray(signal))
    return adapt_1d_func(signal, _shannon_entropy_discrete_1d, axis, bins)
