from torch import nn, Tensor

from typing import Union

from .blocks import DepthSeparableConv2d, conv1x1, conv3x3
from .utils import _init_weights


class ConvDownsample(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        norm_layer: Union[nn.BatchNorm2d, nn.GroupNorm, None] = nn.BatchNorm2d,
        activation: nn.Module = nn.ReLU(inplace=True),
        groups: int = 1,
    ) -> None:
        super().__init__()
        assert isinstance(groups, int) and groups > 0, f"Number of groups should be an integer greater than 0, but got {groups}."
        assert in_channels % groups == 0, f"Number of input channels {in_channels} should be divisible by number of groups {groups}."
        assert out_channels % groups == 0, f"Number of output channels {out_channels} should be divisible by number of groups {groups}."
        self.grouped_conv = groups > 1

        # conv1 is used for downsampling
        # self.conv1 = nn.Conv2d(
        #     in_channels=in_channels,
        #     out_channels=in_channels,
        #     kernel_size=2,
        #     stride=2,
        #     padding=0,
        #     bias=not norm_layer,
        #     groups=groups,
        # )
        # if self.grouped_conv:
        #     self.conv1_1x1 = conv1x1(in_channels, in_channels, stride=1, bias=not norm_layer)
        self.conv1 = nn.AvgPool2d(kernel_size=2, stride=2)  # downsample by 2
        if self.grouped_conv:
            self.conv1_1x1 = nn.Identity()
        
        self.norm1 = norm_layer(in_channels) if norm_layer else nn.Identity()
        self.act1 = activation

        self.conv2 = conv3x3(
            in_channels=in_channels,
            out_channels=in_channels,
            stride=1,
            groups=groups,
            bias=not norm_layer,
        )
        if self.grouped_conv:
            self.conv2_1x1 = conv1x1(in_channels, in_channels, stride=1, bias=not norm_layer)
        
        self.norm2 = norm_layer(in_channels) if norm_layer else nn.Identity()
        self.act2 = activation

        self.conv3 = conv3x3(
            in_channels=in_channels,
            out_channels=out_channels,
            stride=1,
            groups=groups,
            bias=not norm_layer,
        )
        if self.grouped_conv:
            self.conv3_1x1 = conv1x1(out_channels, out_channels, stride=1, bias=not norm_layer)

        self.norm3 = norm_layer(out_channels) if norm_layer else nn.Identity()
        self.act3 = activation

        self.downsample = nn.Sequential(
            nn.AvgPool2d(kernel_size=2, stride=2),  # make sure the spatial sizes match
            conv1x1(in_channels, out_channels, stride=1, bias=not norm_layer),
            norm_layer(out_channels) if norm_layer else nn.Identity(),
        )
        
        self.apply(_init_weights)
    
    def forward(self, x: Tensor) -> Tensor:
        identity = x

        # downsample
        out = self.conv1(x)
        out = self.conv1_1x1(out) if self.grouped_conv else out
        out = self.norm1(out)
        out = self.act1(out)

        out = self.conv2(out)
        out = self.conv2_1x1(out) if self.grouped_conv else out
        out = self.norm2(out)
        out = self.act2(out)

        out = self.conv3(out)
        out = self.conv3_1x1(out) if self.grouped_conv else out
        out = self.norm3(out)

        # shortcut
        out += self.downsample(identity)
        out = self.act3(out)
        return out


class LightConvDownsample(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        norm_layer: Union[nn.BatchNorm2d, nn.GroupNorm, None] = nn.BatchNorm2d,
        activation: nn.Module = nn.ReLU(inplace=True),
    ) -> None:
        super().__init__()
        self.conv1 = DepthSeparableConv2d(
            in_channels=in_channels,
            out_channels=in_channels,
            kernel_size=2,
            stride=2,
            padding=0,
            bias=not norm_layer,
        )
        self.norm1 = norm_layer(in_channels) if norm_layer else nn.Identity()
        self.act1 = activation

        self.conv2 = DepthSeparableConv2d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=3,
            stride=1,
            padding=1,
            bias=not norm_layer,
        )
        self.norm2 = norm_layer(out_channels) if norm_layer else nn.Identity()
        self.act2 = activation

        self.conv3 = DepthSeparableConv2d(
            in_channels=out_channels,
            out_channels=out_channels,
            kernel_size=3,
            stride=1,
            padding=1,
            bias=not norm_layer,
        )
        self.norm3 = norm_layer(out_channels) if norm_layer else nn.Identity()
        self.act3 = activation

        self.downsample = nn.Sequential(
            nn.AvgPool2d(kernel_size=2, stride=2),  # make sure the spatial sizes match
            conv1x1(in_channels, out_channels, stride=1, bias=not norm_layer),
            norm_layer(out_channels) if norm_layer else nn.Identity(),
        )

        self.apply(_init_weights)

    def forward(self, x: Tensor) -> Tensor:
        identity = x

        # downsample
        out = self.conv1(x)
        out = self.norm1(out)
        out = self.act1(out)

        # refine 1
        out = self.conv2(out)
        out = self.norm2(out)
        out = self.act2(out)

        # refine 2
        out = self.conv3(out)
        out = self.norm3(out)

        # shortcut
        out += self.downsample(identity)
        out = self.act3(out)
        return x
    

class LighterConvDownsample(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        norm_layer: Union[nn.BatchNorm2d, nn.GroupNorm, None] = nn.BatchNorm2d,
        activation: nn.Module = nn.ReLU(inplace=True),
    ) -> None:
        super().__init__()
        self.conv1 = DepthSeparableConv2d(
            in_channels=in_channels,
            out_channels=in_channels,
            kernel_size=2,
            stride=2,
            padding=0,
            bias=not norm_layer,
        )
        self.norm1 = norm_layer(in_channels) if norm_layer else nn.Identity()
        self.act1 = activation

        self.conv2 = conv3x3(
            in_channels=in_channels,
            out_channels=in_channels,
            stride=1,
            groups=in_channels,
            bias=not norm_layer,
        )
        self.norm2 = norm_layer(in_channels) if norm_layer else nn.Identity()
        self.act2 = activation

        self.conv3 = conv1x1(
            in_channels=in_channels,
            out_channels=out_channels,
            stride=1,
            bias=not norm_layer,
        )
        self.norm3 = norm_layer(out_channels) if norm_layer else nn.Identity()
        self.act3 = activation

        self.downsample = nn.Sequential(
            nn.AvgPool2d(kernel_size=2, stride=2),  # make sure the spatial sizes match
            conv1x1(in_channels, out_channels, stride=1, bias=not norm_layer),
            norm_layer(out_channels) if norm_layer else nn.Identity(),
        )
    
    def forward(self, x: Tensor) -> Tensor:
        identity = x

        # downsample
        out = self.conv1(x)
        out = self.norm1(out)
        out = self.act1(out)

        # refine, depthwise conv
        out = self.conv2(out)
        out = self.norm2(out)
        out = self.act2(out)

        # refine, pointwise conv
        out = self.conv3(out)
        out = self.norm3(out)

        # shortcut
        out += self.downsample(identity)
        out = self.act3(out)
        return out