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


class WNConv2d(nn.Module):
    """Weight-normalized 2d convolution.
    Args:
        in_channels (int): Number of channels in the input.
        out_channels (int): Number of channels in the output.
        kernel_size (int): Side length of each convolutional kernel.
        padding (int): Padding to add on edges of input.
        bias (bool): Use bias in the convolution operation.
    """

    def __init__(self, in_channels, out_channels, kernel_size, padding, bias=True):
        super(WNConv2d, self).__init__()
        self.conv = nn.utils.weight_norm(
            nn.Conv2d(
                in_channels, out_channels, kernel_size, padding=padding, bias=bias
            )
        )

    def forward(self, x):
        x = self.conv(x)

        return x


class ResidualBlock(nn.Module):
    """ResNet basic block with weight norm."""

    def __init__(self, in_channels, out_channels):
        super(ResidualBlock, self).__init__()

        self.in_norm = nn.BatchNorm2d(in_channels)
        self.in_conv = WNConv2d(
            in_channels, out_channels, kernel_size=3, padding=1, bias=False
        )

        self.out_norm = nn.BatchNorm2d(out_channels)
        self.out_conv = WNConv2d(
            out_channels, out_channels, kernel_size=3, padding=1, bias=True
        )

    def forward(self, x):
        skip = x

        x = self.in_norm(x)
        x = F.relu(x)
        x = self.in_conv(x)

        x = self.out_norm(x)
        x = F.relu(x)
        x = self.out_conv(x)

        x = x + skip

        return x


class ResNet(nn.Module):
    """ResNet for scale and translate factors in Real NVP.

    Args:
        in_channels (int): Number of channels in the input.
        mid_channels (int): Number of channels in the intermediate layers.
        out_channels (int): Number of channels in the output.
        num_blocks (int): Number of residual blocks in the network.
        kernel_size (int): Side length of each filter in convolutional layers.
        padding (int): Padding for convolutional layers.
        double_after_norm (bool): Double input after input BatchNorm.
    """

    def __init__(
        self,
        in_channels,
        mid_channels,
        out_channels,
        num_blocks,
        kernel_size,
        padding,
        double_after_norm,
    ):
        super(ResNet, self).__init__()
        self.in_norm = nn.BatchNorm2d(in_channels)
        self.double_after_norm = double_after_norm
        self.in_conv = WNConv2d(
            2 * in_channels, mid_channels, kernel_size, padding, bias=True
        )
        self.in_skip = WNConv2d(
            mid_channels, mid_channels, kernel_size=1, padding=0, bias=True
        )

        self.blocks = nn.ModuleList(
            [ResidualBlock(mid_channels, mid_channels) for _ in range(num_blocks)]
        )
        self.skips = nn.ModuleList(
            [
                WNConv2d(
                    mid_channels, mid_channels, kernel_size=1, padding=0, bias=True
                )
                for _ in range(num_blocks)
            ]
        )

        self.out_norm = nn.BatchNorm2d(mid_channels)
        self.out_conv = WNConv2d(
            mid_channels, out_channels, kernel_size=1, padding=0, bias=True
        )

    def forward(self, x):
        x = self.in_norm(x)
        if self.double_after_norm:
            x *= 2.0
        x = torch.cat((x, -x), dim=1)
        x = F.relu(x)
        x = self.in_conv(x)
        x_skip = self.in_skip(x)

        for block, skip in zip(self.blocks, self.skips):
            x = block(x)
            x_skip += skip(x)

        x = self.out_norm(x_skip)
        x = F.relu(x)
        x = self.out_conv(x)

        return x
