import typing
from typing import Callable, Literal, Optional

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.axes import Axes

from zyplib.annotation import Event
from zyplib.utils.ensure import ensure_2dims, ensure_ndarray
from zyplib.utils.kv import filter_kwargs
from zyplib.utils.print import warning
from zyplib.vis._plot_utils import _finalize_figure, _prepare_ax
from zyplib.vis.plot_common import plot_line

AUTO_SPECIFIED = object()  # 参数根据上下文由函数内部自动生成


def plot_signals(
    signals: np.ndarray,
    fs: int,
    labels: typing.Sequence[str] | None = None,
    interval: float | Literal['auto'] | None = 'auto',
    fn_value_fmt: Callable[[float], str] | None = lambda x: f'{x:.2f}',
    cmap: str | typing.Sequence[str] | Callable[[int], str] | None = 'black',
    t_start: float = 0,
    events: typing.Sequence[Event] | None = None,
    xlabel: str = 'Time (s)',
    title: str | None = None,
    show_ch_range=False,
    grid: bool = True,
    ax: Optional[Axes] = None,
    show=True,
    save_path=None,
    auto_close: bool | None = None,
    figsize: tuple[float, float] | None = None,
    **kwargs,
):
    """绘制多通道信号波形图

    Parameters
    ----------
    - `signals` : `np.ndarray`, [n_channels, n_samples]; 信号数据
    - `labels` : `list[str]`; 每个信号的标签
    - `fs` : `int`; 采样率
    - `new_fig` : `bool`, optional
        - 是否创建一个新的 figure。默认为 True
    - `show` : `bool`, optional
        - 是否显示绘图。默认为 True
    - `save_path` : `str`, optional
        - 保存绘图的路径。如果为 None，则不保存绘图。默认为 None
    - `interval` : `float | Literal['auto']`, optional
        - 不同通道在 Y 轴上的间隔, 默认 'auto' 表示自动计算
    - `fn_value_fmt` : `Callable[[float], str]`, optional
        - 用于格式化信号值的函数。默认为 `lambda x: f'{x:.2f}'`
    - `cmap` : `str | Callable[[int], str] | None`, optional
        - 用于绘制信号的 colormap。默认为 'black'
        - 也可以为一个 `int -> str` 的函数，int 为通道索引
    - `t_start` : `float`, optional
        - 信号的起始时间。默认为 0
    - `events` : `list[Event] | None`, optional
        - 标注事件。默认为 None
    - `auto_close` : `bool | None`, optional
        - 是否自动关闭图形以释放内存。默认为 None
        - None: 当 show=False 且 save_path 不为 None 时自动关闭
        - True: 总是自动关闭图形
        - False: 永不自动关闭图形
    """
    signals = ensure_ndarray(signals)
    signals = ensure_2dims(signals)

    n_channels, n_samples = signals.shape
    t_start = max(t_start, 0)
    t_end = t_start + n_samples / fs
    time = np.linspace(t_start, t_end, n_samples)

    labels = labels or [f'Ch{i}' for i in range(n_channels)]

    # 计算间隔
    if interval is None or interval == 'auto':
        interval = float(1.5 * np.max(np.abs(signals)))

    # detrends
    signals = np.apply_along_axis(lambda x: x - np.mean(x), 1, signals)
    y_centers = np.arange(n_channels) * interval

    # if new_fig:
    #     # fig = plt.figure(figsize=(15, n_channels * 2))
    #     fig = plt.figure()
    # else:
    #     fig = plt.gcf()
    fig, ax = _prepare_ax(ax, figsize=figsize)

    def get_c(i):
        if isinstance(cmap, list):
            return cmap[i]
        return cmap(i) if callable(cmap) else cmap

    line_plot_kwargs = filter_kwargs(
        kwargs, prohibited=['figsize', 'nrows', 'ncols', 'sharex', 'sharey']
    )
    for i in range(n_channels):
        c = get_c(i)
        ax.plot(
            time, signals[i] + y_centers[i], label=labels[i], color=c, **line_plot_kwargs
        )

    ax.set_xlim(t_start, t_end)
    if xlabel:
        ax.set_xlabel(xlabel)
    if title:
        ax.set_title(title)
    if grid:
        ax.grid(True)

    # 左侧标注通道名称
    ax.set_yticks(y_centers, labels)

    # Y 轴右侧标注 signal 的 min/max
    if show_ch_range:
        ax2 = ax.twinx()
        ax2.set_ylim(ax.get_ylim())
        ax2.set_yticks(y_centers)
        value_formatter = fn_value_fmt or (lambda value: f'{value:.2f}')
        min_max_labels = [
            f'{value_formatter(float(signals[i].min()))}~{value_formatter(float(signals[i].max()))}'
            for i in range(n_channels)
        ]
        ax2.set_yticklabels(min_max_labels)
        # ax2.set_ylabel('Min / Max')

    # 绘制标注
    if events:
        y_lim = ax.get_ylim()
        y_mid = (y_lim[0] + y_lim[1]) / 2
        for event in events:
            onset = event.onset
            duration = event.duration
            label = event.name
            if onset < t_start or onset > t_end:
                warning(f'标注 {label} 的 onset={onset} 超出了图表范围，已跳过')
                continue
            event_end = onset + duration if duration else onset
            if event_end > t_end:
                warning(
                    f'标注 {label} 的 duration={duration} 超出了图表范围，多余的部分将不会显示'
                )
                duration = t_end - onset

            if duration is None:
                ax.axvline(onset, color='red', linestyle='--', label=label)
            else:
                ax.axvspan(onset, onset + duration, color='red', alpha=0.2, label=label)
            if label:
                ax.text(
                    onset,
                    y_mid,  # 显示在 y 轴的中间位置
                    str(label),
                    color='red',
                    # fontsize=10,
                    verticalalignment='center',
                    horizontalalignment='right',
                )

    # plt.tight_layout()

    _finalize_figure(fig, show, save_path, auto_close)

    return fig


def plot_pwelch(
    signals: np.ndarray,
    fs: float,
    mode: Literal['individual', 'average'] = 'average',
    nperseg: int | None = None,
    noverlap: int | None = None,
    nfft: int | None = None,
    max_freq: float | None = None,
    window: str = 'hamming',
    scaling: Literal['density', 'spectrum'] = 'density',
    return_db: bool = True,
    labels: typing.Sequence[str] | None = None,
    error_style: Literal['band', 'bar'] = 'band',
    xlabel: str = 'Frequency (Hz)',
    ylabel: str | object | None = AUTO_SPECIFIED,
    title: str | None = None,
    ax: Optional[Axes] = None,
    show: bool = True,
    save_path: str | None = None,
    auto_close: bool | None = None,
    figsize: tuple[float, float] | None = None,
    **welch_kwargs,
):
    """绘制功率谱密度图 (Welch方法)

    使用 Welch 方法计算信号的功率谱密度并绘制。支持单通道和多通道信号。
    对于多通道信号，可以显示每个通道的PSD，或显示均值和方差区域。

    Parameters
    ----------
    - `signals` : `np.ndarray`
        - 输入信号。一维数组为单通道信号，二维数组为多通道信号 [n_channels, n_samples]
    - `fs` : `float`
        - 采样频率，单位为 Hz
    - `mode` : `Literal['individual', 'average']`, optional
        - 多通道信号的显示模式，by default 'average'
        - 'individual': 显示每个通道的PSD曲线
        - 'average': 显示通道间均值和误差带（推荐用于多通道）
        - 对于单通道信号，此参数无效
    - `nperseg` : `int | None`, optional
        - 每段的长度，默认 `2 * fs`。频率分辨率为 `fs / nperseg`, by default None
    - `noverlap` : `int | None`, optional
        - 段之间的重叠点数。如果为 None，默认为 `nperseg // 2`, by default None
    - `nfft` : `int | None`, optional
        - FFT 长度。如果为 None，默认为大于等于 `nperseg` 的最小 2 的幂次, by default None
    - `max_freq` : `float | None`, optional
        - 显示的最大频率。如果为 None，显示到奈奎斯特频率, by default None
    - `window` : `str`, optional
        - 窗函数类型, by default 'hamming'
    - `scaling` : `Literal['density', 'spectrum']`, optional
        - 'density': 返回功率谱密度 (功率/Hz)；'spectrum': 返回功率谱 (功率), by default 'density'
    - `return_db` : `bool`, optional
        - 是否将功率谱转换为分贝 (dB) 形式, by default True
    - `labels` : `list[str] | None`, optional
        - 通道标签。仅在 `mode='individual'` 时使用, by default None
    - `error_style` : `Literal['band', 'bar']`, optional
        - 误差显示方式。'band'显示误差带，'bar'显示误差棒。仅在 `mode='average'` 时使用, by default 'band'
    - `xlabel` : `str`, optional
        - X轴标签, by default 'Frequency (Hz)'
    - `ylabel` : `str | None`, optional
        - Y轴标签。如果为 None，根据 return_db 自动设置, by default None
    - `title` : `str | None`, optional
        - 图表标题。如果为 None，自动生成, by default None
    - `ax` : `Optional[plt.Axes]`, optional
        - 绘图轴对象。如果为 None 则创建新图表, by default None
    - `show` : `bool`, optional
        - 是否显示图表, by default True
    - `save_path` : `str | None`, optional
        - 保存路径。如果提供则保存图表, by default None
    - `auto_close` : `bool | None`, optional
        - 是否自动关闭图表释放内存, by default None

    Returns
    ----------
    - `fig` : `matplotlib.figure.Figure`
        - 图表对象

    Examples
    --------
    **单通道信号**：
    ```python
    import numpy as np
    from zyplib.vis.plot_eeg import plot_pwelch

    # 生成测试信号
    fs = 1000
    t = np.arange(0, 10, 1/fs)
    signal = np.sin(2*np.pi*10*t) + 0.5*np.sin(2*np.pi*20*t) + np.random.randn(len(t))*0.1

    # 绘制功率谱
    fig = plot_pwelch(signal, fs, max_freq=100)
    ```

    **多通道信号（显示均值和方差）**：
    ```python
    # 生成多通道信号
    n_channels = 10
    signals = np.random.randn(n_channels, len(t)) * 0.1
    for i in range(n_channels):
        signals[i] += np.sin(2*np.pi*10*t) + 0.5*np.sin(2*np.pi*20*t)

    # 显示均值和误差带
    fig = plot_pwelch(signals, fs, mode='average', max_freq=100)
    ```

    **多通道信号（显示各通道）**：
    ```python
    # 显示每个通道的PSD
    fig = plot_pwelch(signals, fs, mode='individual',
                      labels=[f'Ch{i}' for i in range(n_channels)],
                      max_freq=100)
    ```
    """
    # 导入 PSD 计算函数
    from zyplib.signal.psd import welch

    # 确保输入格式正确
    signals = ensure_ndarray(signals)

    # 设置默认ylabel
    if ylabel is AUTO_SPECIFIED:
        ylabel = 'Power (dB/Hz)' if return_db else 'Power'

    # 判断信号维度
    if signals.ndim == 1:
        # 单通道信号
        signals = signals.reshape(1, -1)  # 转换为 [1, n_samples] 格式
        is_multichannel = False
    elif signals.ndim == 2:
        # 多通道信号 [n_channels, n_samples]
        is_multichannel = True
    else:
        raise ValueError('signals 必须是一维或二维数组')

    n_channels, n_samples = signals.shape

    # 计算每个通道的PSD
    psds = []
    for i in range(n_channels):
        f, psd = welch(
            signals[i],
            fs=fs,
            nperseg=nperseg,
            noverlap=noverlap,
            nfft=nfft,
            max_freq=max_freq,
            window=window,
            scaling=scaling,
            return_db=return_db,
            **welch_kwargs,
        )
        psds.append(psd)

    # 将PSD列表转换为数组 [n_channels, n_freqs]
    psds = np.array(psds)

    # 准备绘图数据
    resolved_ylabel = (
        None if ylabel is AUTO_SPECIFIED else typing.cast(str | None, ylabel)
    )

    if is_multichannel and mode == 'average':
        # 多通道：显示均值和误差带
        psd_mean = np.mean(psds, axis=0)
        psd_std = np.std(psds, axis=0)

        # 使用 plot_line 绘制均值和误差带
        fig, ax = plot_line(
            x=f,
            y=psd_mean.reshape(1, -1),  # plot_line 期望 2D 数组
            yerr=psd_std.reshape(1, -1),
            error_style=error_style,
            labels=['Mean PSD'],
            xlabel=xlabel,
            ylabel=resolved_ylabel,
            ax=ax,
            show=False,  # 先不显示，后面统一处理
            auto_close=False,
            figsize=figsize,
        )
    else:
        # 单通道或多通道显示各通道
        if labels is None:
            if is_multichannel:
                labels = [f'Ch{i}' for i in range(n_channels)]
            else:
                labels = ['PSD']

        fig, ax = plot_line(
            x=f,
            y=psds,  # [n_channels, n_freqs]
            labels=labels,
            xlabel=xlabel,
            ylabel=resolved_ylabel,
            ax=ax,
            show=False,  # 先不显示，后面统一处理
            auto_close=False,
            figsize=figsize,
        )

    if title:
        ax.set_title(title)

    # 显示图例
    ax.legend()

    # 最终化图形
    _finalize_figure(fig, show, save_path, auto_close)

    return fig


def plot_stft(
    signals: np.ndarray,
    fs: float,
    averaging: bool = True,
    window: str = 'hamming',
    nperseg: int | None = None,
    noverlap: int | None = None,
    mode: Literal['magnitude', 'phase', 'psd'] = 'magnitude',
    fmin: float = 0,
    fmax: float = 60,
    cmap: str = 'viridis',
    vmin: float | None = None,
    vmax: float | None = None,
    xlabel: str = 'Time (s)',
    ylabel: str = 'Frequency (Hz)',
    title: str | None = None,
    show_colorbar: bool = True,
    colorbar_label: str | object | None = AUTO_SPECIFIED,
    ax: Optional[Axes] = None,
    show: bool = True,
    save_path: str | None = None,
    auto_close: bool | None = None,
    figsize: tuple[float, float] | None = None,
    **stft_kwargs,
):
    """绘制信号的 STFT 时频图

    使用短时傅里叶变换 (STFT) 分析信号的时频特性，并绘制时频图。
    支持单通道和多通道信号，多通道时可以选择显示平均结果。

    Parameters
    ----------
    - `signals` : `np.ndarray`
        - 输入信号。一维数组为单通道信号，二维数组为多通道信号 [n_channels, n_samples]
    - `fs` : `float`
        - 采样频率，单位为 Hz。
    - `averaging` : `bool`, optional
        - 对于多通道信号，是否显示平均结果, by default True
        - True: 显示所有通道的平均STFT结果
        - False: 仅显示第一个通道的STFT（用于兼容）
        - 对于单通道信号，此参数无效
    - `window` : `str`, optional
        - 窗函数类型。默认为 'hamming'。
    - `nperseg` : `int`, optional
        - 每段的长度。默认为 `int(0.5 * fs)`。
    - `noverlap` : `int`, optional
        - 段间重叠的样本数。默认为 `nperseg // 2`。
    - `mode` : `Literal['magnitude', 'phase', 'psd']`, optional
        - 显示模式。
        - 'magnitude': 幅度谱
        - 'phase': 相位谱
        - 'psd': 功率谱密度
        - 默认为 'magnitude'。
    - `fmin` : `float`, optional
        - 显示的最小频率，单位为 Hz。默认为 0。
    - `fmax` : `float`, optional
        - 显示的最大频率，单位为 Hz。默认为 60。
    - `cmap` : `str`, optional
        - 颜色映射。默认为 'viridis'。
    - `vmin` : `float`, optional
        - 颜色映射的最小值。默认为 None（自动）。
    - `vmax` : `float`, optional
        - 颜色映射的最大值。默认为 None（自动）。
    - `xlabel` : `str`, optional
        - X 轴标签。默认为 'Time (s)'。
    - `ylabel` : `str`, optional
        - Y 轴标签。默认为 'Frequency (Hz)'。
    - `title` : `str`, optional
        - 图标题。默认为 None。
    - `show_colorbar` : `bool`, optional
        - 是否显示颜色条。默认为 True。
    - `colorbar_label` : `str`, optional
        - 颜色条标签。默认为 None。
    - `ax` : `plt.Axes`, optional
        - 用于绘图的坐标轴。默认为 None（创建新的）。
    - `show` : `bool`, optional
        - 是否显示图形。默认为 True。
    - `save_path` : `str`, optional
        - 保存路径。默认为 None（不保存）。
    - `auto_close` : `bool`, optional
        - 是否自动关闭图形。默认为 None（自动判断）。
    - `**stft_kwargs`
        - 传递给 STFT 函数的额外参数。

    Returns
    ----------
    - `fig` : `matplotlib.figure.Figure`
        - 生成的图形对象。

    Examples
    ----------
    **单通道信号**：
    ```py
    >>> import numpy as np
    >>> from zyplib.vis.plot_eeg import plot_stft
    >>>
    >>> # 生成测试信号：频率变化的正弦波
    >>> fs = 1000
    >>> t = np.arange(0, 3, 1/fs)
    >>> f1, f2 = 50, 150
    >>> signal = np.sin(2*np.pi*f1*t[:len(t)//2]) + np.sin(2*np.pi*f2*t[len(t)//2:])
    >>>
    >>> # 绘制 STFT 时频图
    >>> fig = plot_stft(signal, fs, mode='magnitude')
    ```

    **多通道信号（平均）**：
    ```py
    >>> # 生成多通道信号
    >>> n_channels = 8
    >>> signals = np.random.randn(n_channels, len(t)) * 0.1
    >>> for i in range(n_channels):
    ...     signals[i] += np.sin(2*np.pi*f1*t[:len(t)//2]) + np.sin(2*np.pi*f2*t[len(t)//2:])
    >>>
    >>> # 显示平均 STFT
    >>> fig = plot_stft(signals, fs, averaging=True, mode='magnitude')
    ```

    **多通道信号（第一个通道）**：
    ```py
    >>> # 仅显示第一个通道的 STFT
    >>> fig = plot_stft(signals, fs, averaging=False, mode='magnitude')
    ```
    """
    # 导入 STFT 函数（延迟导入避免循环依赖）
    from zyplib.signal.time_frequency import stft

    # 确保输入格式正确
    signals = ensure_ndarray(signals)

    # 判断信号维度并处理
    if signals.ndim == 1:
        # 单通道信号
        signal_to_process = signals
        is_multichannel = False
    elif signals.ndim == 2:
        # 多通道信号 [n_channels, n_samples]
        is_multichannel = True
        if averaging:
            # 计算所有通道的平均信号
            signal_to_process = np.mean(signals, axis=0)
        else:
            # 使用第一个通道
            signal_to_process = signals[0]
    else:
        raise ValueError('signals 必须是一维或二维数组')

    # 计算 STFT
    f, t, Zxx = stft(
        signal_to_process,
        fs,
        window=window,
        nperseg=nperseg,
        noverlap=noverlap,
        **stft_kwargs,
    )

    # 根据模式计算显示数据
    if mode == 'magnitude':
        data = np.abs(Zxx)
        default_colorbar_label = 'Magnitude'
        if vmin is None:
            vmin = 0
    elif mode == 'phase':
        data = np.angle(Zxx)
        default_colorbar_label = 'Phase (rad)'
        cmap = cmap if cmap != 'viridis' else 'hsv'  # 相位图更适合用 hsv
    elif mode == 'psd':
        data = np.abs(Zxx) ** 2
        default_colorbar_label = 'Power Spectral Density'
        if vmin is None:
            vmin = 0
    else:
        raise ValueError(f'不支持的模式: {mode}')

    # 频率范围筛选
    freq_mask = (f >= fmin) & (f <= fmax)
    f_filtered = f[freq_mask]
    data_filtered = data[freq_mask, :]

    # 准备绘图
    fig, ax = _prepare_ax(ax, figsize=figsize)

    # 绘制时频图
    im = ax.imshow(
        data_filtered,
        aspect='auto',
        origin='lower',
        extent=(float(t[0]), float(t[-1]), float(f_filtered[0]), float(f_filtered[-1])),
        cmap=cmap,
        vmin=vmin,
        vmax=vmax,
        interpolation='bilinear',
    )

    # 设置标签和标题
    if xlabel:
        ax.set_xlabel(xlabel)
    if ylabel:
        ax.set_ylabel(ylabel)

    if title:
        ax.set_title(title)

    # 添加颜色条
    if colorbar_label is AUTO_SPECIFIED:
        colorbar_label = default_colorbar_label

    if show_colorbar:
        fig.colorbar(im, ax=ax, label=colorbar_label)

    # 完成绘图
    _finalize_figure(fig, show, save_path, auto_close)

    return fig


def plot_cwt(
    signals: np.ndarray,
    fs: float,
    averaging: bool = True,
    wavelet: str = 'cmor',
    n_scales: int = 64,
    fmin: float = 0.5,
    fmax: float = 60,
    mode: Literal['magnitude', 'phase', 'power'] = 'magnitude',
    cmap: str = 'viridis',
    vmin: float | None = None,
    vmax: float | None = None,
    xlabel: str = 'Time (s)',
    ylabel: str = 'Frequency (Hz)',
    title: str | None = None,
    show_colorbar: bool = True,
    colorbar_label: str | object | None = AUTO_SPECIFIED,
    ax: Optional[Axes] = None,
    show: bool = True,
    save_path: str | None = None,
    auto_close: bool | None = None,
    figsize: tuple[float, float] | None = None,
    **cwt_kwargs,
):
    """绘制信号的连续小波变换 (CWT) 时频图

    对信号进行连续小波变换并绘制时频图。特别适用于脑电信号等低频成分重要的信号分析。
    相比STFT，小波变换提供更好的多分辨率分析：低频处有更好的频率分辨率，高频处有更好的时间分辨率。
    支持单通道和多通道信号，多通道时可以选择显示平均结果。

    Parameters
    ----------
    - `signals` : `np.ndarray`
        - 输入信号。一维数组为单通道信号，二维数组为多通道信号 [n_channels, n_samples]
    - `fs` : `float`
        - 采样频率，单位为 Hz。
    - `averaging` : `bool`, optional
        - 对于多通道信号，是否显示平均结果, by default True
        - True: 显示所有通道的平均CWT结果
        - False: 仅显示第一个通道的CWT（用于兼容）
        - 对于单通道信号，此参数无效
    - `wavelet` : `str`, optional
        - 小波基函数名称，默认为 'cmor'（复Morlet小波，最适合时频分析）。
        - 常用选择：'cmor', 'morl', 'mexh', 'gaus1'等
    - `n_scales` : `int`, optional
        - 小波变换的尺度数量，默认为64。增大可获得更精细的频率分辨率，但计算成本更高。
    - `fmin` : `float`, optional
        - 显示的最小频率，单位为 Hz。默认为 0.5（适合脑电低频成分分析）。
    - `fmax` : `float`, optional
        - 显示的最大频率，单位为 Hz。默认为 60（覆盖脑电主要频段）。
    - `mode` : `Literal['magnitude', 'phase', 'power']`, optional
        - 显示模式。
        - 'magnitude': 幅度谱（绝对值）
        - 'phase': 相位谱
        - 'power': 功率谱（平方）
        - 默认为 'magnitude'。
    - `cmap` : `str`, optional
        - 颜色映射。默认为 'viridis'。
    - `vmin` : `float`, optional
        - 颜色映射的最小值。默认为 None（自动）。
    - `vmax` : `float`, optional
        - 颜色映射的最大值。默认为 None（自动）。
    - `xlabel` : `str`, optional
        - X 轴标签。默认为 'Time (s)'。
    - `ylabel` : `str`, optional
        - Y 轴标签。默认为 'Frequency (Hz)'。
    - `title` : `str`, optional
        - 图标题。默认为 None。
    - `show_colorbar` : `bool`, optional
        - 是否显示颜色条。默认为 True。
    - `colorbar_label` : `str`, optional
        - 颜色条标签。默认为 None。
    - `ax` : `plt.Axes`, optional
        - 用于绘图的坐标轴。默认为 None（创建新的）。
    - `show` : `bool`, optional
        - 是否显示图形。默认为 True。
    - `save_path` : `str`, optional
        - 保存路径。默认为 None（不保存）。
    - `auto_close` : `bool`, optional
        - 是否自动关闭图形。默认为 None（自动判断）。
    - `**cwt_kwargs`
        - 传递给 CWT 函数的额外参数。

    Returns
    ----------
    - `fig` : `matplotlib.figure.Figure`
        - 生成的图形对象。

    Examples
    ----------
    **单通道信号**：
    ```py
    >>> import numpy as np
    >>> from zyplib.vis.plot_eeg import plot_cwt
    >>>
    >>> # 生成测试信号：频率变化的正弦波
    >>> fs = 1000
    >>> t = np.arange(0, 3, 1/fs)
    >>> f1, f2 = 10, 30
    >>> signal = np.sin(2*np.pi*f1*t[:len(t)//2]) + np.sin(2*np.pi*f2*t[len(t)//2:])
    >>>
    >>> # 绘制 CWT 时频图
    >>> fig = plot_cwt(signal, fs, wavelet='cmor', fmin=1, fmax=50, mode='magnitude')
    ```

    **多通道信号（平均）**：
    ```py
    >>> # 生成多通道信号
    >>> n_channels = 8
    >>> signals = np.random.randn(n_channels, len(t)) * 0.1
    >>> for i in range(n_channels):
    ...     signals[i] += np.sin(2*np.pi*f1*t[:len(t)//2]) + np.sin(2*np.pi*f2*t[len(t)//2:])
    >>>
    >>> # 显示平均 CWT
    >>> fig = plot_cwt(signals, fs, averaging=True, wavelet='morl')
    ```
    """
    # 导入 CWT 函数（延迟导入避免循环依赖）
    from zyplib.signal.time_frequency import cwt
    from zyplib.utils.ensure import ensure_ndarray

    # 确保输入格式正确
    signals = ensure_ndarray(signals)

    # 判断信号维度并处理
    if signals.ndim == 1:
        # 单通道信号
        signal_to_process = signals
        is_multichannel = False
    elif signals.ndim == 2:
        # 多通道信号 [n_channels, n_samples]
        is_multichannel = True
        if averaging:
            # 计算所有通道的平均信号
            signal_to_process = np.mean(signals, axis=0)
        else:
            # 使用第一个通道
            signal_to_process = signals[0]
    else:
        raise ValueError('signals 必须是一维或二维数组')

    # 计算 CWT
    scales, freqs, t, coef = cwt(
        signal_to_process,
        fs,
        wavelet=wavelet,
        n_scales=n_scales,
        fmin=fmin,
        fmax=fmax,
        **cwt_kwargs,
    )

    # 根据模式计算显示数据
    if mode == 'magnitude':
        data = np.abs(coef)
        default_colorbar_label = 'Magnitude'
        if vmin is None:
            vmin = 0
    elif mode == 'phase':
        data = np.angle(coef)
        default_colorbar_label = 'Phase (rad)'
        cmap = cmap if cmap != 'viridis' else 'hsv'  # 相位图更适合用 hsv
    elif mode == 'power':
        data = np.abs(coef) ** 2
        default_colorbar_label = 'Power'
        if vmin is None:
            vmin = 0
    else:
        raise ValueError(f'不支持的模式: {mode}')

    # 频率范围筛选
    freq_mask = (freqs >= fmin) & (freqs <= fmax)
    freqs_filtered = freqs[freq_mask]
    data_filtered = data[freq_mask, :]

    # 准备绘图
    fig, ax = _prepare_ax(ax, figsize=figsize)

    # 绘制时频图
    im = ax.imshow(
        data_filtered,
        aspect='auto',
        origin='lower',
        extent=(
            float(t[0]),
            float(t[-1]),
            float(freqs_filtered[0]),
            float(freqs_filtered[-1]),
        ),
        cmap=cmap,
        vmin=vmin,
        vmax=vmax,
        interpolation='bilinear',
    )

    # 设置标签和标题
    if xlabel:
        ax.set_xlabel(xlabel)
    if ylabel:
        ax.set_ylabel(ylabel)

    if title:
        ax.set_title(title)
    elif is_multichannel and averaging:
        ax.set_title(f'平均CWT时频图 ({signals.shape[0]}个通道)')
    elif is_multichannel and not averaging:
        ax.set_title('CWT时频图 (通道1)')
    else:
        ax.set_title('CWT时频图')

    # 添加颜色条
    if colorbar_label is AUTO_SPECIFIED:
        colorbar_label = default_colorbar_label

    if show_colorbar:
        fig.colorbar(im, ax=ax, label=colorbar_label)

    # 完成绘图
    _finalize_figure(fig, show, save_path, auto_close)

    return fig


def plot_dwt(
    signals: np.ndarray,
    fs: float,
    averaging: bool = True,
    wavelet: str = 'db4',
    level: Optional[int] = None,
    mode: str = 'symmetric',
    plot_type: Literal['image', 'coef'] = 'image',
    cmap: str = 'viridis',
    vmin: float | None = None,
    vmax: float | None = None,
    xlabel: str = 'Time (s)',
    ylabel: str = 'Level',
    title: str | None = None,
    show_colorbar: bool = True,
    colorbar_label: str | object | None = AUTO_SPECIFIED,
    show_freq_bands: bool = True,
    ax: Optional[Axes] = None,
    show: bool = True,
    save_path: str | None = None,
    auto_close: bool | None = None,
    figsize: tuple[float, float] | None = None,
    **dwt_kwargs,
):
    """绘制信号的离散小波变换 (DWT) 结果

    对信号进行离散小波变换并绘制结果。DWT将信号分解为不同尺度的小波系数，对应不同频率带，
    特别适用于分析信号中的多尺度特征，例如脑电信号的不同频段活动。
    支持单通道和多通道信号，多通道时可以选择显示平均结果。

    Parameters
    ----------
    - `signals` : `np.ndarray`
        - 输入信号。一维数组为单通道信号，二维数组为多通道信号 [n_channels, n_samples]
    - `fs` : `float`
        - 采样频率，单位为 Hz。
    - `averaging` : `bool`, optional
        - 对于多通道信号，是否显示平均结果, by default True
        - True: 显示所有通道的平均DWT结果
        - False: 仅显示第一个通道的DWT（用于兼容）
        - 对于单通道信号，此参数无效
    - `wavelet` : `str`, optional
        - 小波基函数名称，默认为 'db4'（Daubechies 4阶小波）。
        - 常用选择：'db'系列(db1-db38), 'sym'系列, 'coif'系列, 'haar'等
    - `level` : `int`, optional
        - 分解层数。如果为None，则自动选择最大可分解层数。
    - `mode` : `str`, optional
        - 边界处理模式，默认为'symmetric'（对称扩展）。
    - `plot_type` : `Literal['image', 'coef']`, optional
        - 绘图类型, by default 'image'
        - 'image': 绘制时频图像
        - 'coef': 绘制各层系数波形图
    - `cmap` : `str`, optional
        - 颜色映射。默认为 'viridis'。仅在 `plot_type='image'` 时有效。
    - `vmin` : `float`, optional
        - 颜色映射的最小值。默认为 None（自动）。仅在 `plot_type='image'` 时有效。
    - `vmax` : `float`, optional
        - 颜色映射的最大值。默认为 None（自动）。仅在 `plot_type='image'` 时有效。
    - `xlabel` : `str`, optional
        - X 轴标签。默认为 'Time (s)'。
    - `ylabel` : `str`, optional
        - Y 轴标签。默认为 'Level'，仅在 `plot_type='image'` 时有效。
    - `title` : `str`, optional
        - 图标题。默认为 None。
    - `show_colorbar` : `bool`, optional
        - 是否显示颜色条。默认为 True。仅在 `plot_type='image'` 时有效。
    - `colorbar_label` : `str`, optional
        - 颜色条标签。默认为 'Coefficient'。仅在 `plot_type='image'` 时有效。
    - `show_freq_bands` : `bool`, optional
        - 是否显示各层对应的频率范围。默认为 True。
    - `ax` : `plt.Axes`, optional
        - 用于绘图的坐标轴。默认为 None（创建新的）。
        - 注意：在 `plot_type='coef'` 模式下，ax 会被忽略，会创建新的子图组。
    - `show` : `bool`, optional
        - 是否显示图形。默认为 True。
    - `save_path` : `str`, optional
        - 保存路径。默认为 None（不保存）。
    - `auto_close` : `bool`, optional
        - 是否自动关闭图形。默认为 None（自动判断）。
    - `**dwt_kwargs`
        - 传递给 DWT 函数的额外参数。

    Returns
    ----------
    - `fig` : `matplotlib.figure.Figure`
        - 生成的图形对象。

    Examples
    ----------
    **单通道信号时频图**：
    ```py
    >>> import numpy as np
    >>> from zyplib.vis.plot_eeg import plot_dwt
    >>>
    >>> # 生成测试信号
    >>> fs = 1000
    >>> t = np.arange(0, 3, 1/fs)
    >>> signal = np.sin(2*np.pi*10*t) + 0.5*np.sin(2*np.pi*30*t) + 0.2*np.random.randn(len(t))
    >>>
    >>> # 绘制 DWT 时频图
    >>> fig = plot_dwt(signal, fs, wavelet='db4', level=5, plot_type='image')
    ```

    **绘制各层系数**：
    ```py
    >>> # 绘制各层系数波形图
    >>> fig = plot_dwt(signal, fs, wavelet='db4', level=5, plot_type='coef')
    ```
    """
    # 导入 DWT 函数（延迟导入避免循环依赖）
    from zyplib.signal.time_frequency import dwt
    from zyplib.utils.ensure import ensure_ndarray

    # 确保输入格式正确
    signals = ensure_ndarray(signals)

    # 判断信号维度并处理
    if signals.ndim == 1:
        # 单通道信号
        signal_to_process = signals
        is_multichannel = False
    elif signals.ndim == 2:
        # 多通道信号 [n_channels, n_samples]
        is_multichannel = True
        if averaging:
            # 计算所有通道的平均信号
            signal_to_process = np.mean(signals, axis=0)
        else:
            # 使用第一个通道
            signal_to_process = signals[0]
    else:
        raise ValueError('signals 必须是一维或二维数组')

    # 计算 DWT
    coeffs, freqs, t_list = dwt(
        signal_to_process, fs, wavelet=wavelet, level=level, mode=mode, **dwt_kwargs
    )

    actual_level = len(coeffs) - 1  # 实际分解层数

    # 设置默认标题
    if title is None:
        if is_multichannel and averaging:
            default_title = f'DWT分解 ({signals.shape[0]}个通道平均)'
        elif is_multichannel and not averaging:
            default_title = 'DWT分解 (通道1)'
        else:
            default_title = 'DWT分解'
    else:
        default_title = title

    # 根据绘图类型选择不同的绘图方式
    if plot_type == 'image':
        # 时频图绘制

        # 准备绘图
        fig, ax = _prepare_ax(ax, figsize=figsize or (10, 6))

        # 为绘图准备数据
        # 创建统一的时间网格以进行插值
        n_samples = len(signal_to_process)
        t_uniform = np.linspace(0, n_samples / fs, n_samples)

        # 创建图像数据数组
        image_data = np.zeros((actual_level + 1, n_samples))

        # 对每层系数进行插值以匹配原始信号长度
        for i, (coef, t_coef) in enumerate(zip(coeffs, t_list)):
            # 使用线性插值扩展到原始信号长度
            from scipy.interpolate import interp1d

            if len(t_coef) > 1:  # 确保有足够的点进行插值
                f = interp1d(
                    t_coef,
                    coef,
                    kind='linear',
                    bounds_error=False,
                    fill_value=typing.cast(typing.Any, 'extrapolate'),
                )
                image_data[i] = f(t_uniform)
            else:  # 处理边缘情况
                image_data[i] = np.full(n_samples, coef[0])

        # 绘制时频图
        im = ax.imshow(
            image_data,
            aspect='auto',
            origin='lower',
            extent=(0.0, float(n_samples / fs), 0.0, float(actual_level + 1)),
            cmap=cmap,
            vmin=vmin,
            vmax=vmax,
            interpolation='bilinear',
        )

        # 设置标签和刻度
        ax.set_xlabel(xlabel)
        ax.set_ylabel(ylabel)

        # 设置y轴刻度（层数）
        y_ticks = np.arange(0.5, actual_level + 1.5, 1)
        y_labels = ['cA'] + [f'cD{actual_level-i+1}' for i in range(actual_level)]
        ax.set_yticks(y_ticks)
        ax.set_yticklabels(y_labels)

        # 添加频率带标注
        if show_freq_bands:
            for i, freq_band in enumerate(freqs):
                ax.text(
                    n_samples / fs * 1.01,  # 右侧位置
                    i + 0.5,  # 层位置
                    f'{freq_band[0]:.1f}-{freq_band[1]:.1f} Hz',
                    verticalalignment='center',
                    fontsize=8,
                )

        # 设置标题
        ax.set_title(default_title)

        # 添加颜色条
        if colorbar_label is AUTO_SPECIFIED:
            colorbar_label = 'Coefficient'

        if show_colorbar:
            fig.colorbar(im, ax=ax, label=colorbar_label)

    elif plot_type == 'coef':
        # 系数波形图绘制
        # 为每层系数创建子图
        n_subplots = len(coeffs)
        fig_height = n_subplots * 1.5  # 根据层数调整图形高度
        fig = plt.figure(figsize=figsize or (10, fig_height))

        # 绘制每层的系数
        for i, (coef, t_coef, freq_band) in enumerate(zip(coeffs, t_list, freqs)):
            ax = fig.add_subplot(n_subplots, 1, i + 1)
            ax.plot(t_coef, coef)

            # 设置标签
            if i == n_subplots - 1:  # 最后一个子图
                ax.set_xlabel(xlabel)

            # 添加系数类型和频率带标签
            if i == 0:
                coef_type = f'cA{actual_level}'
            else:
                coef_type = f'cD{actual_level-i+1}'

            freq_label = (
                f'{freq_band[0]:.1f}-{freq_band[1]:.1f} Hz' if show_freq_bands else ''
            )
            ax.set_ylabel(f'{coef_type}\n{freq_label}')

            # 设置网格和限制范围
            ax.grid(True, alpha=0.3)
            ax.set_xlim(0, t_coef[-1])

        # 调整子图间距
        plt.tight_layout()

        # 设置总标题
        fig.suptitle(default_title, fontsize=12)
        plt.subplots_adjust(top=0.95)  # 为标题留出空间

    else:
        raise ValueError(f"不支持的绘图类型: {plot_type}，支持的类型: 'image', 'coef'")

    # 完成绘图
    _finalize_figure(fig, show, save_path, auto_close)

    return fig
