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


def build_norm(type=None, channels=None):
    assert type is None or type in ['batch', 'instance', 'pixel']
    if type is None:
        return nn.Identity()
    elif type == 'batch':
        if channels is None:
            raise ValueError('channels is required for batch norm')
        return nn.BatchNorm1d(channels)
    elif type == 'instance':
        if channels is None:
            raise ValueError('channels is required for instance norm')
        return nn.InstanceNorm1d(channels)
    elif type == 'pixel':
        return PixelNorm()


class PixelNorm(nn.Module):
    """Pixel Normalization from

    'PROGRESSIVE GROWING OF GANS FOR IMPROVED QUALITY, STABILITY, AND VARIATION

    沿着通道维度对 Pixel 向量做标准化。如输入为 `[N, 50, 1024]`
    则该方法对 N 个样本中，所有 1024 个点在 50 个通道上组成的 50-dim 向量做归一化
    """

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

    def forward(self, x, epsilon=1e-8):
        # 沿着 Channel 维，对每个 pixel 做一次 norm
        assert len(x.shape) == 3
        return x * (((x**2).mean(dim=1, keepdim=True) + epsilon).rsqrt())


class ConditionalBatchNorm1d(nn.BatchNorm1d):
    """Conditional Batch Normalization"""

    def __init__(
        self,
        num_features,
        eps=1e-05,
        momentum=0.1,
        affine=False,
        track_running_stats=True,
    ):
        super(ConditionalBatchNorm1d, self).__init__(
            num_features, eps, momentum, affine, track_running_stats
        )

    def forward(self, input, weight, bias, **kwargs):
        self._check_input_dim(input)

        exponential_average_factor = 0.0

        if self.training and self.track_running_stats:
            if self.num_batches_tracked is not None:
                self.num_batches_tracked += 1
                if self.momentum is None:  # use cumulative moving average
                    exponential_average_factor = 1.0 / self.num_batches_tracked.item()
                else:  # use exponential moving average
                    exponential_average_factor = self.momentum

        output = F.batch_norm(
            input,
            self.running_mean,
            self.running_var,
            self.weight,
            self.bias,
            self.training or not self.track_running_stats,
            exponential_average_factor,
            self.eps,
        )
        if weight.dim() == 1:
            weight = weight.unsqueeze(0)
        if bias.dim() == 1:
            bias = bias.unsqueeze(0)
        size = output.size()
        weight = weight.unsqueeze(-1).expand(size)
        bias = bias.unsqueeze(-1).expand(size)
        return weight * output + bias


class CategoricalConditionalBatchNorm1d(ConditionalBatchNorm1d):
    def __init__(
        self,
        num_classes,
        num_features,
        eps=1e-5,
        momentum=0.1,
        affine=False,
        track_running_stats=True,
    ):
        super(CategoricalConditionalBatchNorm1d, self).__init__(
            num_features, eps, momentum, affine, track_running_stats
        )
        self.weights = nn.Embedding(num_classes, num_features)
        self.biases = nn.Embedding(num_classes, num_features)

        self._initialize()

    def _initialize(self):
        nn.init.ones_(self.weights.weight.data)
        nn.init.zeros_(self.biases.weight.data)

    def forward(self, input, c, **kwargs):
        weight = self.weights(c)
        bias = self.biases(c)
        return super(CategoricalConditionalBatchNorm1d, self).forward(input, weight, bias)


class RMSNorm(nn.Module):
    """Root Mean Square Layer Normalization

    RMS 归一化是一种简化的层归一化，只进行缩放而不进行平移。
    在现代大语言模型（如 LLaMA、T5）中广泛使用，计算效率更高。

    References
    ----------
    - **Paper**: Zhang, B., & Sennrich, R. (2019). "Root Mean Square Layer Normalization."
      *NeurIPS 2019*. https://arxiv.org/abs/1910.07467
    - **LLaMA**: Touvron, H., et al. (2023). "LLaMA: Open and Efficient Foundation Language Models."
      https://arxiv.org/abs/2302.13971
    - **T5**: Raffel, C., et al. (2019). "Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer."
      https://arxiv.org/abs/1910.10683

    Examples
    ----------
    ```py
    # 用于序列数据 [N, C, T]
    rms_norm = RMSNorm(64)  # 64 个特征维度
    x = torch.randn(32, 64, 128)  # [batch, features, time]
    output = rms_norm(x)

    # 用于 Transformer 中
    class TransformerLayer(nn.Module):
        def __init__(self, dim):
            super().__init__()
            self.norm1 = RMSNorm(dim)
            self.norm2 = RMSNorm(dim)
            self.attn = nn.MultiheadAttention(dim, num_heads=8)
            self.mlp = nn.Sequential(nn.Linear(dim, dim*4), nn.GELU(), nn.Linear(dim*4, dim))

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

    def __init__(self, normalized_shape: int, eps: float = 1e-6):
        """初始化 RMSNorm 层

        Parameters
        ----------
        - `normalized_shape` : `int`
            - 归一化的特征维度大小
        - `eps` : `float`, optional
            - 数值稳定性的小常数, by default 1e-6
        """
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(normalized_shape))

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

        Parameters
        ----------
        - `x` : `torch.Tensor`
            - 输入张量，形状为 [..., normalized_shape]

        Returns
        ----------
        - `torch.Tensor`
            - 归一化后的张量，形状与输入相同
        """
        # 计算 RMS
        rms = torch.sqrt(torch.mean(x**2, dim=-1, keepdim=True) + self.eps)
        # 归一化并缩放
        return x / rms * self.weight

    def extra_repr(self) -> str:
        return f'normalized_shape={self.weight.shape[0]}, eps={self.eps}'


class LayerScale(nn.Module):
    """Layer Scale 层缩放模块

    为每个残差连接添加可学习的缩放参数，有助于训练深层网络。
    在 Vision Transformer 和其他深层架构中广泛使用。

    References
    ----------
    - **Paper**: Touvron, H., et al. (2021). "Going deeper with Image Transformers."
      *ICCV 2021*. https://arxiv.org/abs/2103.17239
    - **CaiT**: Touvron, H., et al. (2021). "Class-Attention in Image Transformers."
      *NeurIPS 2021*. https://arxiv.org/abs/2103.17239
    - **timm Implementation**: https://github.com/rwightman/pytorch-image-models

    Examples
    ----------
    ```py
    # 基本使用
    layer_scale = LayerScale(dim=64, init_values=1e-4)
    x = torch.randn(32, 64, 128)
    scaled_x = layer_scale(x)

    # 在 Transformer 块中使用
    class TransformerBlock(nn.Module):
        def __init__(self, dim, init_values=1e-4):
            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.ls1 = LayerScale(dim, init_values)
            self.ls2 = LayerScale(dim, init_values)

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

    def __init__(self, dim: int, init_values: float = 1e-4):
        """初始化 LayerScale 层

        Parameters
        ----------
        - `dim` : `int`
            - 特征维度大小
        - `init_values` : `float`, optional
            - 缩放参数的初始值, by default 1e-4
        """
        super().__init__()
        self.gamma = nn.Parameter(init_values * torch.ones(dim))

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

        Parameters
        ----------
        - `x` : `torch.Tensor`
            - 输入张量，最后一维应为 dim

        Returns
        ----------
        - `torch.Tensor`
            - 缩放后的张量，形状与输入相同
        """
        return x * self.gamma

    def extra_repr(self) -> str:
        return f'dim={self.gamma.shape[0]}, init_values={self.gamma.data[0].item():.2e}'
