import functools
from .init import init_fn
import torch
import torch.nn as nn
import torch.nn.functional as F

class ResBlock(nn.Module):
    def __init__(self, in_channels, mid_channels, out_channels):
        super(ResBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=in_channels, out_channels=mid_channels, kernel_size=3, stride=1, padding=1)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(in_channels=mid_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1)

    def forward(self, x):
        output = self.conv2(self.relu(self.conv1(x)))
        output = torch.add(output, x)
        return output

class ResBlocks(nn.Module):
    def __init__(self, input_channels, num_resblocks, num_channels):
        super(ResBlocks, self).__init__()
        self.input_channels = input_channels
        self.first_conv = nn.Conv2d(in_channels=self.input_channels, out_channels=num_channels, kernel_size=3, stride=1, padding=1)

        modules = []
        for _ in range(num_resblocks):
            modules.append(ResBlock(in_channels=num_channels, mid_channels=num_channels, out_channels=num_channels))
        self.resblocks = nn.Sequential(*modules)

        fn = functools.partial(init_fn, init_type='kaiming_normal', init_bn_type='uniform', gain=0.2)
        self.apply(fn)

    def forward(self, h):
        shallow_feature = self.first_conv(h)
        new_h = self.resblocks(shallow_feature)
        return new_h

class Conv2DMod(nn.Module):
    def __init__(self, in_chan, out_chan, kernel, style_ch=82, demod=True, stride=1, padding=1, dilation=1, eps = 1e-8, **kwargs):
        super().__init__()
        self.filters = out_chan
        self.demod = demod
        self.kernel = kernel
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.weight = nn.Parameter(torch.randn((out_chan, in_chan, kernel, kernel)))
        self.eps = eps
        self.to_style = nn.Linear(style_ch, in_chan)
        nn.init.kaiming_normal_(self.weight, a=0, mode='fan_in', nonlinearity='leaky_relu')

    # def _get_same_padding(self, size, kernel, dilation, stride):
    #     return ((size - 1) * (stride - 1) + dilation * (kernel - 1)) // 2

    def forward(self, x, scale):
        b, c, h, w = x.shape

        y = self.to_style(scale)
        w1 = y[:, None, :, None, None]
        # print('w1--', w1.shape)
        w2 = self.weight[None, :, :, :, :]
        # print('w2--',w2.shape)
        weights = w2 * (w1 + 1)

        if self.demod:
            d = torch.rsqrt((weights ** 2).sum(dim=(2, 3, 4), keepdim=True) + self.eps)
            weights = weights * d

        x = x.reshape(1, -1, h, w)

        _, _, *ws = weights.shape
        weights = weights.reshape(b * self.filters, *ws)

        # padding = self._get_same_padding(h, self.kernel, self.dilation, self.stride)
        x = F.conv2d(x, weights, padding=self.padding, groups=b)
        # print(x.shape)

        x = x.reshape(-1, self.filters, h, w)
        return x

class ResBlocks_mod(nn.Module):
    def __init__(self, input_channels, num_resblocks, num_channels):
        super(ResBlocks_mod, self).__init__()
        self.input_channels = input_channels
        self.first_conv = nn.Conv2d(in_channels=self.input_channels, out_channels=num_channels, kernel_size=3, stride=1, padding=1)
        self.num_resblocks = num_resblocks
        modules = []
        for _ in range(num_resblocks):
            modules.append(ResBlock(in_channels=num_channels, mid_channels=num_channels, out_channels=num_channels))
        self.resblocks = nn.Sequential(*modules)
        fn = functools.partial(init_fn, init_type='kaiming_normal', init_bn_type='uniform', gain=0.2)
        self.apply(fn)

        modc = []
        for _ in range(num_resblocks // 5):
            modc.append(Conv2DMod(num_channels, num_channels, 3))
        self.mbconv = nn.Sequential(*modc)

    def forward(self, h, scale):
        new_h = self.first_conv(h)
        for i in range(self.num_resblocks//5):
            new_h = self.mbconv[i](new_h, scale)
            new_h = self.resblocks[i*5:(i+1)*5](new_h)
        return new_h

class D(nn.Module):
    def __init__(self, in_channels, mid_channels, out_channels):
        super(D, self).__init__()
        layers = []
        layers.append(nn.Conv2d(in_channels=in_channels, out_channels=mid_channels, kernel_size=3, stride=1, padding=1))
        layers.append(nn.ReLU())
        layers.append(nn.Conv2d(in_channels=mid_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1))
        self.convs = nn.Sequential(*layers)

        fn = functools.partial(init_fn, init_type='kaiming_normal', init_bn_type='uniform', gain=0.2)
        self.apply(fn)

    def forward(self, x):
        x = self.convs(x)
        return x