from typing import Any, Callable

import numpy as np
from typing import cast

from zyplib.utils.ensure import ensure_npy


def segment_signal(
    signal: np.ndarray,
    fs: int,
    seg_duration: float,
    overlap: float = 0,
    begin: float = 0,
    end: float | None = None,
):
    """将信号分割成多个片段

    在信号最后一个维度(时间维度)上分割，如果最后一个段的长度小于 seg_duration，则丢弃最后一个段。

    Parameters
    ----------
    - `signal` : `np.ndarray`
        - 待分段的信号
        - 可以是1-D或2-D数组，最后一个维度为时间维度
    - `fs` : `int`
        - 输入信号的采样频率 (Hz)
    - `seg_duration` : `float`
        - 每个段的持续时间 (秒)
    - `overlap` : `float`, optional
        - 相邻段之间的重叠时间 (秒), by default 0
    - `begin` : `float`, optional
        - 信号的起始时间 (秒), by default 0
    - `end` : `float`, optional
        - 信号的结束时间 (秒), by default None
        - 如果为 None，则使用信号的最后一个时间点作为结束时间

    Returns
    ----------
    - `np.ndarray`
        - 分段后的信号，形状为 [N, ..., L]
        - N 为段数，L 为每段的长度

    Raises
    ----------
    - `ValueError`
        - 当 seg_duration 不为正数时
        - 当 overlap 为负数时
        - 当 overlap 大于等于 seg_duration 时
        - 当 end 小于等于 begin 时
        - 当信号长度太短无法分段时

    Notes
    ----------
    使用 numpy 的 as_strided 函数实现高效的信号分段，返回的是原始数据的视图。
    为避免修改原始数据，最后会创建一个副本返回。
    """
    if seg_duration <= 0:
        raise ValueError('seg_duration must be positive.')
    if overlap < 0:
        raise ValueError('overlap must be non-negative.')
    if overlap >= seg_duration:
        raise ValueError('overlap must be less than seg_duration.')

    seg_len = int(seg_duration * fs)
    overlap_len = int(overlap * fs)
    step = seg_len - overlap_len
    if step <= 0:
        raise ValueError('overlap must be less than seg_duration.')

    if end is None:
        end = signal.shape[-1] / fs
    if end <= begin:
        raise ValueError('end must be greater than begin.')

    signal = signal[..., int(begin * fs) : int(end * fs)]

    # 计算分段后的信号长度
    length = signal.shape[-1]
    n_seg = (length - overlap_len) // step
    if n_seg <= 0:
        raise ValueError('signal is too short.')

    signal = cast(np.ndarray, ensure_npy(signal))

    # 分段
    segments = np.lib.stride_tricks.as_strided(
        signal,
        shape=(n_seg, *signal.shape[:-1], seg_len),
        strides=(step * signal.strides[-1], *signal.strides),
        writeable=False,
    )
    # 从 view 中拷贝数据, 避免修改原始数据
    segments = segments.copy()
    return segments


def _fn_identity(x):
    return x


def sliding_window(
    signal: np.ndarray,
    n_window_points: int,
    n_step_points: int | None = None,  # 默认等于 n_window_points
    idx_begin_point: int = 0,
    idx_end_point: int | None = None,
    axis: int = -1,
    fn_process: Callable[[np.ndarray], Any] = _fn_identity,
):
    """滑动窗口算法

    Parameters
    ----------
    - `signal` : `np.ndarray`
        - 输入信号
    - `n_window_points` : `int`
        - 窗口大小（采样点数）
    - `n_step_points` : `int`, optional
        - 步长（采样点数），默认等于 n_window_points
    - `idx_begin` : `int`, optional
        - 开始索引，默认为 0
    - `idx_end` : `int`, optional
        - 结束索引，默认为 None（使用信号末尾）
    - `axis` : `int`, optional
        - 操作的轴，默认为 -1（最后一个轴）
    - `fn_process` : `Callable`, optional
        - 对每个窗口进行处理的函数，默认为恒等函数

    Returns
    -------
    - `list`
        - 处理后的窗口结果列表
    """
    if n_step_points is None:
        n_step_points = n_window_points

    if axis < 0:
        axis = signal.ndim + axis

    # 获取指定轴的长度
    signal_length = signal.shape[axis]

    # 处理结束索引
    if idx_end_point is None:
        idx_end_point = signal_length

    # 确保索引范围有效
    idx_begin_point = max(0, idx_begin_point)
    idx_end_point = min(signal_length, idx_end_point)

    results = []

    current_idx = idx_begin_point
    while current_idx + n_window_points <= idx_end_point:
        # 创建切片对象
        slices = [slice(None)] * signal.ndim
        slices[axis] = slice(current_idx, current_idx + n_window_points)
        window_data = signal[tuple(slices)]

        # 处理窗口数据
        processed_result = fn_process(window_data)
        results.append(processed_result)

        # 移动到下一个窗口
        current_idx += n_step_points

    return results
