import typing

# PyTorch-compatible transforms for EEG signal augmentation
from abc import ABC, abstractmethod

import torch
import torch.nn as nn

__all__ = ['BaseTransform', 'AugmentationTransformer']


class BaseTransform(nn.Module, ABC):
    """EEG信号变换的基类

    继承自nn.Module以获得设备管理、状态跟踪等特性，
    同时兼容torchvision transforms生态系统。
    """

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

    @abstractmethod
    def forward(self, X: torch.Tensor, *args) -> torch.Tensor:
        """前向传播方法

        Parameters
        ----------
        - `X` : `torch.Tensor`
            - 输入EEG数据，形状为 (C, T) 或 (N, C, T)
        - `*args` :
            - 额外的参数，用于适配其他输入格式

        Returns
        -------
        - `torch.Tensor`
            - 变换后的数据，形状与输入相同
        """
        pass


class AugmentationTransformer(BaseTransform):
    """增广变换基类

    为所有数据增广变换提供统一的参数管理和状态跟踪功能。
    自动将传入的参数注册为模块缓冲区，确保在设备迁移时保持一致性。
    """

    def __init__(self, **params):
        """初始化增广变换

        Parameters
        ----------
        - `**params` :
            - 变换的配置参数，将自动注册为模块缓冲区
        """
        super().__init__()
        # 添加为 Module 状态
        for key, value in params.items():
            if isinstance(value, torch.Tensor):
                self.register_buffer(key, value)
            else:
                # 对于非Tensor类型，直接设置为属性
                setattr(self, key, value)
        self._hparams = list(params.keys())

        self._hparams_range = {}
        self._initial_hpmarams = params.copy()

    def define_hparam_range(self, **param_range):
        """设置超参数范围，方便后面实现 Augment 算法的时候通过超参数 magnitude 来控制变换强度。

        Parameters
        ----------
        - `**param_range` :
            - 超参数范围定义，格式为 param_name=(min_val, max_val)

        Raises
        ------
        - `ValueError`
            - 当参数不是该增广算法的有效超参数时
        - `ValueError`
            - 当参数范围不是两个值的元组或列表时
        - `ValueError`
            - 当参数当前值不是数字类型时
        - `ValueError`
            - 当参数当前值超出指定范围时
        """
        for key, value in param_range.items():
            if key not in self._hparams:
                raise ValueError(
                    f"'{key}' 并非增广算法 {self.__class__.__name__} 的超参数."
                )
            if not isinstance(value, (tuple, list)) or len(value) != 2:
                raise ValueError(f"超参数 '{key}' 的范围必须是两个值的元组或列表.")
            cur_val = getattr(self, key, None)
            if cur_val is not None and not isinstance(cur_val, (int, float)):
                raise ValueError(
                    f"超参数 '{key}' 的当前值 ({cur_val}) 不是数字，无法设置."
                )
            if not (value[0] <= cur_val <= value[1]):
                raise ValueError(
                    f"超参数 '{key}' 的当前值 ({cur_val}) 超出了指定范围 {value}."
                )
            self._hparams_range[key] = tuple(value)

    def _translate_magnitude_to_hparams(self, magnitude: float = 1.0):
        """将magnitude转换为具体的超参数值

        Parameters
        ----------
        - `magnitude` : `float`, optional
            - 控制变换强度的系数，范围在 0.0 到 1.0 之间，by default 1.0
            - 0.0 表示最小变换强度，1.0 表示最大变换强度

        Returns
        -------
        - `dict`
            - 根据magnitude计算出的超参数值字典

        Raises
        ------
        - `ValueError`
            - 当magnitude不在[0.0, 1.0]范围内时
        - `RuntimeError`
            - 当没有定义参数范围时
        """
        if not (0.0 <= magnitude <= 1.0):
            raise ValueError(f'magnitude必须在[0.0, 1.0]范围内，得到: {magnitude}')

        if not self._hparams_range:
            raise RuntimeError(
                f'增广算法 {self.__class__.__name__} 尚未定义参数范围，请先调用 define_hparam_range()'
            )

        new_params = {}
        for param_name, (min_val, max_val) in self._hparams_range.items():
            # 线性插值计算新参数值
            new_val = min_val + magnitude * (max_val - min_val)
            new_params[param_name] = new_val

        return new_params

    def update_hparams(
        self, magnitude: float | None = None, new_params: typing.Mapping | None = None
    ):
        """更新当前实例的超参数

        Parameters
        ----------
        - `new_params` :
            - 要更新的超参数及其新值
        - `magnitude` : `float`, optional
            - 控制变换强度的系数，范围在 0.0 到 1.0 之间，by default None
        Raises
        ------
        - `ValueError`
            - 当尝试更新不存在的参数时
        """
        if new_params is None and magnitude is None:
            raise ValueError('必须提供 new_params 或 magnitude 之一来更新超参数.')
        if new_params is not None and magnitude is not None:
            raise ValueError('只能提供 new_params 或 magnitude 之一，不能同时提供.')

        if magnitude is not None:
            new_params = self._translate_magnitude_to_hparams(magnitude)

        typed_new_params = typing.cast(typing.Mapping, new_params)
        for key, value in typed_new_params.items():
            if key not in self._hparams:
                raise ValueError(f"'{key}' 不是 {self.__class__.__name__} 的有效超参数")

            # 更新参数值
            if isinstance(getattr(self, key), torch.Tensor):
                # 如果原参数是Tensor，需要重新注册buffer
                if isinstance(value, torch.Tensor):
                    self.register_buffer(key, value)
                else:
                    self.register_buffer(key, torch.tensor(value, dtype=torch.float32))
            else:
                setattr(self, key, value)

    @property
    def augment_params(self):
        return {k: getattr(self, k) for k in self._hparams}
