import typing
import warnings
from typing import Any, Callable, Optional, Tuple, Union, cast

import torch
import torch.nn as nn

from ._base import BaseTransform
from .signal_aug_transforms import __all__ as sig_aug_transforms_all

__all__ = sig_aug_transforms_all + [
    'ComposedTransform',
    'RandSelectiveTransform',
    'RandSignalMaskTransform',
]


def __getattr__(name: str):
    import importlib

    if name in sig_aug_transforms_all:
        module = importlib.import_module('.signal_aug_transforms', __name__)
        return getattr(module, name)
    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")


class ComposedTransform(BaseTransform):
    """组合多个变换的容器

    允许将多个变换按顺序应用于输入数据。

    在没有安装 torchvision 时，提供一个简单的接口来组合变换。
    """

    def __init__(self, transforms: typing.Sequence):
        """
        Parameters
        ----------
        - `transforms` : `list`
            - 一个由变换对象组成的列表。
        """
        super().__init__()
        # 使用 nn.ModuleList 确保变换被正确注册为模块
        self.transforms = nn.ModuleList(transforms)

    def forward(
        self, X: torch.Tensor, y: Optional[torch.Tensor] = None
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
        """
        按顺序应用所有变换。

        Parameters
        ----------
        - `X` : `torch.Tensor`
            - 输入数据。
        - `y` : `Optional[torch.Tensor]`, optional
            - 对应的标签张量, by default None.

        Returns
        -------
        - `Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]`
            - 如果调用时 `y` 为 None，只返回变换后的 `X`。
            - 如果调用时 `y` 不为 None，返回一个元组 `(transformed_X, transformed_y)`。
        """
        y_was_provided = y is not None

        for t in self.transforms:
            # 将 X 和 y (即使是 None) 传递给每个变换
            result = t(X, y)

            # 根据返回值的类型更新 X 和 y
            if isinstance(result, tuple):
                # 假设返回 (X, y)
                X, y = result
            else:
                # 如果只返回一个张量，则认为是 X
                X = result

        if y_was_provided:
            return cast(Tuple[torch.Tensor, torch.Tensor], (X, y))
        else:
            return X


class RandSelectiveTransform(BaseTransform):
    """随机选择性变换

    在 forward 过程中，以指定的概率 `prob` 随机选择一部分样本施加变换，
    其他样本保持不变。此变换可以同时作用于数据 `X` 和标签 `y`。

    Example
    ----------
    ```python
    # 示例1: 只对 X 应用变换
    noise_transform = AddWhiteNoise(std=0.1)
    selective_noise = RandSelectiveTransform(transform=noise_transform, prob=0.5)
    augmented_X = selective_noise(batch_X)

    # 示例2: 同时对 X 和 y 应用变换 (假设有一个会影响标签的变换)
    # class Mixup(BaseTransform):
    #     def forward(self, X, y):
    #         # ... Mixup 逻辑 ...
    #         return mixed_X, mixed_y
    #
    # selective_mixup = RandSelectiveTransform(transform=Mixup(), prob=1.0)
    # augmented_X, augmented_y = selective_mixup(batch_X, batch_y)
    ```
    """

    def __init__(
        self,
        transform: Callable | BaseTransform,
        prob: float = 0.4,
        random_state: Optional[int] = None,
    ):
        """随机选择性变换

        Parameters
        ----------
        - `transform` : `Callable | BaseTransform`
            - 要应用的变换函数或变换对象。
            - 如果只变换 `X`，该函数应接收 `X` 并返回 `transformed_X`。
            - 如果同时变换 `X` 和 `y`，该函数应接收 `X` 和 `y`，并返回一个元组 `(transformed_X, transformed_y)`。
        - `prob` : `float`, optional
            - 每个样本被选中进行变换的概率，by default 0.4
        - `random_state` : `Optional[int]`, optional
            - 随机种子，用于可重复的随机选择，by default None
        """
        super().__init__()
        self.transform = transform
        self.prob = prob
        self.batch_axis = 0
        self.random_state = random_state

        if random_state is not None:
            self._rng = torch.Generator()
            self._rng.manual_seed(random_state)
        else:
            self._rng = None

    def forward(
        self, X: torch.Tensor, y: Optional[torch.Tensor] = None
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
        """
        应用随机选择性变换。

        Parameters
        ----------
        - `X` : `torch.Tensor`
            - 输入数据，形状通常为 (N, C, T) 或其他批次格式。
        - `y` : `Optional[torch.Tensor]`, optional
            - 对应的标签张量，形状通常为 (N,) 或 (N, ...)。
            - 如果提供，变换也将尝试应用于标签。by default None.

        Returns
        -------
        - `Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]`
            - 如果 `y` 为 None，只返回变换后的 `X`。
            - 如果 `y` 不为 None，返回一个元组 `(transformed_X, transformed_y)`。
        """
        # 如果变换概率为0，或输入是标量，则直接返回
        if self.prob == 0 or X.dim() == 0:
            return (X, y) if y is not None else X

        # 检查输入维度是否包含批次轴，若不包含，则按概率决定是否对单个样本变换
        if X.dim() <= self.batch_axis:
            # 为单个样本生成一个随机决策
            apply_transform = torch.rand(1, generator=self._rng).item() < self.prob
            if apply_transform:
                # 准备变换函数的参数
                transform_args = (X, y) if y is not None else (X,)
                return cast(
                    Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
                    self.transform(*transform_args),
                )
            else:
                return (X, y) if y is not None else X

        # --- 以下为处理批次数据的核心逻辑 ---
        batch_size = X.shape[self.batch_axis]

        # 生成随机掩码，决定哪些样本需要变换
        if self._rng is not None:
            mask = torch.rand(batch_size, generator=self._rng) < self.prob
        else:
            mask = torch.rand(batch_size, device=X.device) < self.prob

        # 如果没有样本被选中，直接返回原数据，避免不必要的clone
        if not mask.any():
            return (X, y) if y is not None else X

        # 创建输出张量的副本，以避免原地修改
        output_X = X.clone()
        output_y = y.clone() if y is not None else None

        # 构建用于高级索引的切片
        full_indices: list[Any] = [slice(None)] * X.dim()
        full_indices[self.batch_axis] = mask  # 使用布尔掩码直接索引

        # 提取被选中的数据和标签
        selected_X = X[tuple(full_indices)]
        selected_y = y[mask] if y is not None else None

        # 准备变换函数的参数
        transform_args = (
            (selected_X, selected_y) if selected_y is not None else (selected_X,)
        )

        # 对选中的样本应用变换 (简化了调用逻辑)
        transformed_result = self.transform(*transform_args)

        # 根据是否提供了y，解包变换结果
        if y is not None:
            transformed_X, transformed_y = cast(
                tuple[torch.Tensor, torch.Tensor], transformed_result
            )
            # 将变换后的标签放回输出张量
            assert output_y is not None
            output_y[mask] = transformed_y
        else:
            transformed_X = cast(torch.Tensor, transformed_result)

        # 将变换后的样本放回输出张量
        output_X[tuple(full_indices)] = transformed_X

        # 根据输入时是否提供了y，决定返回内容
        if y is not None:
            assert output_y is not None
            return output_X, output_y
        return output_X


class RandSignalMaskTransform(BaseTransform):
    """随机信号掩码变换

    在脑电信号上时间序列维度上随机选择一个时间段，对该时间段应用指定的变换。
    支持固定或随机的时间段选择策略。

    支持的输入格式：
    - `[batch, time]`: 视为 `[batch, 1, time]`
    - `[batch, channel, time]`: 多通道信号

    Example
    ----------
    ```python
    # 示例1: 固定时间段变换
    noise_transform = AddWhiteNoise(std=0.1)
    mask_transform = RandSignalMaskTransform(
        transform=noise_transform,
        srate=100.0,
        fixed_start=2.0,  # 从2秒开始
        fixed_width=1.0   # 持续1秒
    )
    augmented_X = mask_transform(batch_X)

    # 示例2: 随机时间段变换
    mask_transform = RandSignalMaskTransform(
        transform=noise_transform,
        srate=100.0,
        rand_start_range=(0.0, 5.0),   # 起始时间0-5秒
        rand_width_range=(0.5, 2.0),  # 持续时间0.5-2秒
        sync_across_channels=False     # 每个通道独立mask
    )
    augmented_X = mask_transform(batch_X)
    ```
    """

    def __init__(
        self,
        transform: Callable | BaseTransform,
        srate: float,
        fixed_start: Optional[float] = None,
        fixed_width: Optional[float] = None,
        rand_start_range: Optional[Tuple[float, float]] = None,
        rand_width_range: Optional[Tuple[float, float]] = None,
        sync_across_channels: bool = True,
        random_state: Optional[int] = None,
    ):
        """随机信号掩码变换

        Parameters
        ----------
        - `transform` : `Callable | BaseTransform`
            - 要应用的变换函数或变换对象。
        - `srate` : `float`
            - 信号采样率（Hz）。
        - `fixed_start` : `Optional[float]`, optional
            - 固定起始时间（秒），与 `rand_start_range` 互斥, by default None
        - `fixed_width` : `Optional[float]`, optional
            - 固定时间段宽度（秒），与 `rand_width_range` 互斥, by default None
        - `rand_start_range` : `Optional[Tuple[float, float]]`, optional
            - 随机起始时间范围（秒），格式为 (low, high), by default None
        - `rand_width_range` : `Optional[Tuple[float, float]]`, optional
            - 随机时间段宽度范围（秒），格式为 (low, high), by default None
        - `sync_across_channels` : `bool`, optional
            - 是否所有通道使用相同的时间段mask, by default True
        - `random_state` : `Optional[int]`, optional
            - 随机种子，用于可重复的随机选择, by default None
        """
        super().__init__()
        self.transform = transform
        self.srate = srate
        self.fixed_start = fixed_start
        self.fixed_width = fixed_width
        self.rand_start_range = rand_start_range
        self.rand_width_range = rand_width_range
        self.sync_across_channels = sync_across_channels
        self.random_state = random_state

        # 参数验证
        self._validate_params()

        # 设置随机数生成器
        if random_state is not None:
            self._rng = torch.Generator()
            self._rng.manual_seed(random_state)
        else:
            self._rng = None

    def _validate_params(self):
        """验证参数的合法性"""
        # 检查 fixed_* 和 rand_*_range 参数的互斥性
        if self.fixed_start is not None and self.rand_start_range is not None:
            raise ValueError('fixed_start 和 rand_start_range 参数不能同时设置')

        if self.fixed_width is not None and self.rand_width_range is not None:
            raise ValueError('fixed_width 和 rand_width_range 参数不能同时设置')

        # 检查范围参数的合法性
        if self.rand_start_range is not None:
            low, high = self.rand_start_range
            if low < 0 or high <= low:
                raise ValueError('rand_start_range 必须满足 0 <= low < high')

        if self.rand_width_range is not None:
            low, high = self.rand_width_range
            if low <= 0 or high <= low:
                raise ValueError('rand_width_range 必须满足 0 < low < high')

        # 检查固定参数的合法性
        if self.fixed_start is not None and self.fixed_start < 0:
            raise ValueError('fixed_start 必须 >= 0')

        if self.fixed_width is not None and self.fixed_width <= 0:
            raise ValueError('fixed_width 必须 > 0')

    def _determine_time_axis(self, X: torch.Tensor) -> int:
        """确定时间轴的索引"""
        if X.dim() == 2:
            return 1  # [batch, time]
        elif X.dim() == 3:
            return 2  # [batch, channel, time]
        else:
            raise ValueError(f'不支持的输入维度: {X.dim()}，期望 2D 或 3D')

    def _select_time_segment(self, signal_length: int) -> Tuple[int, int]:
        """选择时间段，返回 (start_idx, end_idx)"""
        # 确定宽度
        if self.fixed_width is not None:
            width_samples = int(self.fixed_width * self.srate)
        else:
            if self.rand_width_range is not None:
                low, high = self.rand_width_range
                if self._rng is not None:
                    width_seconds = (
                        torch.rand(1, generator=self._rng).item() * (high - low) + low
                    )
                else:
                    width_seconds = torch.rand(1).item() * (high - low) + low
            else:
                # 默认随机宽度：10% - 50% 的信号长度
                min_width = max(1, signal_length // 10)
                max_width = signal_length // 2
                if self._rng is not None:
                    width_samples = int(
                        torch.randint(
                            min_width, max_width + 1, (1,), generator=self._rng
                        ).item()
                    )
                else:
                    width_samples = int(
                        torch.randint(min_width, max_width + 1, (1,)).item()
                    )
                width_samples = min(width_samples, signal_length)
                return self._select_start_position(signal_length, width_samples)

            width_samples = int(width_seconds * self.srate)

        # 确保宽度不超过信号长度
        width_samples = min(width_samples, signal_length)

        return self._select_start_position(signal_length, width_samples)

    def _select_start_position(
        self, signal_length: int, width_samples: int
    ) -> Tuple[int, int]:
        """根据宽度选择起始位置"""
        # 确定起始位置
        if self.fixed_start is not None:
            start_samples = int(self.fixed_start * self.srate)
        else:
            if self.rand_start_range is not None:
                low, high = self.rand_start_range
                low_samples = int(low * self.srate)
                high_samples = int(high * self.srate)
            else:
                # 默认随机起始：整个信号范围
                low_samples = 0
                high_samples = signal_length

            # 确保起始位置加上宽度不超过信号长度
            max_start = signal_length - width_samples
            high_samples = min(high_samples, max_start)

            if low_samples > high_samples:
                # 截断策略：调整宽度
                width_samples = signal_length - low_samples
                if width_samples <= 0:
                    width_samples = min(signal_length, int(0.1 * signal_length))
                    low_samples = 0
                    high_samples = signal_length - width_samples
                warnings.warn(
                    f'时间段参数冲突，已调整宽度为 {width_samples/self.srate:.3f} 秒'
                )

            if self._rng is not None:
                start_samples = int(
                    torch.randint(
                        low_samples, high_samples + 1, (1,), generator=self._rng
                    ).item()
                )
            else:
                start_samples = int(
                    torch.randint(low_samples, high_samples + 1, (1,)).item()
                )

        # 最终边界检查和截断
        start_samples = max(0, min(start_samples, signal_length - 1))
        end_samples = min(start_samples + width_samples, signal_length)

        if end_samples - start_samples < width_samples:
            warnings.warn(
                f'时间段被截断：从 {width_samples/self.srate:.3f} 秒调整为 {(end_samples - start_samples)/self.srate:.3f} 秒'
            )

        return start_samples, end_samples

    def forward(
        self, X: torch.Tensor, *args
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]:
        """
        应用随机信号掩码变换。

        Parameters
        ----------
        - `X` : `torch.Tensor`
            - 输入信号，形状为 `[batch, time]` 或 `[batch, channel, time]`。
        - `*args`
            - 其他参数，将原路返回。

        Returns
        -------
        - `Union[torch.Tensor, Tuple[torch.Tensor, ...]]`
            - 如果没有额外参数，返回变换后的 `X`。
            - 如果有额外参数，返回 `(transformed_X, *args)`。
        """
        time_axis = self._determine_time_axis(X)
        signal_length = X.shape[time_axis]

        # 如果是2D，转换为3D处理
        if X.dim() == 2:
            X = X.unsqueeze(1)  # [batch, 1, time]
            time_axis = 2
            was_2d = True
        else:
            was_2d = False

        n_channels = X.shape[1]
        # 创建输出张量
        output_X = X.clone()

        if self.sync_across_channels:
            # 所有batch和通道使用相同的时间段
            start_idx, end_idx = self._select_time_segment(signal_length)
            # 提取时间段 [batch, channel, time_segment]
            segment = X[:, :, start_idx:end_idx]
            # 应用变换
            transformed_segment = cast(torch.Tensor, self.transform(segment))
            # 放回原位置
            output_X[:, :, start_idx:end_idx] = transformed_segment
        else:
            # 每个通道独立选择时间段，但同一通道在所有batch中使用相同时间段
            for c in range(n_channels):
                start_idx, end_idx = self._select_time_segment(signal_length)
                # 提取时间段 [batch, 1, time_segment]
                segment = X[:, c : c + 1, start_idx:end_idx]
                # 应用变换
                transformed_segment = cast(torch.Tensor, self.transform(segment))
                # 放回原位置
                output_X[:, c : c + 1, start_idx:end_idx] = transformed_segment

        # 如果原来是2D，转换回2D
        if was_2d:
            output_X = output_X.squeeze(1)

        # 根据是否有额外参数决定返回格式
        if args:
            return (output_X,) + args
        else:
            return output_X
