"""正则化模块

实现现代深度学习中常用的正则化技术，包括随机深度等。
"""

from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F

from .base import BaseModule


class DropPath(BaseModule):
    """随机深度 (Stochastic Depth) 正则化层

    在训练时随机跳过某些层的计算，在推理时保持所有层激活。
    这种技术可以减少过拟合，加速训练，并提高模型泛化能力。

    References
    ----------
    - **Paper**: Huang, G., Sun, Y., Liu, Z., Sedra, D., & Weinberger, K. Q. (2016).
      "Deep Networks with Stochastic Depth." *ECCV 2016*.
      https://arxiv.org/abs/1603.09382
    - **Vision Transformer**: Dosovitskiy, A., et al. (2020). "An Image is Worth 16x16 Words:
      Transformers for Image Recognition at Scale." https://arxiv.org/abs/2010.11929
    - **timm Implementation**: https://github.com/rwightman/pytorch-image-models

    Examples
    ----------
    ```py
    # 基本使用
    drop_path = DropPath(drop_prob=0.1)
    x = torch.randn(32, 64, 128)  # [N, C, T] 或任意形状
    output = drop_path(x)

    # 在 Transformer 块中使用
    class TransformerBlock(nn.Module):
        def __init__(self, dim, drop_path_prob=0.1):
            super().__init__()
            self.attn = nn.MultiheadAttention(dim, num_heads=8)
            self.mlp = nn.Sequential(nn.Linear(dim, dim*4), nn.GELU(), nn.Linear(dim*4, dim))
            self.drop_path = DropPath(drop_path_prob)

        def forward(self, x):
            x = x + self.drop_path(self.attn(x, x, x)[0])
            x = x + self.drop_path(self.mlp(x))
            return x
    ```
    """

    def __init__(self, drop_prob: float = 0.0):
        """初始化 DropPath 层

        Parameters
        ----------
        - `drop_prob` : `float`, optional
            - 随机跳过的概率，范围 [0, 1], by default 0.0
        """
        super().__init__()
        self.drop_prob = drop_prob

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """前向传播

        Parameters
        ----------
        - `x` : `torch.Tensor`
            - 输入张量，任意形状

        Returns
        ----------
        - `torch.Tensor`
            - 输出张量，形状与输入相同
        """
        if self.drop_prob == 0.0 or not self.training:
            return x

        keep_prob = 1 - self.drop_prob
        shape = (x.shape[0],) + (1,) * (
            x.ndim - 1
        )  # work with diff dim tensors, not just 2D ConvNets
        random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
        random_tensor.floor_()  # binarize
        output = x.div(keep_prob) * random_tensor
        return output

    def extra_repr(self) -> str:
        return f'drop_prob={self.drop_prob}'


class StochasticDepth(BaseModule):
    """随机深度的另一种实现

    与 DropPath 类似，但使用不同的实现策略。
    这个版本更接近原始论文的描述。

    References
    ----------
    - **Paper**: Huang, G., Sun, Y., Liu, Z., Sedra, D., & Weinberger, K. Q. (2016).
      "Deep Networks with Stochastic Depth." *ECCV 2016*.
      https://arxiv.org/abs/1603.09382

    Examples
    ----------
    ```py
    # 基本使用
    stoch_depth = StochasticDepth(survival_prob=0.9)
    x = torch.randn(32, 64, 128)
    output = stoch_depth(x)

    # 在残差块中使用
    class ResidualBlock(nn.Module):
        def __init__(self, dim, survival_prob=0.9):
            super().__init__()
            self.conv1 = nn.Conv1d(dim, dim, 3, padding=1)
            self.conv2 = nn.Conv1d(dim, dim, 3, padding=1)
            self.stoch_depth = StochasticDepth(survival_prob)

        def forward(self, x):
            residual = x
            out = F.relu(self.conv1(x))
            out = self.conv2(out)
            out = self.stoch_depth(out)
            return out + residual
    ```
    """

    def __init__(self, survival_prob: float = 0.9):
        """初始化 StochasticDepth 层

        Parameters
        ----------
        - `survival_prob` : `float`, optional
            - 层存活概率，范围 [0, 1], by default 0.9
        """
        super().__init__()
        self.survival_prob = survival_prob

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """前向传播

        Parameters
        ----------
        - `x` : `torch.Tensor`
            - 输入张量，任意形状

        Returns
        ----------
        - `torch.Tensor`
            - 输出张量，形状与输入相同
        """
        if not self.training or self.survival_prob == 1.0:
            return x

        # 生成伯努利随机数
        batch_size = x.shape[0]
        random_tensor = torch.rand(batch_size, 1, 1, device=x.device, dtype=x.dtype)
        binary_tensor = (random_tensor < self.survival_prob).float()

        # 在训练时进行缩放以保持期望值不变
        output = x * binary_tensor / self.survival_prob
        return output

    def extra_repr(self) -> str:
        return f'survival_prob={self.survival_prob}'
