import torch
from typing import Optional
from .config import TransformerConfig


class MaskedGroupNorm(torch.nn.GroupNorm):
    def forward(
        self, x: torch.Tensor, padding_mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        if padding_mask is None:
            return super().forward(x)
        B, C, L = x.shape
        G = self.num_groups
        x_grouped = x.view(B, G, C // G, L)
        padding_mask_grouped = padding_mask.reshape(B, G, C // G, L).bool()
        mean = torch.masked.mean(
            x_grouped, mask=padding_mask_grouped, dim=(2, 3), keepdim=True
        )
        var = torch.masked.var(
            x_grouped,
            mask=padding_mask_grouped,
            dim=(2, 3),
            keepdim=True,
            unbiased=False,
        )
        x_norm = (x_grouped - mean) / torch.sqrt(var + self.eps)
        x_norm = x_norm.view(B, C, L)
        if self.affine:
            x_norm = x_norm * self.weight.view(1, -1, 1) + self.bias.view(1, -1, 1)
        return x_norm * padding_mask


class ConvBlock1d(torch.nn.Module):
    def __init__(self, config: TransformerConfig):
        super().__init__()
        self.groupnorm = MaskedGroupNorm(num_groups=1, num_channels=config.hidden_size)
        self.activation = torch.nn.SiLU()
        self.project = torch.nn.Conv1d(
            in_channels=config.hidden_size,
            out_channels=config.hidden_size,
            kernel_size=3,
            padding="same",
        )

    def forward(
        self, x: torch.Tensor, padding_mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        x = self.groupnorm(x, padding_mask=padding_mask)
        x = self.activation(x)
        return self.project(x)


class ResnetBlock1d(torch.nn.Module):
    def __init__(self, config: TransformerConfig):
        super().__init__()
        self.block1 = ConvBlock1d(config)
        self.block2 = ConvBlock1d(config)

    def forward(
        self, x: torch.Tensor, padding_mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        if padding_mask is not None:
            padding_mask = padding_mask.unsqueeze(1).expand_as(x)
        h = self.block1(x, padding_mask=padding_mask)
        h = self.block2(h, padding_mask=padding_mask)
        return h + x