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):
    """
    cin
    cmid
    cout
    nblocks
    kernel
    padding
    double_after_norm
    """

    def __init__(self, cin, cmid, cout, nblocks, kernel, padding, double_after_norm):
        super(ResNet, self).__init__()

        # Attributes
        self.batch_norm_in = nn.BatchNorm2d(cin)
        self.double_after_norm = double_after_norm
        self.conv_in = WNConv2d(2 * cin, cmid, kernel, padding, bias=True)
        self.skip_in = WNConv2d(cmid, cmid, kernel_size=1, padding=0, bias=True)

        self.blocks = nn.ModuleList([ResidualBlock(cmid, cmid) for _ in range(nblocks)])

        self.skips = nn.ModuleList([WNConv2d(cmid, cmid, kernel_size=1, padding=0, bias=True) for _ in range(nblocks)])

        self.batch_norm_out = nn.BatchNorm2d(cmid)

        self.conv_out = WNConv2d(cmid, cout, kernel_size=1, padding=0, bias=True)

    def forward(self, x):
        x = self.batch_norm_in(x)
        if self.double_after_norm:
            x = x * 2
        x = torch.cat((x, -x), dim=1)
        x = F.relu(x)
        x = self.conv_in(x)
        x_skip = self.skip_in(x)

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

        return x