"""高级时频分析特征提取算法。(简化数据结构版)

提供各种高级时频分析方法，用于捕捉脑电信号的时频特性:
- 经验模态分解 (EMD)
- 集合经验模态分解 (EEMD)
- 希尔伯特变换特征
- 希尔伯特-黄变换

本模块依赖 PyEMD 库进行经验模态分解，请确保已安装该库：
pip install EMD-signal
"""

import typing

from typing import Any, Dict, List, Optional, Tuple, cast

import numpy as np

try:
    from PyEMD import EEMD, EMD
except ImportError:
    raise ImportError(
        "PyEMD package is required. Install it with 'pip install EMD-signal'"
    )
from scipy import signal as sig

from zyplib.utils.ensure import ensure_ndarray


def _decompose_1d(
    x_1d: np.ndarray,
    n_imfs_requested: int,
    is_eemd: bool,
    emd_lib_common_kwargs: typing.Mapping[str, Any],
    eemd_specific_kwargs: Optional[typing.Mapping[str, Any]] = None,
) -> np.ndarray:
    """
    (内部辅助函数) 使用EMD或EEMD分解单个1D信号。
    将IMFs填充到 n_imfs_requested 数量，并与残差堆叠。
    返回: shape 为 (n_imfs_requested + 1, n_times_in_x_1d) 的数组
    """
    if x_1d.ndim != 1:
        raise ValueError(f'_decompose_1d expects a 1D array, got shape {x_1d.shape}')

    time_pts = x_1d.shape[0]

    if time_pts == 0:
        imfs_padded = np.zeros((n_imfs_requested, 0), dtype=x_1d.dtype)
        residue_1d_arr = np.zeros((1, 0), dtype=x_1d.dtype)
        return np.vstack((imfs_padded, residue_1d_arr))

    current_eemd_constructor_kwargs = (
        typing.cast(Any, eemd_specific_kwargs).copy()
        if eemd_specific_kwargs is not None
        else {}
    )
    if is_eemd:
        current_eemd_constructor_kwargs.update(emd_lib_common_kwargs)
        noise_seed_val = current_eemd_constructor_kwargs.pop('noise_seed_val', None)
        decomposer = EEMD(**current_eemd_constructor_kwargs)
        if noise_seed_val is not None:
            decomposer.noise_seed(noise_seed_val)
    else:
        decomposer = EMD(**emd_lib_common_kwargs)

    t = np.arange(time_pts)
    all_components_from_pyemd = decomposer.emd(x_1d, t, max_imf=n_imfs_requested)

    if (
        not isinstance(all_components_from_pyemd, np.ndarray)
        or all_components_from_pyemd.size == 0
    ):
        imfs_1d_extracted = np.empty((0, time_pts), dtype=x_1d.dtype)
        residue_1d = np.copy(x_1d)
    elif all_components_from_pyemd.ndim == 1:
        imfs_1d_extracted = np.empty((0, time_pts), dtype=x_1d.dtype)
        residue_1d = all_components_from_pyemd
    elif all_components_from_pyemd.shape[0] == 0:
        imfs_1d_extracted = np.empty((0, time_pts), dtype=x_1d.dtype)
        residue_1d = np.copy(x_1d)
    else:
        imfs_1d_extracted = all_components_from_pyemd[:-1]
        residue_1d = all_components_from_pyemd[-1]

    if residue_1d.shape != (time_pts,):
        residue_1d = (
            residue_1d.reshape(time_pts) if residue_1d.size == time_pts else np.copy(x_1d)
        )

    num_found_imfs = imfs_1d_extracted.shape[0] if imfs_1d_extracted.ndim == 2 else 0

    if num_found_imfs < n_imfs_requested:
        padding_shape = (n_imfs_requested - num_found_imfs, time_pts)
        padding = np.zeros(padding_shape, dtype=x_1d.dtype)
        imfs_1d_padded = (
            np.vstack((imfs_1d_extracted, padding)) if num_found_imfs > 0 else padding
        )
    elif num_found_imfs > n_imfs_requested:
        imfs_1d_padded = imfs_1d_extracted[:n_imfs_requested]
    else:
        imfs_1d_padded = imfs_1d_extracted

    if n_imfs_requested == 0:
        imfs_1d_padded = np.empty((0, time_pts), dtype=x_1d.dtype)
    elif imfs_1d_padded.ndim == 1 and n_imfs_requested == 1:
        imfs_1d_padded = imfs_1d_padded.reshape(1, time_pts)
    elif imfs_1d_padded.shape != (n_imfs_requested, time_pts):
        imfs_1d_padded = np.zeros((n_imfs_requested, time_pts), dtype=x_1d.dtype)

    return np.vstack((imfs_1d_padded, residue_1d.reshape(1, time_pts)))


def emd(
    signal: np.ndarray, n_imfs: int = 5, axis: int = -1, **emd_kwargs: Any
) -> Tuple[np.ndarray, np.ndarray]:
    """使用经验模态分解提取信号的内在模态函数(IMFs)

    Parameters
    ----------
    - `signal` : `np.ndarray`
        - 输入信号，形状可以是:
            - (n_times,) - 一维时间序列
            - (n_channels, n_times) - 多通道时间序列
            - (n_samples, n_channels, n_times) - 多样本多通道时间序列
    - `n_imfs` : `int`, optional
        - 要提取的IMF数量，默认为5。如果实际分解得到的IMF少于此数，将用零填充以确保输出IMFs数组的第一个维度为此值（除非n_imfs为0）。
    - `axis` : `int`, optional
        - 进行EMD分解的轴，通常是时间轴，默认为最后一个维度。
    - `**emd_kwargs` : `Any`
        - 传递给 `PyEMD.EMD` 构造函数的额外参数。

    Returns
    -------
    - `Tuple[np.ndarray, np.ndarray]`
        - 一个元组 `(imfs, residue)`:
            - `imfs`: 内在模态函数数组。
                      形状为 `(n_imfs, d1, d2, ..., time_len_at_axis, ..., dN)` (如果n_imfs > 0),
                      或 `(0, d1, d2, ..., time_len_at_axis, ..., dN)` (如果n_imfs == 0)。
                      其中 `(d1, ..., dN)` 是原始信号 `signal` 的维度结构，`time_len_at_axis` 是沿 `axis` 轴的长度。
                      简单来说，`imfs` 的形状是 `(n_imfs, ...signal_shape...)`，其中 `signal_shape` 与输入 `signal` 的形状相同。
            - `residue`: 残差信号，其形状与输入 `signal` 相同。
    """
    signal_arr = typing.cast(np.ndarray, ensure_ndarray(signal))

    if n_imfs < 0:
        raise ValueError('n_imfs must be non-negative.')

    if signal_arr.size == 0:
        imf_final_shape = [n_imfs if n_imfs > 0 else 0] + list(signal_arr.shape)
        return (
            np.empty(tuple(imf_final_shape), dtype=signal_arr.dtype),
            np.empty_like(signal_arr),
        )

    normalized_axis = axis if axis >= 0 else signal_arr.ndim + axis

    if signal_arr.ndim == 1:
        components_stacked_1d = _decompose_1d(
            signal_arr,
            n_imfs,
            is_eemd=False,
            emd_lib_common_kwargs=emd_kwargs,
            eemd_specific_kwargs=None,
        )
        # components_stacked_1d shape is (n_imfs + 1, time_pts)
        # Reshape for consistency with multi-dim output:
        # imfs -> (n_imfs, time_pts), residue -> (time_pts,)
        if n_imfs == 0:
            final_imfs_shaped = np.empty((0, signal_arr.shape[0]), dtype=signal_arr.dtype)
            final_residue_shaped = components_stacked_1d[0, ...]
        else:
            final_imfs_shaped = components_stacked_1d[:-1, ...]
            final_residue_shaped = components_stacked_1d[-1, ...]
        return final_imfs_shaped, final_residue_shaped

    # For multi-dimensional signal:
    # components_multidim shape from apply_along_axis: (d0, ..., d_axis-1, n_imfs+1, time_len_of_original_axis, d_axis+1, ...)
    components_multidim = np.apply_along_axis(
        _decompose_1d,
        axis=normalized_axis,
        arr=signal_arr,
        n_imfs_requested=n_imfs,
        is_eemd=False,
        emd_lib_common_kwargs=emd_kwargs,
        eemd_specific_kwargs=None,
    )
    # Move components dim (n_imfs+1), currently at `normalized_axis`, to the front.
    components_leading_dim = np.moveaxis(components_multidim, normalized_axis, 0)
    # Shape of components_leading_dim: (n_imfs+1, d0, ..., time_len_of_original_axis_now_at_norm_axis+1_pos, ...)
    # The original signal dimensions (except the one replaced by components) are preserved.
    # The time dimension is now part of the trailing dimensions.
    # We need to reconstruct the original signal shape for residue and for each IMF slice.

    # Example: signal (S, C, T), axis = -1 (or 2). normalized_axis = 2.
    # _decompose_1d returns (M, T_len) where M = n_imfs+1.
    # components_multidim shape: (S, C, M, T_len).
    # components_leading_dim shape: (M, S, C, T_len). This is the desired structure.

    if n_imfs == 0:
        imf_dims = [0] + list(components_leading_dim.shape[1:])  # (0, S, C, T_len)
        final_imfs = np.empty(tuple(imf_dims), dtype=signal_arr.dtype)
        final_residue = components_leading_dim[0, ...]  # (S, C, T_len)
    else:
        final_imfs = components_leading_dim[:-1, ...]  # (n_imfs, S, C, T_len)
        final_residue = components_leading_dim[-1, ...]  # (S, C, T_len)

    return final_imfs, final_residue


def ensemble_emd(
    signal: np.ndarray,
    n_imfs: int = 5,
    noise_strength: float = 0.2,
    ensemble_size: int = 10,
    noise_seed: Optional[int] = None,
    axis: int = -1,
    **emd_lib_common_kwargs: Any,
) -> Tuple[np.ndarray, np.ndarray]:
    """使用集合经验模态分解提取信号的内在模态函数(IMFs)

    Parameters
    ----------
    - `signal` : `np.ndarray`
        - 输入信号，形状可以是:
            - (n_times,) - 一维时间序列
            - (n_channels, n_times) - 多通道时间序列
            - (n_samples, n_channels, n_times) - 多样本多通道时间序列
    - `n_imfs` : `int`, optional
        - 要提取的IMF数量，默认为5。如果实际分解得到的IMF少于此数，将用零填充。
    - `noise_strength` : `float`, optional
        - EEMD中添加的噪声强度（标准差与信号标准差之比），默认为0.2。
    - `ensemble_size` : `int`, optional
        - EEMD中的集合大小（试验次数），默认为10。
    - `noise_seed` : `Optional[int]`, optional
        - 用于EEMD噪声生成的随机种子，默认为None。
    - `axis` : `int`, optional
        - 进行EEMD分解的轴，默认为-1。
    - `**emd_lib_common_kwargs` : `Any`
        - 传递给底层 `PyEMD.EMD` (通过 `PyEMD.EEMD` 使用) 的额外参数。

    Returns
    -------
    - `Tuple[np.ndarray, np.ndarray]`
        - 一个元组 `(imfs, residue)`，结构与 `empirical_mode_decomposition` 返回值相同。
    """
    signal_arr = typing.cast(np.ndarray, ensure_ndarray(signal))

    if n_imfs < 0:
        raise ValueError('n_imfs must be non-negative.')

    eemd_specific_params = {
        'trials': ensemble_size,
        'noise_width': noise_strength,
        'noise_seed_val': noise_seed,
    }

    if signal_arr.size == 0:
        imf_final_shape = [n_imfs if n_imfs > 0 else 0] + list(signal_arr.shape)
        return (
            np.empty(tuple(imf_final_shape), dtype=signal_arr.dtype),
            np.empty_like(signal_arr),
        )

    normalized_axis = axis if axis >= 0 else signal_arr.ndim + axis

    if signal_arr.ndim == 1:
        components_stacked_1d = _decompose_1d(
            signal_arr,
            n_imfs,
            is_eemd=True,
            emd_lib_common_kwargs=emd_lib_common_kwargs,
            eemd_specific_kwargs=eemd_specific_params,
        )
        if n_imfs == 0:
            final_imfs_shaped = np.empty((0, signal_arr.shape[0]), dtype=signal_arr.dtype)
            final_residue_shaped = components_stacked_1d[0, ...]
        else:
            final_imfs_shaped = components_stacked_1d[:-1, ...]
            final_residue_shaped = components_stacked_1d[-1, ...]
        return final_imfs_shaped, final_residue_shaped

    components_multidim = np.apply_along_axis(
        _decompose_1d,
        axis=normalized_axis,
        arr=signal_arr,
        n_imfs_requested=n_imfs,
        is_eemd=True,
        emd_lib_common_kwargs=emd_lib_common_kwargs,
        eemd_specific_kwargs=eemd_specific_params,
    )
    components_leading_dim = np.moveaxis(components_multidim, normalized_axis, 0)

    if n_imfs == 0:
        imf_dims = [0] + list(components_leading_dim.shape[1:])
        final_imfs = np.empty(tuple(imf_dims), dtype=signal_arr.dtype)
        final_residue = components_leading_dim[0, ...]
    else:
        final_imfs = components_leading_dim[:-1, ...]
        final_residue = components_leading_dim[-1, ...]

    return final_imfs, final_residue


def hilbert_features(
    signal: np.ndarray,
    fs: float = 1.0,
    features: Optional[typing.Sequence[str]] = None,
    axis: int = -1,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """计算基于希尔伯特变换的特征

    Parameters
    ----------
    - `signal` : `np.ndarray`
        - 输入信号，形状可以是:
            - (n_times,) - 一维时间序列
            - (n_channels, n_times) - 多通道时间序列
            - (n_samples, n_channels, n_times) - 多样本多通道时间序列
    - `fs` : `float`, optional
        - 采样频率，默认为1.0 Hz。
    - `features` : `Optional[List[str]]`, optional
        - 需要计算的特征列表，可选 `['inst_amp', 'inst_phase', 'inst_freq']`。
        - 默认为None，计算所有三个特征。如果某个特征未在列表中指定，则返回的对应数组为全零。
    - `axis` : `int`, optional
        - 进行希尔伯特变换的轴（通常是时间轴），默认为-1。

    Returns
    -------
    - `Tuple[np.ndarray, np.ndarray, np.ndarray]`
        - 一个元组 `(inst_amp, inst_phase, inst_freq)`:
            - `inst_amp`: 瞬时幅度数组，形状与输入 `signal` 相同。
            - `inst_phase`: 瞬时相位数组（已解缠绕），形状与输入 `signal` 相同。
            - `inst_freq`: 瞬时频率数组，形状与输入 `signal` 相同。
        - 如果某个特征未在 `features` 参数中请求，则其对应的返回数组将为全零。
    """
    signal_arr = typing.cast(
        np.ndarray, ensure_ndarray(signal, dtype=np.dtype(np.float64))
    )

    if features is None:
        features_to_compute = {'inst_amp', 'inst_phase', 'inst_freq'}
    else:
        features_to_compute = set(features)
        valid_features = {'inst_amp', 'inst_phase', 'inst_freq'}
        if not features_to_compute.issubset(valid_features):
            invalid = features_to_compute - valid_features
            raise ValueError(
                f'Invalid feature(s) requested: {invalid}. Valid are {valid_features}'
            )

    if signal_arr.size == 0:
        empty_arr = np.empty_like(signal_arr, dtype=np.float64)
        return empty_arr, empty_arr, empty_arr

    # Normalize axis for array processing
    # For scalar input (ndim=0), processing it as a 1-element array with axis=0
    if signal_arr.ndim == 0:
        # Convert scalar to 1D array for hilbert and other ops
        # Results will be reshaped back to scalar if appropriate or handled as 1-element arrays
        signal_arr_proc = np.array(
            [signal_arr.item()]
        )  # Use .item() to get Python scalar if 0-dim
        normalized_axis_proc = 0
        is_scalar_input = True
    else:
        signal_arr_proc = signal_arr
        normalized_axis_proc = axis if axis >= 0 else signal_arr_proc.ndim + axis
        if not (0 <= normalized_axis_proc < signal_arr_proc.ndim):
            raise np.AxisError(
                f'Axis {axis} is out of bounds for array of dimension {signal_arr_proc.ndim}'
            )
        is_scalar_input = False

    analytic_signal = np.asarray(sig.hilbert(signal_arr_proc, axis=normalized_axis_proc))

    res_inst_amp = np.zeros_like(signal_arr_proc, dtype=np.float64)
    res_inst_phase = np.zeros_like(signal_arr_proc, dtype=np.float64)
    res_inst_freq = np.zeros_like(signal_arr_proc, dtype=np.float64)

    calculated_phase = None

    if 'inst_amp' in features_to_compute:
        res_inst_amp = np.abs(analytic_signal)

    if 'inst_phase' in features_to_compute or 'inst_freq' in features_to_compute:
        angles = np.angle(analytic_signal)
        calculated_phase = np.unwrap(angles, axis=normalized_axis_proc)
        if 'inst_phase' in features_to_compute:
            res_inst_phase = calculated_phase

    if 'inst_freq' in features_to_compute:
        if calculated_phase is None:
            angles = np.angle(analytic_signal)
            calculated_phase = np.unwrap(angles, axis=normalized_axis_proc)

        can_diff = signal_arr_proc.shape[normalized_axis_proc] >= 2
        if not can_diff:
            pass  # inst_freq remains zeros
        else:
            inst_freq_short = (
                np.diff(calculated_phase, axis=normalized_axis_proc) / (2.0 * np.pi) * fs
            )
            last_vals = np.take(inst_freq_short, indices=-1, axis=normalized_axis_proc)
            last_vals_expanded = np.expand_dims(last_vals, axis=normalized_axis_proc)
            res_inst_freq = np.concatenate(
                (inst_freq_short, last_vals_expanded), axis=normalized_axis_proc
            )

    if is_scalar_input:  # Convert back to 0-dim arrays if input was scalar
        res_inst_amp = res_inst_amp[0] if res_inst_amp.ndim > 0 else res_inst_amp
        res_inst_phase = res_inst_phase[0] if res_inst_phase.ndim > 0 else res_inst_phase
        res_inst_freq = res_inst_freq[0] if res_inst_freq.ndim > 0 else res_inst_freq
        # Ensure they are 0-dim arrays, not Python scalars
        res_inst_amp = np.array(res_inst_amp)
        res_inst_phase = np.array(res_inst_phase)
        res_inst_freq = np.array(res_inst_freq)

    return res_inst_amp, res_inst_phase, res_inst_freq


def hilbert_huang_transform(
    signal: np.ndarray,  # Changed from input_signal for consistency
    fs: float = 1.0,
    n_imfs: int = 5,
    axis: int = -1,
    emd_type: str = 'emd',
    **emd_processing_kwargs: Any,
) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray], Optional[np.ndarray]]:
    """计算希尔伯特-黄变换

    Parameters
    ----------
    - `signal` : `np.ndarray`
        - 输入信号，形状可以是:
            - (n_times,) - 一维时间序列
            - (n_channels, n_times) - 多通道时间序列
            - (n_samples, n_channels, n_times) - 多样本多通道时间序列
    - `fs` : `float`, optional
        - 采样频率，默认为1.0 Hz。
    - `n_imfs` : `int`, optional
        - EMD/EEMD 中要提取的IMF数量，默认为5。
          EMD/EEMD函数将确保返回的IMFs数组的第一个维度为此值（通过零填充）。
    - `axis` : `int`, optional
        - 进行EMD/EEMD和希尔伯特变换的轴（通常是时间轴），默认为-1。
    - `emd_type` : `str`, optional
        - 使用的EMD类型，可选 'emd' 或 'eemd'，默认为 'emd'。
    - `**emd_processing_kwargs` : `Any`
        - 传递给 `empirical_mode_decomposition` (如果 `emd_type` 是 'emd')
          或 `ensemble_emd` (如果 `emd_type` 是 'eemd') 的额外参数。
        - 例如，对于 'eemd'，可以包含 `noise_strength`, `ensemble_size`, `noise_seed`。
        - `PyEMD.EMD` 的参数 (如 `spline_kind`, `extrema_detection`) 也可以在此处提供。

    Returns
    -------
    - `Tuple[np.ndarray, np.ndarray, Optional[np.ndarray], Optional[np.ndarray]]`
        - 一个元组 `(all_imf_amplitudes, all_imf_frequencies, residue_amplitude, residue_frequency)`:
            - `all_imf_amplitudes`: 所有实际提取的非零IMFs的瞬时幅度堆叠而成的数组。
                                     其形状为 `(k, ...)`，其中 `k` 是实际找到的非零IMF的数量
                                     （`k <= n_imfs`），`...` 代表单个IMF的形状（与输入`signal`形状相同）。
                                     如果未找到非零IMF，则k为0，数组为空但保持维度。
            - `all_imf_frequencies`: 所有实际提取的非零IMFs的瞬时频率堆叠而成的数组。形状同上。
            - `residue_amplitude`: 残差的瞬时幅度数组，形状与输入`signal`相同。
                                    如果残差本身为None或全零，则此项为None。
            - `residue_frequency`: 残差的瞬时频率数组，形状与输入`signal`相同。
                                    如果残差本身为None或全零，则此项为None。
    """
    signal_arr = cast(np.ndarray, ensure_ndarray(signal))

    if emd_type.lower() == 'emd':
        imfs_array, residue_array = emd(signal_arr, n_imfs, axis, **emd_processing_kwargs)
    elif emd_type.lower() == 'eemd':
        current_eemd_kwargs = emd_processing_kwargs.copy()
        noise_strength = current_eemd_kwargs.pop('noise_strength', 0.2)
        ensemble_size = current_eemd_kwargs.pop('ensemble_size', 10)
        noise_seed = current_eemd_kwargs.pop('noise_seed', None)

        imfs_array, residue_array = ensemble_emd(
            signal_arr,
            n_imfs,
            noise_strength=noise_strength,
            ensemble_size=ensemble_size,
            noise_seed=noise_seed,
            axis=axis,
            **current_eemd_kwargs,
        )
    else:
        raise ValueError(f"Invalid emd_type: {emd_type}. Choose 'emd' or 'eemd'.")

    all_imf_amps_list: List[np.ndarray] = []
    all_imf_freqs_list: List[np.ndarray] = []

    num_potential_imfs = imfs_array.shape[0]  # This is n_imfs (due to padding)

    for i in range(num_potential_imfs):
        imf_component = imfs_array[i, ...]  # Shape: (signal_shape)

        if not np.any(imf_component):
            continue

        # Pass the original 'axis' because imf_component maintains original signal's other dimensions
        inst_amp, _, inst_freq = hilbert_features(
            imf_component, fs, ['inst_amp', 'inst_freq'], axis=axis
        )
        all_imf_amps_list.append(inst_amp)
        all_imf_freqs_list.append(inst_freq)

    # Determine the shape for an empty stack of IMF features
    # It should be (0, ...original_signal_shape...)
    empty_stack_shape = [0] + list(signal_arr.shape)

    if not all_imf_amps_list:
        stacked_imf_amps = np.empty(tuple(empty_stack_shape), dtype=np.float64)
        stacked_imf_freqs = np.empty(tuple(empty_stack_shape), dtype=np.float64)
    else:
        stacked_imf_amps = np.stack(all_imf_amps_list, axis=0)
        stacked_imf_freqs = np.stack(all_imf_freqs_list, axis=0)

    res_amp_arr: Optional[np.ndarray] = None
    res_freq_arr: Optional[np.ndarray] = None
    if residue_array is not None and np.any(residue_array):
        # Pass original 'axis' as residue_array has same ndim as signal_arr
        res_amp_arr, _, res_freq_arr = hilbert_features(
            residue_array, fs, ['inst_amp', 'inst_freq'], axis=axis
        )

    return stacked_imf_amps, stacked_imf_freqs, res_amp_arr, res_freq_arr
