import typing

# Data augmentation for EEG signals - PyTorch only version
from typing import Literal, Optional, Tuple, cast, overload

import numpy as np
import torch

from zyplib.signal.filter import notch
from zyplib.utils.ensure import ensure_3dims

__all__ = [
    'sign_flip',
    'time_reverse',
    'add_white_noise',
    'scale_amplitude',
    'add_spike',
    'bandstop_aug',
    'ft_surrogate',
]


def _apply_notch(
    X: np.ndarray, srate: float, f0: float, Q: float, axis: int
) -> np.ndarray:
    return cast(np.ndarray, notch(X, srate, f0=f0, Q=Q, axis=axis))


def _use_random_state(random_state=None, device=None) -> torch.Generator:
    """检查和转换随机状态对象

    为 PyTorch 提供统一的随机数生成接口

    Parameters
    ----------
    - `random_state` : `None | int | torch.Generator`, optional
        随机状态或种子，默认为 None
    - `device` : `str | torch.device`, optional
        PyTorch 设备，用于生成设备感知的随机数，默认为 None

    Returns
    -------
    - `torch.Generator`
        随机状态对象
    """
    if isinstance(random_state, torch.Generator):
        return random_state
    elif random_state is None:
        return torch.Generator(device=device)
    elif isinstance(random_state, int):
        generator = torch.Generator(device=device)
        generator.manual_seed(random_state)
        return generator
    else:
        raise TypeError(f'Unsupported random_state type: {type(random_state)}')


def sign_flip(X: torch.Tensor) -> torch.Tensor:
    """上下颠倒信号

    将信号进行垂直翻转（乘以-1）。

    Parameters
    ----------
    - `X` : `torch.Tensor`
        - 输入数据，形状为 (C, T) 或 (N, C, T)

    Returns
    ----------
    - `torch.Tensor`
        - 翻转后的数据，形状与输入相同
    """
    return -X


def time_reverse(X: torch.Tensor) -> torch.Tensor:
    """时间反转

    将信号在时间维度上反转。

    Parameters
    ----------
    - `X` : `torch.Tensor`
        - 输入数据，形状为 (C, T) 或 (N, C, T)

    Returns
    ----------
    - `torch.Tensor`
        - 时间反转后的数据，形状与输入相同
    """
    return torch.flip(X, dims=[-1])


def add_white_noise(
    X: torch.Tensor,
    scale: float = 0.1,
    rel_power: bool = False,
    random_state: Optional[int] = None,
) -> torch.Tensor:
    """添加白噪声

    向信号中添加高斯白噪声。提供两种噪声生成方式：
    1. 绝对噪声强度（rel_power=False）：生成指定标准差的噪声
    2. 相对噪声强度（rel_power=True）：生成相对于原始信号功率的噪声

    Parameters
    ----------
    - `X` : `torch.Tensor`
        - 输入数据，形状为 (C, T) 或 (N, C, T)
    - `std` : `float`, optional
        - 噪声的标准差或相对功率，by default 0.1
        - 当 rel_power=False 时，表示噪声的绝对标准差
        - 当 rel_power=True 时，表示相对于信号功率的比例
    - `rel_power` : `bool`, optional
        - 是否使用相对功率模式，by default False
    - `random_state` : `Optional[int]`, optional
        - 随机种子，用于生成可重复的噪声，by default None

    Returns
    ----------
    - `torch.Tensor`
        - 添加噪声后的数据，形状与输入相同
    """
    device = X.device
    generator = _use_random_state(random_state, device=device)

    if rel_power:
        signal_power = torch.var(X, dim=-1, keepdim=True)
        noise_std = torch.sqrt(signal_power * scale)
        noise = (
            torch.randn(X.shape, generator=generator, device=device, dtype=X.dtype)
            * noise_std
        )
    else:
        noise = (
            torch.randn(X.shape, generator=generator, device=device, dtype=X.dtype)
            * scale
        )

    return X + noise


def scale_amplitude(
    X: torch.Tensor,
    scale_range: Tuple[float, float] = (0.9, 1.1),
    granularity: Literal['global', 'per_epoch', 'per_channel'] = 'per_channel',
    random_state: Optional[int] = None,
) -> torch.Tensor:
    """
    随机缩放信号的振幅。

    默认的行为是，随机扰动各个通道的幅值，在 [0.9, 1.1] 范围内。

    Parameters
    ----------
    - `X` : `torch.Tensor`
        - 输入数据，形状为 (C, T) 或 (N, C, T)。
    - `scale_range` : `Tuple[float, float]`, optional
        - 缩放因子的随机采样范围, by default (0.8, 1.2)。
    - `granularity` : `Literal['global', 'per_epoch', 'per_channel']`, optional
        - 缩放因子的作用粒度, by default 'per_channel'。
        - 'global': 整个输入使用一个随机缩放因子。
        - 'per_epoch': 每个 epoch (试验) 使用一个独立的随机因子。
        - 'per_channel': 每个通道使用一个独立的随机因子。
    - `random_state` : `Optional[int]`, optional
        - 随机种子，用于生成可复现的缩放因子, by default None。

    Returns
    -------
    - `torch.Tensor`
        - 振幅缩放后的数据，形状与输入相同。
    """
    device = X.device
    generator = _use_random_state(random_state, device=device)
    original_shape = X.shape
    X_proc = cast(torch.Tensor, ensure_3dims(X, newaxis='batch'))  # 返回 (N, C, T)

    N, C, T = X_proc.shape

    # 1. 根据清晰的粒度参数确定缩放因子的形状
    if granularity == 'global':
        shape = (1,)
    elif granularity == 'per_epoch':
        shape = (N,)
    elif granularity == 'per_channel':
        shape = (N, C)
    else:
        raise ValueError(f'不支持的粒度: {granularity}')

    # 2. 生成随机缩放因子
    scale_factor = torch.rand(shape, generator=generator, device=device, dtype=X.dtype)
    scale_factor = scale_factor * (scale_range[1] - scale_range[0]) + scale_range[0]

    # 3. 调整形状以进行广播
    while scale_factor.ndim < X_proc.ndim:
        scale_factor = scale_factor.unsqueeze(-1)

    # 4. 应用缩放并返回正确形状
    return (X_proc * scale_factor).reshape(original_shape)


def _make_gaussian_spike(
    t_width: float,
    srate: int,
    amplitude: float,
    *,
    device: Optional[torch.device] = None,
    dtype: Optional[torch.dtype] = None,
) -> torch.Tensor:
    """生成尖峰波形

    使用高斯包络生成持续时间极短的尖峰信号。

    Parameters
    ----------
    - `t_width` : `float`
        - 尖峰宽度（秒）
    - `srate` : `int`
        - 采样率（Hz）
    - `amplitude` : `float`
        - 尖峰峰值振幅，可正可负
    - `device` : `torch.device`, optional
        - PyTorch 设备
    - `dtype` : `torch.dtype`, optional
        - 数据类型

    Returns
    ----------
    - `torch.Tensor`
        - 长度为 round(t_width * srate) 的 1-D 波形
    """
    # Ensure the spike occupies at least one sample
    length = max(1, int(round(t_width * srate)))

    # Prefer an odd length so that the spike is centred
    if length % 2 == 0:
        length += 1

    t = torch.arange(length, device=device) - length // 2
    sigma = length / 6.0  # ≈99.7% of Gaussian within ±3σ (here ±length/2)
    spike = amplitude * torch.exp(-0.5 * (t / sigma) ** 2)
    if dtype is not None:
        spike = spike.to(dtype=dtype)
    return spike


@overload
def add_spike(
    X: torch.Tensor,
    srate: int,
    number: int,
    t_width: float,
    *,
    amplitude_mode: Literal[
        'fixed', 'random_absolute', 'random_relative', 'relative'
    ] = 'random_relative',
    amplitude: Tuple[float, float] | float = (-1.5, 1.5),
    position_granularity: Literal['per_epoch', 'per_channel'] = 'per_epoch',
    random_state: Optional[int] = None,
    return_spike_pos: Literal[False] = False,
) -> torch.Tensor: ...


@overload
def add_spike(
    X: torch.Tensor,
    srate: int,
    number: int,
    t_width: float,
    *,
    amplitude_mode: Literal[
        'fixed', 'random_absolute', 'random_relative', 'relative'
    ] = 'random_relative',
    amplitude: Tuple[float, float] | float = (-1.5, 1.5),
    position_granularity: Literal['per_epoch', 'per_channel'] = 'per_epoch',
    random_state: Optional[int] = None,
    return_spike_pos: Literal[True],
) -> tuple[torch.Tensor, torch.Tensor]: ...


def add_spike(
    X: torch.Tensor,
    srate: int,
    number: int,
    t_width: float,
    *,
    amplitude_mode: Literal[
        'fixed', 'random_absolute', 'random_relative', 'relative'
    ] = 'random_relative',
    amplitude: Tuple[float, float] | float = (-1.5, 1.5),
    position_granularity: Literal['per_epoch', 'per_channel'] = 'per_epoch',
    random_state: Optional[int] = None,
    return_spike_pos: bool = False,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
    """
    添加尖峰噪声。

    Parameters
    ----------
    - `X` : `torch.Tensor`
        - 输入数据，形状为 (C, T) 或 (N, C, T)。
    - `srate` : `int`
        - 采样率（Hz）。
    - `number` : `int`
        - 每个粒度单位（epoch或channel）插入的尖峰数量。
    - `t_width` : `float`
        - 单个尖峰的持续时间（秒）。
    - `amplitude_mode` : `Literal['fixed', 'random_absolute', 'random_relative', 'relative']`, optional
        - 振幅生成模式, by default 'random_relative'。
        - 'fixed': 使用固定的绝对振幅。
        - 'random_absolute': 从范围内随机采样一个绝对振幅。
        - 'random_relative': 从范围内随机采样一个比例因子，乘以通道标准差作为振幅。
        - 'relative': 使用固定的相对比例，乘以通道标准差作为振幅。
        - **注意**: 所有的相对模式都是相对于所在的通道，而非完整的 epoch 信号的标准差。
    - `amplitude` : `Tuple[float, float] | float`, optional
        - 根据 amplitude_mode 提供的值, by default (-1.5, 1.5)。
        - 如果 mode='fixed' 或 'relative', 应提供 float。
        - 如果 mode 是 'random_*', 应提供一个 (min, max) 的 tuple。
    - `position_granularity` : `Literal['per_epoch', 'per_channel']`, optional
        - 尖峰位置的生成粒度, by default 'per_epoch'。
        - 'per_epoch': 在每个 epoch 内生成一组共享的位置，应用到所有通道。
        - 'per_channel': 为每个通道独立生成一组位置。
    - `random_state` : `Optional[int]`, optional
        - 随机种子，保证可复现
    - `return_spike_pos` : `bool`, optional
        - 是否返回尖峰位置，by default False。

    Returns
    -------
    - `torch.Tensor`
        - 添加尖峰后的数据，形状与输入相同。
    """
    device = X.device
    dtype = X.dtype
    generator = _use_random_state(random_state, device=device)

    original_shape = X.shape
    X_proc = cast(torch.Tensor, ensure_3dims(X.clone(), newaxis='batch'))
    N, C, T = X_proc.shape

    actual_spike_len = len(_make_gaussian_spike(t_width, srate, 1.0, device=device))
    if T < actual_spike_len:
        raise ValueError(f'信号长度 {T} 小于尖峰长度 {actual_spike_len}')

    def _get_amplitude(batch_idx: int, ch_idx: int) -> float:
        if amplitude_mode == 'fixed':
            if not isinstance(amplitude, (float, int)):
                raise TypeError("For 'fixed' mode, amplitude must be a float or int.")
            return float(amplitude)
        elif amplitude_mode == 'relative':
            if not isinstance(amplitude, (float, int)):
                raise TypeError("For 'relative' mode, amplitude must be a float or int.")
            return float(amplitude) * torch.std(X_proc[batch_idx, ch_idx]).item()

        if not isinstance(amplitude, tuple) or len(amplitude) != 2:
            raise TypeError('For random modes, amplitude must be a (min, max) tuple.')

        rand_val = torch.rand(1, generator=generator, device=device).item()
        rand_val = rand_val * (amplitude[1] - amplitude[0]) + amplitude[0]

        if amplitude_mode == 'random_absolute':
            return rand_val
        elif amplitude_mode == 'random_relative':
            return rand_val * torch.std(X_proc[batch_idx, ch_idx]).item()
        else:
            raise ValueError(f'Unknown amplitude_mode: {amplitude_mode}')

    def _get_random_positions() -> torch.Tensor:
        if number <= (T - actual_spike_len):
            # 使用 torch.randperm 来避免重复
            positions = torch.randperm(
                T - actual_spike_len, generator=generator, device=device
            )[:number]
        else:
            positions = torch.randint(
                0, T - actual_spike_len, (number,), generator=generator, device=device
            )
        return positions

    for n in range(N):
        if position_granularity == 'per_channel':
            for c in range(C):
                positions = _get_random_positions()
                for pos in positions:
                    pos_int = int(pos.item())
                    amp = _get_amplitude(n, c)
                    spike = cast(
                        torch.Tensor,
                        _make_gaussian_spike(
                            t_width, srate, amp, device=device, dtype=dtype
                        ),
                    )
                    end_pos = pos_int + len(spike)
                    X_proc[n, c, pos_int:end_pos] += spike
        elif position_granularity == 'per_epoch':
            # 先为整个 epoch 生成共享的位置
            positions = _get_random_positions()
            for pos in positions:
                pos_int = int(pos.item())
                # 振幅仍然是每个通道独立计算的，这通常是期望的行为
                for c in range(C):
                    amp = _get_amplitude(n, c)
                    spike = cast(
                        torch.Tensor,
                        _make_gaussian_spike(
                            t_width, srate, amp, device=device, dtype=dtype
                        ),
                    )
                    end_pos = pos_int + len(spike)
                    X_proc[n, c, pos_int:end_pos] += spike
        else:
            raise ValueError(f'不支持的粒度: {position_granularity}')

    if return_spike_pos:
        return X_proc.reshape(original_shape), cast(torch.Tensor, positions / srate)
    return X_proc.reshape(original_shape)


def bandstop_aug(
    X: torch.Tensor,
    srate: int,
    bandwidth: float = 2,
    freqs_to_notch: Optional[int | typing.Sequence[float]] = None,
    random_state: Optional[int] = None,
    random_freq_range: Optional[Tuple[float, float]] = None,
    axis=-1,
) -> torch.Tensor:
    """
    带阻滤波器数据增强

    使用陷波滤波器（notch filter）对信号中的特定频率进行抑制，常用于去除工频干扰等。
    对每个目标频率，带阻宽度保持恒定。

    Parameters
    ----------
    - `X` : `torch.Tensor`
        - 输入数据，形状为 (C, T) 或 (N, C, T)
    - `srate` : `int`
        - 采样率（Hz）
    - `bandwidth` : `float`, optional
        - 带阻宽度（Hz），by default 2
        - 定义为每个频率点两侧的衰减带宽
    - `freqs_to_notch` : `Optional[list[float] | int]`, optional
        - 需要抑制的频率列表（Hz），by default [50]
        - 可指定多个频率，如 [50, 60] 同时抑制 50Hz 和 60Hz
        - 指定 int, 随机选择频带
    - `random_state` : `Optional[int]`, optional
        - 随机种子，用于生成可重复的随机数
    - `random_freq_range` : `Optional[Tuple[float, float]]`, optional
        - 随机选择频带的频率范围
    - `axis` : `int`, optional
        - 应用滤波的轴，by default -1

    Returns
    ----------
    - `torch.Tensor`
        - 滤波后的数据，形状与输入相同
    """
    if bandwidth <= 0:
        return X

    # 修复可变默认参数
    if freqs_to_notch is None:
        notch_freqs: typing.Sequence[float] = [50.0]
    elif isinstance(freqs_to_notch, int):
        freq_range = random_freq_range or (0.5, srate / 2 - 1)  # 浮点数上限

        # PyTorch 随机数生成
        generator = _use_random_state(random_state, device=X.device)
        sampled_freqs = torch.rand(freqs_to_notch, generator=generator, device=X.device)
        sampled_freqs = sampled_freqs * (freq_range[1] - freq_range[0]) + freq_range[0]
        sampled_freqs = torch.clamp(sampled_freqs, 0.1, srate / 2 - 0.1)
        notch_freqs = cast(list[float], sampled_freqs.cpu().numpy().tolist())
    else:
        notch_freqs = freqs_to_notch

    # PyTorch + NumPy 混合模式处理
    device = X.device
    dtype = X.dtype

    # 转换为 NumPy 进行滤波，确保连续内存布局
    X_np = X.detach().cpu().numpy()
    if not X_np.flags['C_CONTIGUOUS']:
        X_np = X_np.copy()

    # 对每个频率点分别应用陷波滤波
    for freq in notch_freqs:
        # 为每个频率正确计算 Q 值: Q = f0 / bandwidth
        Q = freq / bandwidth
        X_filtered = _apply_notch(X_np, srate, f0=freq, Q=Q, axis=axis)
        # 确保结果是连续的
        if not X_filtered.flags['C_CONTIGUOUS']:
            X_filtered = X_filtered.copy()
        X_np = X_filtered

    # 转换回 PyTorch
    return torch.from_numpy(X_np).to(dtype=dtype, device=device)


def _construct_symmetric_phase(
    positive_freq_phases: torch.Tensor, n_samples: int
) -> torch.Tensor:
    """
    根据给定的正频率相位，构建一个完整的共轭对称相位数组。

    Parameters
    ----------
    - positive_freq_phases : torch.Tensor
        一个复数数组，代表正频率部分的相位 (例如，由 j*φ 构成)。
        形状应为 (batch_size, channels, n_positive_freqs)。
    - n_samples : int
        原始信号的时间点总数。

    Returns
    -------
    - torch.Tensor
        完整的、具有共轭对称性的相位数组。
    """
    device = positive_freq_phases.device
    dtype = positive_freq_phases.dtype

    # 直流分量 (k=0) 的相位为0
    zeros = torch.zeros(
        (positive_freq_phases.shape[0], positive_freq_phases.shape[1], 1),
        dtype=dtype,
        device=device,
    )

    if n_samples % 2 == 0:  # 偶数长度
        # 奈奎斯特频率 (k=N/2) 的相位也为0
        full_phase = torch.cat(
            [
                zeros,
                positive_freq_phases,
                zeros,
                -torch.flip(positive_freq_phases, dims=[-1]),
            ],
            dim=-1,
        )
    else:  # 奇数长度
        full_phase = torch.cat(
            [zeros, positive_freq_phases, -torch.flip(positive_freq_phases, dims=[-1])],
            dim=-1,
        )
    return full_phase


def ft_surrogate(
    X: torch.Tensor,
    phase_noise_magnitude: float = 1.0,
    random_state=None,
    phase: torch.Tensor | typing.Sequence[float] | float | None = None,
) -> torch.Tensor:
    """FT surrogate 数据增强

    通过在频域中添加相位扰动来生成时间序列数据的替代样本，同时保留原始数据的功率谱。
    允许用户手动指定相位，或使用随机相位。

    Parameters
    ----------
    - `X` : `torch.Tensor`
        EEG 输入样本或批次，形状为 (C, T) 或 (N, C, T)
    - `phase_noise_magnitude` : `float`, optional
        相位扰动的全局缩放因子。
        - 当使用随机相位时，它控制扰动范围 [0, `magnitude` * 2 * `pi`]。
        - 当手动指定相位时，它作为乘数应用于给定的相位。默认为 1.0。
    - `random_state` : `int | torch.Generator`, optional
        当`phase`为`None`时，用于生成随机相位的种子。默认为 `None`。
    - `phase` : `float | list | torch.Tensor`, optional
        手动指定的相位（以弧度为单位）。
        - 如果是 `float`，所有频率分量将使用相同的相位。
        - 如果是 `list/tensor`，则为每个正频率分量指定相位。
          其长度必须为 `(n_samples - 1) // 2` (奇数) 或 `n_samples // 2 - 1` (偶数)。
        - 如果为 `None` (默认)，则生成随机相位。

    Returns
    -------
    - `torch.Tensor`
        变换后的输入数据，形状与输入相同。
    """
    device = X.device
    needs_squeeze = X.ndim != 3
    original_shape = X.shape
    X_proc = cast(torch.Tensor, ensure_3dims(X, newaxis='batch'))

    f = torch.fft.fft(X_proc, dim=-1)

    n_samples = X_proc.shape[-1]
    batch_size = X_proc.shape[0]
    base_channels = 1  # 为所有通道生成一个相位模式，确保通道同步性不被破坏

    # 根据信号长度确定需要的相位值的数量
    n_positive_freqs = (n_samples - 1) // 2 if n_samples % 2 != 0 else n_samples // 2 - 1

    # --- 步骤1: 准备正频率部分的相位数据 ---
    if phase is None:
        # 模式1: 生成随机相位
        if not (0 <= phase_noise_magnitude <= 1):
            raise ValueError(
                f'当使用随机相位时, phase_noise_magnitude 必须是 0到1之间的浮点数。得到 {phase_noise_magnitude}。'
            )
        generator = _use_random_state(random_state, device=device)
        random_values = torch.rand(
            (batch_size, base_channels, n_positive_freqs),
            generator=generator,
            device=device,
        )
        positive_phases = phase_noise_magnitude * 2 * torch.pi * random_values * 1j
    else:
        # 模式2: 处理手动指定的相位
        if isinstance(phase, (float, int)):
            phase_values = torch.full((n_positive_freqs,), float(phase), device=device)
        elif isinstance(phase, list):
            phase_values = torch.as_tensor(phase, dtype=torch.float32, device=device)
            if phase_values.size(0) != n_positive_freqs:
                raise ValueError(
                    f"对于长度为 {n_samples} 的信号, 'phase' 数组需要包含 "
                    f'{n_positive_freqs} 个元素, 但收到了 {phase_values.size(0)} 个。'
                )
        elif isinstance(phase, torch.Tensor):
            phase_values = phase.to(device=device)
            if phase_values.size(0) != n_positive_freqs:
                raise ValueError(
                    f"对于长度为 {n_samples} 的信号, 'phase' 张量需要包含 "
                    f'{n_positive_freqs} 个元素, 但收到了 {phase_values.size(0)} 个。'
                )
        else:
            raise TypeError(
                f"'phase' 必须是 float, int, list 或 torch tensor, 但收到 {type(phase)}。"
            )

        # 应用全局缩放并转换为复数形式
        positive_phases = phase_noise_magnitude * phase_values.reshape(1, 1, -1) * 1j

    # --- 步骤2: 调用核心函数构建完整的对称相位 ---
    phase_perturbation = _construct_symmetric_phase(positive_phases, n_samples)

    # 将单通道相位模式复制到所有实际通道
    if X_proc.shape[1] > 1:
        phase_perturbation = phase_perturbation.repeat(1, X_proc.shape[1], 1)

    # 应用相位噪声
    f_shifted = f * torch.exp(phase_perturbation)

    # IFFT
    shifted = torch.fft.ifft(f_shifted, dim=-1)
    transformed_X = shifted.real

    # 恢复原始形状
    if needs_squeeze:
        transformed_X = transformed_X.reshape(original_shape)

    return transformed_X
