from typing import Tuple

import numpy as np


def hjorth_activity(
    signal: np.ndarray, axis: int = -1, keepdims: bool = False, ddof: int = 0
) -> np.ndarray:
    """计算 Hjorth Activity 参数 (信号方差)

    Parameters
    ----------
    - `signal` : `np.ndarray`
        - 输入信号
    - `axis` : `int`, optional
        - 计算轴，默认为最后一个维度
    - `keepdims` : `bool`, optional
        - 是否保持输出维度与输入一致，默认为 False
    - `ddof` : `int`, optional
        - 自由度，默认为 0

    Returns
    -------
    - `np.ndarray`
        - Activity 参数
    """
    return np.var(signal, axis=axis, keepdims=keepdims, ddof=ddof)


def hjorth_mobility(
    signal: np.ndarray, axis: int = -1, keepdims: bool = False, ddof: int = 0
) -> np.ndarray:
    """计算 Hjorth Mobility 参数 (信号一阶导数方差与信号方差的比值的平方根)

    Parameters
    ----------
    - `signal` : `np.ndarray`
        - 输入信号
    - `axis` : `int`, optional
        - 计算轴，默认为最后一个维度
    - `keepdims` : `bool`, optional
        - 是否保持输出维度与输入一致，默认为 False
    - `ddof` : `int`, optional
        - 自由度，默认为 0

    Returns
    -------
    - `np.ndarray`
        - Mobility 参数
    """
    # 计算差分
    diff = np.diff(signal, axis=axis)

    # 为了保持维度一致，需要在差分的维度上填充一个值
    # 使用最后一个差分值进行填充
    pad_width = [(0, 0)] * signal.ndim
    pad_width[axis] = (0, 1)
    diff_padded = np.pad(diff, pad_width, mode='edge')

    # 添加小常数防止除零
    eps = np.finfo(signal.dtype).eps
    signal_std = np.std(signal, axis=axis, keepdims=keepdims, ddof=ddof)
    diff_std = np.std(diff_padded, axis=axis, keepdims=keepdims, ddof=ddof)

    return diff_std / (signal_std + eps)


def hjorth_complexity(
    signal: np.ndarray, axis: int = -1, keepdims: bool = False, ddof: int = 0
) -> np.ndarray:
    """计算 Hjorth Complexity 参数 (信号二阶导数的 mobility 与一阶导数的 mobility 的比值)

    Parameters
    ----------
    - `signal` : `np.ndarray`
        - 输入信号
    - `axis` : `int`, optional
        - 计算轴，默认为最后一个维度
    - `keepdims` : `bool`, optional
        - 是否保持输出维度与输入一致，默认为 False
    - `ddof` : `int`, optional
        - 自由度，默认为 0

    Returns
    -------
    - `np.ndarray`
        - Complexity 参数
    """
    # 计算一阶差分
    diff1 = np.diff(signal, axis=axis)
    pad_width = [(0, 0)] * signal.ndim
    pad_width[axis] = (0, 1)
    diff1_padded = np.pad(diff1, pad_width, mode='edge')

    # 计算二阶差分
    diff2 = np.diff(diff1_padded, axis=axis)
    diff2_padded = np.pad(diff2, pad_width, mode='edge')

    # 添加小常数防止除零
    eps = np.finfo(signal.dtype).eps
    mob1 = hjorth_mobility(signal, axis=axis, keepdims=keepdims, ddof=ddof)
    mob2 = hjorth_mobility(diff1_padded, axis=axis, keepdims=keepdims, ddof=ddof)

    return mob2 / (mob1 + eps)


def hjorth_parameters(
    signal: np.ndarray, axis: int = -1, keepdims: bool = False, ddof: int = 0
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """一次性计算所有Hjorth参数

    Parameters
    ----------
    - `signal` : `np.ndarray`
        - 输入信号
    - `axis` : `int`, optional
        - 计算轴，默认为最后一个维度
    - `keepdims` : `bool`, optional
        - 是否保持输出维度与输入一致，默认为 False
    - `ddof` : `int`, optional
        - 自由度，默认为 0

    Returns
    -------
    - `Tuple[np.ndarray, np.ndarray, np.ndarray]`
        - (Activity, Mobility, Complexity) 参数元组
    """
    activity = hjorth_activity(signal, axis, keepdims, ddof)
    mobility = hjorth_mobility(signal, axis, keepdims, ddof)
    complexity = hjorth_complexity(signal, axis, keepdims, ddof)

    return activity, mobility, complexity
