from typing import Callable

import numpy as np


def adapt_1d_func(
    signal: np.ndarray, func: Callable, axis: int = -1, *args
) -> np.ndarray:
    """将用于一维输入信号的的特征提取函数适配到多维信号张量

    Parameters
    ----------
    - `signal` : `np.ndarray`
        - 输入信号，形状可以是:
          - (n_times,) - 一维时间序列
          - (n_channels, n_times) - 多通道时间序列
          - (n_samples, n_channels, n_times) - 多样本多通道时间序列
    - `func` : `callable`
        - 要应用的函数，应接受一维数组作为第一个参数
    - `axis` : `int`
        - 计算轴, 默认 -1
    - `*args` :
        - 传递给func的其他参数

    Returns
    -------
    - `np.ndarray`
        - 函数应用的结果，形状取决于输入信号的维度:
          - 如果输入是1D，则输出是标量
          - 如果输入是2D，则输出是1D数组，形状为(n_channels,)
          - 如果输入是3D，则输出是2D数组，形状为(n_samples, n_channels)
    """
    if signal.ndim == 1:
        return func(signal, *args)
    elif signal.ndim == 2 or signal.ndim == 3:
        return np.apply_along_axis(lambda x: func(x, *args), axis, signal)
    else:
        raise ValueError(f'Unsupported signal dimension: {signal.ndim}')
