import typing
from typing import Literal, Optional, Tuple, Union, cast

import torch

from ._base import AugmentationTransformer
from .signal_augment import (
    add_spike,
    add_white_noise,
    bandstop_aug,
    ft_surrogate,
    scale_amplitude,
    sign_flip,
    time_reverse,
)

__all__ = [
    'SignFlip',
    'TimeReverse',
    'AddWhiteNoise',
    'ScaleAmplitude',
    'AddSpike',
    'BandstopAug',
    'FTSurrogate',
]


# ========================= Augmentation ========================= #


class SignFlip(AugmentationTransformer):
    """信号上下翻转变换

    将信号进行垂直翻转（乘以-1）。这是一个简单但有效的数据增强方法，

    Example
    -------
    ```python
    transform = SignFlip()
    augmented_data = transform(eeg_data)
    ```
    """

    def __init__(self):
        super().__init__()

    def forward(self, X: torch.Tensor, *args) -> torch.Tensor:
        """应用信号翻转

        Parameters
        ----------
        - `X` : `torch.Tensor`
            - 输入数据，形状为 (C, T) 或 (N, C, T)
        - `*args` :
            - 额外参数（此transform不使用）

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


class TimeReverse(AugmentationTransformer):
    """时间反转变换

    将信号在时间维度上反转。这种变换可以帮助模型学习到
    时间顺序不敏感的特征表示。

    Example
    -------
    ```python
    transform = TimeReverse()
    augmented_data = transform(eeg_data)
    ```
    """

    def __init__(self):
        super().__init__()

    def forward(self, X: torch.Tensor, *args) -> torch.Tensor:
        """应用时间反转

        Parameters
        ----------
        - `X` : `torch.Tensor`
            - 输入数据，形状为 (C, T) 或 (N, C, T)
        - `*args` :
            - 额外参数（此transform不使用）

        Returns
        -------
        - `torch.Tensor`
            - 时间反转后的数据，形状与输入相同
        """
        return time_reverse(X)


class AddWhiteNoise(AugmentationTransformer):
    """添加白噪声变换

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

    Example
    -------
    ```python
    # 绝对噪声
    transform = AddWhiteNoise(std=0.1, rel_power=False)
    # 相对噪声
    transform = AddWhiteNoise(std=0.1, rel_power=True)
    augmented_data = transform(eeg_data)
    ```
    """

    def __init__(
        self,
        std: float = 0.1,
        rel_power: bool = False,
        random_state: Optional[int] = None,
    ):
        """初始化添加白噪声变换

        Parameters
        ----------
        - `std` : `float`, optional
            - 噪声的标准差或相对功率，by default 0.1
        - `rel_power` : `bool`, optional
            - 是否使用相对功率模式，by default False
        - `random_state` : `Optional[int]`, optional
            - 随机种子，用于生成可重复的噪声，by default None
        """
        super().__init__(std=std, rel_power=rel_power, random_state=random_state)

    def forward(self, X: torch.Tensor, *args) -> torch.Tensor:
        """添加白噪声

        Parameters
        ----------
        - `X` : `torch.Tensor`
            - 输入数据，形状为 (C, T) 或 (N, C, T)
        - `*args` :
            - 额外参数（此transform不使用）

        Returns
        -------
        - `torch.Tensor`
            - 添加噪声后的数据，形状与输入相同
        """
        return add_white_noise(
            X,
            scale=cast(float, self.std),
            rel_power=cast(bool, self.rel_power),
            random_state=cast(Optional[int], self.random_state),
        )


class ScaleAmplitude(AugmentationTransformer):
    """振幅缩放变换

    随机缩放信号的振幅，支持不同粒度的缩放策略。

    Example
    -------
    ```python
    transform = ScaleAmplitude(
        scale_range=(0.8, 1.2),
        granularity='per_channel'
    )
    augmented_data = transform(eeg_data)
    ```
    """

    def __init__(
        self,
        scale_range: Tuple[float, float] = (0.8, 1.2),
        granularity: Literal['global', 'per_epoch', 'per_channel'] = 'per_channel',
        random_state: Optional[int] = None,
    ):
        """初始化振幅缩放变换

        Parameters
        ----------
        - `scale_range` : `Tuple[float, float]`, optional
            - 缩放因子的随机采样范围，by default (0.8, 1.2)
        - `granularity` : `Literal['global', 'per_epoch', 'per_channel']`, optional
            - 缩放因子的作用粒度，by default 'per_channel'
        - `random_state` : `Optional[int]`, optional
            - 随机种子，用于生成可复现的缩放因子，by default None
        """
        super().__init__(
            scale_range=scale_range, granularity=granularity, random_state=random_state
        )

    def forward(self, X: torch.Tensor, *args) -> torch.Tensor:
        """应用振幅缩放

        Parameters
        ----------
        - `X` : `torch.Tensor`
            - 输入数据，形状为 (C, T) 或 (N, C, T)
        - `*args` :
            - 额外参数（此transform不使用）

        Returns
        -------
        - `torch.Tensor`
            - 振幅缩放后的数据，形状与输入相同
        """
        return scale_amplitude(
            X,
            scale_range=cast(Tuple[float, float], self.scale_range),
            granularity=cast(
                Literal['global', 'per_epoch', 'per_channel'], self.granularity
            ),
            random_state=cast(Optional[int], self.random_state),
        )


class AddSpike(AugmentationTransformer):
    """添加尖峰噪声变换

    在信号中添加尖峰噪声，模拟EEG中的伪迹现象。

    Example
    -------
    ```python
    transform = AddSpike(
        srate=250,
        number=2,
        t_width=0.01,
        amplitude_mode='random_relative'
    )
    augmented_data = transform(eeg_data)
    ```
    """

    def __init__(
        self,
        srate: int,
        number: int,
        t_width: float,
        amplitude_mode: Literal[
            'fixed', 'random_absolute', 'random_relative'
        ] = 'random_relative',
        amplitude: Union[float, Tuple[float, float]] = (-1.5, 1.5),
        position_granularity: Literal['per_epoch', 'per_channel'] = 'per_epoch',
        random_state: Optional[int] = None,
    ):
        """初始化添加尖峰噪声变换

        Parameters
        ----------
        - `srate` : `int`
        - 采样率（Hz）
        - `number` : `int`
            - 每个粒度单位插入的尖峰数量
        - `t_width` : `float`
            - 单个尖峰的持续时间（秒）
        - `amplitude_mode` : `Literal['fixed', 'random_absolute', 'random_relative']`, optional
            - 振幅生成模式，by default 'random_relative'
        - `amplitude` : `Union[float, Tuple[float, float]]`, optional
            - 振幅参数，by default (-1.5, 1.5)
        - `position_granularity` : `Literal['per_epoch', 'per_channel']`, optional
            - 尖峰位置的生成粒度，by default 'per_epoch'
        - `random_state` : `Optional[int]`, optional
            - 随机种子，by default None
        """
        super().__init__(
            srate=srate,
            number=number,
            t_width=t_width,
            amplitude_mode=amplitude_mode,
            amplitude=amplitude,
            position_granularity=position_granularity,
            random_state=random_state,
        )

    def forward(self, X: torch.Tensor, *args) -> torch.Tensor:
        """添加尖峰噪声

        Parameters
        ----------
        - `X` : `torch.Tensor`
            - 输入数据，形状为 (C, T) 或 (N, C, T)
        - `*args` :
            - 额外参数（此transform不使用）

        Returns
        -------
        - `torch.Tensor`
            - 添加尖峰后的数据，形状与输入相同
        """
        return cast(
            torch.Tensor,
            add_spike(
                X,
                srate=cast(int, self.srate),
                number=cast(int, self.number),
                t_width=cast(float, self.t_width),
                amplitude_mode=cast(
                    Literal['fixed', 'random_absolute', 'random_relative', 'relative'],
                    self.amplitude_mode,
                ),
                amplitude=cast(Union[float, Tuple[float, float]], self.amplitude),
                position_granularity=cast(
                    Literal['per_epoch', 'per_channel'], self.position_granularity
                ),
                random_state=cast(Optional[int], self.random_state),
            ),
        )


class BandstopAug(AugmentationTransformer):
    """带阻滤波器数据增强变换

    使用陷波滤波器对信号中的特定频率进行抑制，
    常用于去除工频干扰等。

    Example
    -------
    ```python
    # 固定频率抑制
    transform = BandstopAug(srate=250, freqs_to_notch=[50, 60])
    # 随机频率抑制
    transform = BandstopAug(srate=250, freqs_to_notch=3,
                           random_freq_range=(1, 100))
    augmented_data = transform(eeg_data)
    ```
    """

    def __init__(
        self,
        srate: int,
        bandwidth: float = 2,
        freqs_to_notch: Optional[Union[int, typing.Sequence[float]]] = None,
        random_state: Optional[int] = None,
        random_freq_range: Optional[Tuple[float, float]] = None,
        axis: int = -1,
    ):
        """初始化带阻滤波器数据增强变换

        Parameters
        ----------
        - `srate` : `int`
            - 采样率（Hz）
        - `bandwidth` : `float`, optional
            - 带阻宽度（Hz），by default 2
        - `freqs_to_notch` : `Optional[Union[int, list[float]]]`, optional
            - 需要抑制的频率列表（Hz），by default None
        - `random_state` : `Optional[int]`, optional
            - 随机种子，by default None
        - `random_freq_range` : `Optional[Tuple[float, float]]`, optional
            - 随机选择频带的频率范围，by default None
        - `axis` : `int`, optional
            - 应用滤波的轴，by default -1
        """
        super().__init__(
            srate=srate,
            bandwidth=bandwidth,
            freqs_to_notch=freqs_to_notch,
            random_state=random_state,
            random_freq_range=random_freq_range,
            axis=axis,
        )

    def forward(self, X: torch.Tensor, *args) -> torch.Tensor:
        """应用带阻滤波

        Parameters
        ----------
        - `X` : `torch.Tensor`
            - 输入数据，形状为 (C, T) 或 (N, C, T)
        - `*args` :
            - 额外参数（此transform不使用）

        Returns
        -------
        - `torch.Tensor`
            - 滤波后的数据，形状与输入相同
        """
        return bandstop_aug(
            X,
            srate=cast(int, self.srate),
            bandwidth=cast(float, self.bandwidth),
            freqs_to_notch=cast(
                Optional[Union[int, typing.Sequence[float]]], self.freqs_to_notch
            ),
            random_state=cast(Optional[int], self.random_state),
            random_freq_range=cast(Optional[Tuple[float, float]], self.random_freq_range),
            axis=cast(int, self.axis),
        )


class FTSurrogate(AugmentationTransformer):
    """FT surrogate 数据增强变换

    通过在频域中添加相位扰动来生成时间序列数据的替代样本，
    同时保留原始数据的功率谱。

    Example
    -------
    ```python
    # 随机相位扰动
    transform = FTSurrogate(phase_noise_magnitude=0.5)
    # 固定相位扰动
    transform = FTSurrogate(phase=0.5)
    augmented_data = transform(eeg_data)
    ```
    """

    def __init__(
        self,
        phase_noise_magnitude: float = 1.0,
        random_state: Optional[int] = None,
        phase: Optional[Union[float, typing.Sequence[float], torch.Tensor]] = None,
    ):
        """初始化FT surrogate数据增强变换

        Parameters
        ----------
        - `phase_noise_magnitude` : `float`, optional
            - 相位扰动的全局缩放因子，by default 1.0
        - `random_state` : `Optional[int]`, optional
            - 随机种子，by default None
        - `phase` : `Optional[Union[float, list[float], torch.Tensor]]`, optional
            - 手动指定的相位，by default None
        """
        super().__init__(
            phase_noise_magnitude=phase_noise_magnitude,
            random_state=random_state,
            phase=phase,
        )

    def forward(self, X: torch.Tensor, *args) -> torch.Tensor:
        """应用FT surrogate变换

        Parameters
        ----------
        - `X` : `torch.Tensor`
            - 输入数据，形状为 (C, T) 或 (N, C, T)
        - `*args` :
            - 额外参数（此transform不使用）

        Returns
        -------
        - `torch.Tensor`
            - 变换后的数据，形状与输入相同
        """
        return ft_surrogate(
            X,
            phase_noise_magnitude=cast(float, self.phase_noise_magnitude),
            random_state=cast(Optional[int], self.random_state),
            phase=cast(
                Optional[Union[float, typing.Sequence[float], torch.Tensor]], self.phase
            ),
        )
