import torch
import torch.nn as nn
import math

from .DCNv2.dcn_v2 import DCN_ID


class DenseBlock(torch.nn.Module):
    def __init__(self, input_size, output_size, bias=True, activation='relu', norm='batch'):
        super(DenseBlock, self).__init__()
        self.fc = torch.nn.Linear(input_size, output_size, bias=bias)

        self.norm = norm
        if self.norm =='batch':
            self.bn = torch.nn.BatchNorm1d(output_size)
        elif self.norm == 'instance':
            self.bn = torch.nn.InstanceNorm1d(output_size)

        self.activation = activation
        if self.activation == 'relu':
            self.act = torch.nn.ReLU(True)
        elif self.activation == 'prelu':
            self.act = torch.nn.PReLU()
        elif self.activation == 'lrelu':
            self.act = torch.nn.LeakyReLU(0.2, True)
        elif self.activation == 'tanh':
            self.act = torch.nn.Tanh()
        elif self.activation == 'sigmoid':
            self.act = torch.nn.Sigmoid()

    def forward(self, x):
        if self.norm is not None:
            out = self.bn(self.fc(x))
        else:
            out = self.fc(x)

        if self.activation is not None:
            return self.act(out)
        else:
            return out


class ConvBlock(torch.nn.Module):
    def __init__(self, input_size, output_size, kernel_size=3, stride=1, padding=1, bias=True, activation='prelu', norm=None):
        super(ConvBlock, self).__init__()
        self.conv = torch.nn.Conv2d(input_size, output_size, kernel_size, stride, padding, bias=bias)
        torch.nn.init.kaiming_normal_(self.conv.weight)
        if bias:
            self.conv.bias.data.zero_()

        self.norm = norm
        if self.norm =='batch':
            self.bn = torch.nn.BatchNorm2d(output_size)
        elif self.norm == 'instance':
            self.bn = torch.nn.InstanceNorm2d(output_size)

        self.activation = activation
        if self.activation == 'relu':
            self.act = torch.nn.ReLU(True)
        elif self.activation == 'prelu':
            self.act = torch.nn.PReLU()
        elif self.activation == 'lrelu':
            self.act = torch.nn.LeakyReLU(0.2, True)
        elif self.activation == 'tanh':
            self.act = torch.nn.Tanh()
        elif self.activation == 'sigmoid':
            self.act = torch.nn.Sigmoid()

    def forward(self, x):
        if self.norm is not None:
            out = self.bn(self.conv(x))
        else:
            out = self.conv(x)

        if self.activation is not None:
            return self.act(out)
        else:
            return out


class DeconvBlock(torch.nn.Module):
    def __init__(self, input_size, output_size, kernel_size=4, stride=2, padding=1, bias=True, activation='prelu', norm=None):
        super(DeconvBlock, self).__init__()
        self.deconv = torch.nn.ConvTranspose2d(input_size, output_size, kernel_size, stride, padding, bias=bias)
        torch.nn.init.kaiming_normal_(self.deconv.weight)
        if bias:
            self.deconv.bias.data.zero_()

        self.norm = norm
        if self.norm == 'batch':
            self.bn = torch.nn.BatchNorm2d(output_size)
        elif self.norm == 'instance':
            self.bn = torch.nn.InstanceNorm2d(output_size)

        self.activation = activation
        if self.activation == 'relu':
            self.act = torch.nn.ReLU(True)
        elif self.activation == 'prelu':
            self.act = torch.nn.PReLU()
        elif self.activation == 'lrelu':
            self.act = torch.nn.LeakyReLU(0.2, True)
        elif self.activation == 'tanh':
            self.act = torch.nn.Tanh()
        elif self.activation == 'sigmoid':
            self.act = torch.nn.Sigmoid()

    def forward(self, x):
        if self.norm is not None:
            out = self.bn(self.deconv(x))
        else:
            out = self.deconv(x)

        if self.activation is not None:
            return self.act(out)
        else:
            return out
        
        
class ResnetBlock(torch.nn.Module):
    def __init__(self, num_filter, kernel_size=3, stride=1, padding=1, bias=True, activation='prelu', norm='batch'):
        super(ResnetBlock, self).__init__()
        self.conv1 = torch.nn.Conv2d(num_filter, num_filter, kernel_size, stride, padding, bias=bias)
        self.conv2 = torch.nn.Conv2d(num_filter, num_filter, kernel_size, stride, padding, bias=bias)

        self.norm = norm
        if self.norm == 'batch':
            self.bn = torch.nn.BatchNorm2d(num_filter)
        elif norm == 'instance':
            self.bn = torch.nn.InstanceNorm2d(num_filter)

        self.activation = activation
        if self.activation == 'relu':
            self.act = torch.nn.ReLU(True)
        elif self.activation == 'prelu':
            self.act = torch.nn.PReLU()
        elif self.activation == 'lrelu':
            self.act = torch.nn.LeakyReLU(0.2, True)
        elif self.activation == 'tanh':
            self.act = torch.nn.Tanh()
        elif self.activation == 'sigmoid':
            self.act = torch.nn.Sigmoid()

    def forward(self, x):
        residual = x
        if self.norm is not None:
            out = self.bn(self.conv1(x))
        else:
            out = self.conv1(x)

        if self.activation is not None:
            out = self.act(out)

        if self.norm is not None:
            out = self.bn(self.conv2(out))
        else:
            out = self.conv2(out)

        out = torch.add(out, residual)
        
        if self.activation is not None:
            out = self.act(out)
            
        return out


class ResnetBlock_fixup_init(torch.nn.Module):
    def __init__(self, num_filter, kernel_size=3, stride=1, padding=1, bias=True, activation='prelu', norm='batch', num_residual=20):
        super(ResnetBlock_fixup_init, self).__init__()
        self.conv1 = torch.nn.Conv2d(num_filter, num_filter, kernel_size, stride, padding, bias=False)
        self.conv2 = torch.nn.Conv2d(num_filter, num_filter, kernel_size, stride, padding, bias=False)

        torch.nn.init.kaiming_normal_(self.conv1.weight)
        for param in self.conv1.parameters():
            param = param / ( num_residual ** (1 / 2) )
        self.conv2.weight.data.zero_()

        if bias:
            self.conv1_bias = nn.Parameter(torch.zeros(1))
            self.conv2_bias = nn.Parameter(torch.zeros(1))
            self.act1_bias = nn.Parameter(torch.zeros(1))
            self.act2_bias = nn.Parameter(torch.zeros(1))

        self.multiplier = nn.Parameter(torch.ones(1))


        self.norm = norm
        if self.norm == 'batch':
            self.bn = torch.nn.BatchNorm2d(num_filter)
        elif norm == 'instance':
            self.bn = torch.nn.InstanceNorm2d(num_filter)

        self.activation = activation
        if self.activation == 'relu':
            self.act = torch.nn.ReLU(True)
        elif self.activation == 'prelu':
            self.act = torch.nn.PReLU()
        elif self.activation == 'lrelu':
            self.act = torch.nn.LeakyReLU(0.2, True)
        elif self.activation == 'tanh':
            self.act = torch.nn.Tanh()
        elif self.activation == 'sigmoid':
            self.act = torch.nn.Sigmoid()

    def forward(self, x):
        residual = x
        if self.norm is not None:
            out = self.bn(self.conv1(x))
        else:
            out = self.conv1(x + self.conv1_bias)

        if self.activation is not None:
            out = self.act(out + self.act1_bias)

        if self.norm is not None:
            out = self.bn(self.conv2(out))
        else:
            out = self.conv2(out + self.conv2_bias)

        out = out * self.multiplier

        out = torch.add(out, residual)
        
        if self.activation is not None:
            out = self.act(out + self.act2_bias)
            
        return out

class UpBlock(torch.nn.Module):
    def __init__(self, num_filter, kernel_size=8, stride=4, padding=2, bias=True, activation='prelu', norm=None):
        super(UpBlock, self).__init__()
        self.up_conv1 = DeconvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)
        self.up_conv2 = ConvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)
        self.up_conv3 = DeconvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)        

    def forward(self, x):
    	h0 = self.up_conv1(x)
    	l0 = self.up_conv2(h0)
    	h1 = self.up_conv3(l0 - x)
    	return h1 + h0

class UpBlockPix(torch.nn.Module):
    def __init__(self, num_filter, kernel_size=8, stride=4, padding=2, scale=4, bias=True, activation='prelu', norm=None):
        super(UpBlockPix, self).__init__()
        self.up_conv1 = Upsampler(scale,num_filter)
        self.up_conv2 = ConvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)
        self.up_conv3 = Upsampler(scale,num_filter)        

    def forward(self, x):
    	h0 = self.up_conv1(x)
    	l0 = self.up_conv2(h0)
    	h1 = self.up_conv3(l0 - x)
    	return h1 + h0
        
class D_UpBlock(torch.nn.Module):
    def __init__(self, num_filter, kernel_size=8, stride=4, padding=2, num_stages=1, bias=True, activation='prelu', norm=None):
        super(D_UpBlock, self).__init__()
        self.conv = ConvBlock(num_filter*num_stages, num_filter, 1, 1, 0, activation, norm=None)
        self.up_conv1 = DeconvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)
        self.up_conv2 = ConvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)
        self.up_conv3 = DeconvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)        

    def forward(self, x):
    	x = self.conv(x)
    	h0 = self.up_conv1(x)
    	l0 = self.up_conv2(h0)
    	h1 = self.up_conv3(l0 - x)
    	return h1 + h0

class D_UpBlockPix(torch.nn.Module):
    def __init__(self, num_filter, kernel_size=8, stride=4, padding=2, num_stages=1, scale=4, bias=True, activation='prelu', norm=None):
        super(D_UpBlockPix, self).__init__()
        self.conv = ConvBlock(num_filter*num_stages, num_filter, 1, 1, 0, activation, norm=None)
        self.up_conv1 = Upsampler(scale,num_filter)
        self.up_conv2 = ConvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)
        self.up_conv3 = Upsampler(scale,num_filter)

    def forward(self, x):
    	x = self.conv(x)
    	h0 = self.up_conv1(x)
    	l0 = self.up_conv2(h0)
    	h1 = self.up_conv3(l0 - x)
    	return h1 + h0

class DownBlock(torch.nn.Module):
    def __init__(self, num_filter, kernel_size=8, stride=4, padding=2, bias=True, activation='prelu', norm=None):
        super(DownBlock, self).__init__()
        self.down_conv1 = ConvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)
        self.down_conv2 = DeconvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)
        self.down_conv3 = ConvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)

    def forward(self, x):
    	l0 = self.down_conv1(x)
    	h0 = self.down_conv2(l0)
    	l1 = self.down_conv3(h0 - x)
    	return l1 + l0

class DownBlockPix(torch.nn.Module):
    def __init__(self, num_filter, kernel_size=8, stride=4, padding=2, scale=4,bias=True, activation='prelu', norm=None):
        super(DownBlockPix, self).__init__()
        self.down_conv1 = ConvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)
        self.down_conv2 = Upsampler(scale,num_filter)
        self.down_conv3 = ConvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)

    def forward(self, x):
    	l0 = self.down_conv1(x)
    	h0 = self.down_conv2(l0)
    	l1 = self.down_conv3(h0 - x)
    	return l1 + l0

class D_DownBlock(torch.nn.Module):
    def __init__(self, num_filter, kernel_size=8, stride=4, padding=2, num_stages=1, bias=True, activation='prelu', norm=None):
        super(D_DownBlock, self).__init__()
        self.conv = ConvBlock(num_filter*num_stages, num_filter, 1, 1, 0, activation, norm=None)
        self.down_conv1 = ConvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)
        self.down_conv2 = DeconvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)
        self.down_conv3 = ConvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)

    def forward(self, x):
    	x = self.conv(x)
    	l0 = self.down_conv1(x)
    	h0 = self.down_conv2(l0)
    	l1 = self.down_conv3(h0 - x)
    	return l1 + l0

class D_DownBlockPix(torch.nn.Module):
    def __init__(self, num_filter, kernel_size=8, stride=4, padding=2, num_stages=1, scale=4, bias=True, activation='prelu', norm=None):
        super(D_DownBlockPix, self).__init__()
        self.conv = ConvBlock(num_filter*num_stages, num_filter, 1, 1, 0, activation, norm=None)
        self.down_conv1 = ConvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)
        self.down_conv2 = Upsampler(scale,num_filter)
        self.down_conv3 = ConvBlock(num_filter, num_filter, kernel_size, stride, padding, activation, norm=None)

    def forward(self, x):
    	x = self.conv(x)
    	l0 = self.down_conv1(x)
    	h0 = self.down_conv2(l0)
    	l1 = self.down_conv3(h0 - x)
    	return l1 + l0

class PSBlock(torch.nn.Module):
    def __init__(self, input_size, output_size, scale_factor, kernel_size=3, stride=1, padding=1, bias=True, activation='prelu', norm='batch'):
        super(PSBlock, self).__init__()
        self.conv = torch.nn.Conv2d(input_size, output_size * scale_factor**2, kernel_size, stride, padding, bias=bias)
        self.ps = torch.nn.PixelShuffle(scale_factor)

        self.norm = norm
        if self.norm == 'batch':
            self.bn = torch.nn.BatchNorm2d(output_size)
        elif norm == 'instance':
            self.bn = torch.nn.InstanceNorm2d(output_size)

        self.activation = activation
        if self.activation == 'relu':
            self.act = torch.nn.ReLU(True)
        elif self.activation == 'prelu':
            self.act = torch.nn.PReLU()
        elif self.activation == 'lrelu':
            self.act = torch.nn.LeakyReLU(0.2, True)
        elif self.activation == 'tanh':
            self.act = torch.nn.Tanh()
        elif self.activation == 'sigmoid':
            self.act = torch.nn.Sigmoid()

    def forward(self, x):
        if self.norm is not None:
            out = self.bn(self.ps(self.conv(x)))
        else:
            out = self.ps(self.conv(x))

        if self.activation is not None:
            out = self.act(out)
        return out


class Upsampler(torch.nn.Module):
    def __init__(self, scale, n_feat, bn=False, act='prelu', bias=True):
        super(Upsampler, self).__init__()
        modules = []
        for _ in range(int(math.log(scale, 2))):
            modules.append(ConvBlock(n_feat, 4 * n_feat, 3, 1, 1, bias, activation=None, norm=None))
            modules.append(torch.nn.PixelShuffle(2))
            if bn: modules.append(torch.nn.BatchNorm2d(n_feat))
            #modules.append(torch.nn.PReLU())
        self.up = torch.nn.Sequential(*modules)
        
        self.activation = act
        if self.activation == 'relu':
            self.act = torch.nn.ReLU(True)
        elif self.activation == 'prelu':
            self.act = torch.nn.PReLU()
        elif self.activation == 'lrelu':
            self.act = torch.nn.LeakyReLU(0.2, True)
        elif self.activation == 'tanh':
            self.act = torch.nn.Tanh()
        elif self.activation == 'sigmoid':
            self.act = torch.nn.Sigmoid()

    def forward(self, x):
        out = self.up(x)
        if self.activation is not None:
            out = self.act(out)
        return out
             

class Upsample2xBlock(torch.nn.Module):
    def __init__(self, input_size, output_size, bias=True, upsample='deconv', activation='relu', norm='batch'):
        super(Upsample2xBlock, self).__init__()
        scale_factor = 2
        # 1. Deconvolution (Transposed convolution)
        if upsample == 'deconv':
            self.upsample = DeconvBlock(input_size, output_size,
                                        kernel_size=4, stride=2, padding=1,
                                        bias=bias, activation=activation, norm=norm)

        # 2. Sub-pixel convolution (Pixel shuffler)
        elif upsample == 'ps':
            self.upsample = PSBlock(input_size, output_size, scale_factor=scale_factor,
                                    bias=bias, activation=activation, norm=norm)

        # 3. Resize and Convolution
        elif upsample == 'rnc':
            self.upsample = torch.nn.Sequential(
                torch.nn.Upsample(scale_factor=scale_factor, mode='nearest'),
                ConvBlock(input_size, output_size,
                          kernel_size=3, stride=1, padding=1,
                          bias=bias, activation=activation, norm=norm)
            )

    def forward(self, x):
        out = self.upsample(x)
        return out


class PyramidModule(torch.nn.Module):
    def __init__(self, num_inchannels, activation='prelu'):
        super(PyramidModule, self).__init__()
        
        self.l1_1 = ResnetBlock(num_inchannels, kernel_size=3, stride=1, padding=1, bias=True, activation=activation, norm=None)
        self.l1_2 = ResnetBlock(num_inchannels, kernel_size=3, stride=1, padding=1, bias=True, activation=activation, norm=None)
        self.l1_3 = ResnetBlock(num_inchannels, kernel_size=3, stride=1, padding=1, bias=True, activation=activation, norm=None)
        self.l1_4 = ResnetBlock(num_inchannels, kernel_size=3, stride=1, padding=1, bias=True, activation=activation, norm=None)
        self.l1_5 = ResnetBlock(num_inchannels, kernel_size=3, stride=1, padding=1, bias=True, activation=activation, norm=None)
        
        self.l2_1 = ResnetBlock(num_inchannels*2, kernel_size=3, stride=1, padding=1, bias=True, activation=activation, norm=None)
        self.l2_2 = ResnetBlock(num_inchannels*2, kernel_size=3, stride=1, padding=1, bias=True, activation=activation, norm=None)
        self.l2_3 = ResnetBlock(num_inchannels*2, kernel_size=3, stride=1, padding=1, bias=True, activation=activation, norm=None)
        self.l2_4 = ResnetBlock(num_inchannels*2, kernel_size=3, stride=1, padding=1, bias=True, activation=activation, norm=None)
        
        self.l3_1 = ResnetBlock(num_inchannels*4, kernel_size=3, stride=1, padding=1, bias=True, activation=activation, norm=None)
        self.l3_2 = ResnetBlock(num_inchannels*4, kernel_size=3, stride=1, padding=1, bias=True, activation=activation, norm=None)
        self.l3_3 = ResnetBlock(num_inchannels*4, kernel_size=3, stride=1, padding=1, bias=True, activation=activation, norm=None)
        
        self.down1 = ConvBlock(num_inchannels, num_inchannels*2, 4, 2, 1, bias=True, activation=activation, norm=None)
        self.down2 = ConvBlock(num_inchannels*2, num_inchannels*4, 4, 2, 1, bias=True, activation=activation, norm=None)
        
        self.up1 = DeconvBlock(num_inchannels*2, num_inchannels, 4, 2, 1, bias=True,  activation=activation, norm=None)
        self.up2 = DeconvBlock(num_inchannels*4, num_inchannels*2, 4, 2, 1, bias=True, activation=activation, norm=None)
        
        self.final = ConvBlock(num_inchannels, num_inchannels, 3, 1, 1, bias=True, activation=activation, norm=None)
        
    def forward(self, x):
        out1_1 = self.l1_1(x)
        out2_1 = self.l2_1(self.down1(out1_1))
        out3_1 = self.l3_1(self.down2(out2_1))
        
        out1_2 = self.l1_2(out1_1 + self.up1(out2_1))
        out2_2 = self.l2_2(out2_1 + self.down1(out1_2) + self.up2(out3_1))
        out3_2 = self.l3_2(out3_1 + self.down2(out2_2))
        
        out1_3 = self.l1_3(out1_2 + self.up1(out2_2))
        out2_3 = self.l2_3(out2_2 + self.down1(out1_3) + self.up2(out3_2))
        out3_3 = self.l3_3(out3_2 + self.down2(out2_3))
        
        out1_4 = self.l1_4(out1_3 + self.up1(out2_3))
        out2_4 = self.l2_4(out2_3 + self.down1(out1_4) + self.up2(out3_3))
        
        out1_5 = self.l1_5(out1_4 + self.up1(out2_4))
        
        final = self.final(out1_5)
        
        return final
    
    
class PyramidModule_fixup_init(torch.nn.Module):
    def __init__(self, num_inchannels, activation='prelu'):
        super(PyramidModule_fixup_init, self).__init__()
        
        self.l1_1 = ResnetBlock_fixup_init(num_inchannels, kernel_size=3, stride=1, padding=1, bias=True, activation=activation, norm=None)
        self.l1_2 = ResnetBlock_fixup_init(num_inchannels, kernel_size=3, stride=1, padding=1, bias=True, activation=activation, norm=None)
        self.l1_3 = ResnetBlock_fixup_init(num_inchannels, kernel_size=3, stride=1, padding=1, bias=True, activation=activation, norm=None)
        self.l1_4 = ResnetBlock_fixup_init(num_inchannels, kernel_size=3, stride=1, padding=1, bias=True, activation=activation, norm=None)
        self.l1_5 = ResnetBlock_fixup_init(num_inchannels, kernel_size=3, stride=1, padding=1, bias=True, activation=activation, norm=None)
        
        self.l2_1 = ResnetBlock_fixup_init(num_inchannels*2, kernel_size=3, stride=1, padding=1, bias=True, activation=activation, norm=None)
        self.l2_2 = ResnetBlock_fixup_init(num_inchannels*2, kernel_size=3, stride=1, padding=1, bias=True, activation=activation, norm=None)
        self.l2_3 = ResnetBlock_fixup_init(num_inchannels*2, kernel_size=3, stride=1, padding=1, bias=True, activation=activation, norm=None)
        self.l2_4 = ResnetBlock_fixup_init(num_inchannels*2, kernel_size=3, stride=1, padding=1, bias=True, activation=activation, norm=None)
        
        self.l3_1 = ResnetBlock_fixup_init(num_inchannels*4, kernel_size=3, stride=1, padding=1, bias=True, activation=activation, norm=None)
        self.l3_2 = ResnetBlock_fixup_init(num_inchannels*4, kernel_size=3, stride=1, padding=1, bias=True, activation=activation, norm=None)
        self.l3_3 = ResnetBlock_fixup_init(num_inchannels*4, kernel_size=3, stride=1, padding=1, bias=True, activation=activation, norm=None)
        
        self.down1 = ConvBlock(num_inchannels, num_inchannels*2, 4, 2, 1, bias=True, activation=activation, norm=None)
        self.down2 = ConvBlock(num_inchannels*2, num_inchannels*4, 4, 2, 1, bias=True, activation=activation, norm=None)
        
        self.up1 = DeconvBlock(num_inchannels*2, num_inchannels, 4, 2, 1, bias=True,  activation=activation, norm=None)
        self.up2 = DeconvBlock(num_inchannels*4, num_inchannels*2, 4, 2, 1, bias=True, activation=activation, norm=None)
        
        self.final = ConvBlock(num_inchannels, num_inchannels, 3, 1, 1, bias=True, activation=activation, norm=None)
        
    def forward(self, x):
        out1_1 = self.l1_1(x)
        out2_1 = self.l2_1(self.down1(out1_1))
        out3_1 = self.l3_1(self.down2(out2_1))
        
        out1_2 = self.l1_2(out1_1 + self.up1(out2_1))
        out2_2 = self.l2_2(out2_1 + self.down1(out1_2) + self.up2(out3_1))
        out3_2 = self.l3_2(out3_1 + self.down2(out2_2))
        
        out1_3 = self.l1_3(out1_2 + self.up1(out2_2))
        out2_3 = self.l2_3(out2_2 + self.down1(out1_3) + self.up2(out3_2))
        out3_3 = self.l3_3(out3_2 + self.down2(out2_3))
        
        out1_4 = self.l1_4(out1_3 + self.up1(out2_3))
        out2_4 = self.l2_4(out2_3 + self.down1(out1_4) + self.up2(out3_3))
        
        out1_5 = self.l1_5(out1_4 + self.up1(out2_4))
        
        final = self.final(out1_5)
        
        return final


# borrow from https://github.com/xinntao/EDVR
class PCDAlignment(nn.Module):
    """Alignment module using Pyramid, Cascading and Deformable convolution
    (PCD). It is used in EDVR.
    Ref:
        EDVR: Video Restoration with Enhanced Deformable Convolutional Networks
    Args:
        num_feat (int): Channel number of middle features. Default: 64.
        deformable_groups (int): Deformable groups. Defaults: 8.
    """

    def __init__(self, num_feat=64, deformable_groups=8, activation='prelu'):
        super(PCDAlignment, self).__init__()

        # Pyramid has three levels:
        # L3: level 3, 1/4 spatial size
        # L2: level 2, 1/2 spatial size
        # L1: level 1, original spatial size

        self.conv_l2 = nn.Sequential(
            ConvBlock(num_feat, num_feat, kernel_size=3, stride=2, padding=1, norm=None, activation=activation),
            ConvBlock(num_feat, num_feat, kernel_size=3, stride=1, padding=1, norm=None, activation=activation),
        )

        self.conv_l3 = nn.Sequential(
            ConvBlock(num_feat, num_feat, kernel_size=3, stride=2, padding=1, norm=None, activation=activation),
            ConvBlock(num_feat, num_feat, kernel_size=3, stride=1, padding=1, norm=None, activation=activation),
        )


        self.offset_conv1 = nn.ModuleDict()
        self.offset_conv2 = nn.ModuleDict()
        self.offset_conv3 = nn.ModuleDict()
        self.dcn_pack = nn.ModuleDict()
        self.feat_conv = nn.ModuleDict()

        # Pyramids
        for i in range(3, 0, -1):
            level = f'l{i}'
            self.offset_conv1[level] = nn.Conv2d(num_feat * 2, num_feat, 3, 1,
                                                 1)
            if i == 3:
                self.offset_conv2[level] = nn.Conv2d(num_feat, num_feat, 3, 1,
                                                     1)
            else:
                self.offset_conv2[level] = nn.Conv2d(num_feat * 2, num_feat, 3,
                                                     1, 1)
                self.offset_conv3[level] = nn.Conv2d(num_feat, num_feat, 3, 1,
                                                     1)
            self.dcn_pack[level] = DCN_ID(num_feat, num_feat, num_feat, 3, 1, padding=1, deformable_groups=deformable_groups)

            if i < 3:
                self.feat_conv[level] = nn.Conv2d(num_feat * 2, num_feat, 3, 1,
                                                  1)

        # Cascading dcn
        self.cas_offset_conv1 = nn.Conv2d(num_feat * 2, num_feat, 3, 1, 1)
        self.cas_offset_conv2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
        self.cas_dcnpack = DCN_ID(num_feat, num_feat, num_feat, 3, 1, padding=1, deformable_groups=deformable_groups)

        self.upsample = nn.Upsample(
            scale_factor=2, mode='bilinear', align_corners=False)
        
        if activation == 'relu':
            self.act = nn.ReLU(True)
        elif activation == 'prelu':
            self.act = nn.PReLU()
        elif activation == 'lrelu':
            self.act = nn.LeakyReLU(negative_slope=0.1, inpcale=True)
        else:
            raise NotImplementedError


    def forward(self, base_feat, neigbor_feat):
        """Align neighboring frame features to the reference frame features.
        Args:
            nbr_feat_l (list[Tensor]): Neighboring feature list. It
                contains three pyramid levels (L1, L2, L3),
                each with shape (b, c, h, w).
            ref_feat_l (list[Tensor]): Reference feature list. It
                contains three pyramid levels (L1, L2, L3),
                each with shape (b, c, h, w).
        Returns:
            Tensor: Aligned features.
        """

        nbr_feat_l = [neigbor_feat]
        ref_feat_l = [base_feat]

        nbr_feat_l.append(self.conv_l2(nbr_feat_l[0]))
        nbr_feat_l.append(self.conv_l3(nbr_feat_l[1]))

        ref_feat_l.append(self.conv_l2(ref_feat_l[0]))
        ref_feat_l.append(self.conv_l3(ref_feat_l[1]))

        # Pyramids
        upsampled_offset, upsampled_feat = None, None
        for i in range(3, 0, -1):
            level = f'l{i}'
            offset = torch.cat([nbr_feat_l[i - 1], ref_feat_l[i - 1]], dim=1)
            offset = self.act(self.offset_conv1[level](offset))
            if i == 3:
                offset = self.act(self.offset_conv2[level](offset))
            else:
                offset = self.act(self.offset_conv2[level](torch.cat(
                    [offset, upsampled_offset], dim=1)))
                offset = self.act(self.offset_conv3[level](offset))

            feat = self.dcn_pack[level](nbr_feat_l[i - 1], offset)
            if i < 3:
                feat = self.feat_conv[level](
                    torch.cat([feat, upsampled_feat], dim=1))
            if i > 1:
                feat = self.act(feat)

            if i > 1:  # upsample offset and features
                # x2: when we upsample the offset, we should also enlarge
                # the magnitude.
                upsampled_offset = self.upsample(offset) * 2
                upsampled_feat = self.upsample(feat)

        # Cascading
        offset = torch.cat([feat, ref_feat_l[0]], dim=1)
        offset = self.act(
            self.cas_offset_conv2(self.act(self.cas_offset_conv1(offset))))
        feat = self.act(self.cas_dcnpack(feat, offset))
        return feat


class TSAFusion(nn.Module):
    """Temporal Spatial Attention (TSA) fusion module.
    Temporal: Calculate the correlation between center frame and
        neighboring frames;
    Spatial: It has 3 pyramid levels, the attention is similar to SFT.
        (SFT: Recovering realistic texture in image super-resolution by deep
            spatial feature transform.)
    Args:
        num_feat (int): Channel number of middle features. Default: 64.
        num_frame (int): Number of frames. Default: 5.
        center_frame_idx (int): The index of center frame. Default: 2.
    """

    def __init__(self, num_feat=64, num_frame=5, center_frame_idx=2):
        super(TSAFusion, self).__init__()
        self.center_frame_idx = center_frame_idx
        # temporal attention (before fusion conv)
        self.temporal_attn1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
        self.temporal_attn2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
        self.feat_fusion = nn.Conv2d(num_frame * num_feat, num_feat, 1, 1)

        # spatial attention (after fusion conv)
        self.max_pool = nn.MaxPool2d(3, stride=2, padding=1)
        self.avg_pool = nn.AvgPool2d(3, stride=2, padding=1)
        self.spatial_attn1 = nn.Conv2d(num_frame * num_feat, num_feat, 1)
        self.spatial_attn2 = nn.Conv2d(num_feat * 2, num_feat, 1)
        self.spatial_attn3 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
        self.spatial_attn4 = nn.Conv2d(num_feat, num_feat, 1)
        self.spatial_attn5 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
        self.spatial_attn_l1 = nn.Conv2d(num_feat, num_feat, 1)
        self.spatial_attn_l2 = nn.Conv2d(num_feat * 2, num_feat, 3, 1, 1)
        self.spatial_attn_l3 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
        self.spatial_attn_add1 = nn.Conv2d(num_feat, num_feat, 1)
        self.spatial_attn_add2 = nn.Conv2d(num_feat, num_feat, 1)

        self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
        self.upsample = nn.Upsample(
            scale_factor=2, mode='bilinear', align_corners=False)

    def forward(self, aligned_feat):
        """
        Args:
            aligned_feat (Tensor): Aligned features with shape (b, t, c, h, w).
        Returns:
            Tensor: Features after TSA with the shape (b, c, h, w).
        """
        b, t, c, h, w = aligned_feat.size()
        # temporal attention
        embedding_ref = self.temporal_attn1(
            aligned_feat[:, self.center_frame_idx, :, :, :].clone())
        embedding = self.temporal_attn2(aligned_feat.view(-1, c, h, w))
        embedding = embedding.view(b, t, -1, h, w)  # (b, t, c, h, w)

        corr_l = []  # correlation list
        for i in range(t):
            emb_neighbor = embedding[:, i, :, :, :]
            corr = torch.sum(emb_neighbor * embedding_ref, 1)  # (b, h, w)
            corr_l.append(corr.unsqueeze(1))  # (b, 1, h, w)
        corr_prob = torch.sigmoid(torch.cat(corr_l, dim=1))  # (b, t, h, w)
        corr_prob = corr_prob.unsqueeze(2).expand(b, t, c, h, w)
        corr_prob = corr_prob.contiguous().view(b, -1, h, w)  # (b, t*c, h, w)
        aligned_feat = aligned_feat.view(b, -1, h, w) * corr_prob

        # fusion
        feat = self.lrelu(self.feat_fusion(aligned_feat))

        # spatial attention
        attn = self.lrelu(self.spatial_attn1(aligned_feat))
        attn_max = self.max_pool(attn)
        attn_avg = self.avg_pool(attn)
        attn = self.lrelu(
            self.spatial_attn2(torch.cat([attn_max, attn_avg], dim=1)))
        # pyramid levels
        attn_level = self.lrelu(self.spatial_attn_l1(attn))
        attn_max = self.max_pool(attn_level)
        attn_avg = self.avg_pool(attn_level)
        attn_level = self.lrelu(
            self.spatial_attn_l2(torch.cat([attn_max, attn_avg], dim=1)))
        attn_level = self.lrelu(self.spatial_attn_l3(attn_level))
        attn_level = self.upsample(attn_level)

        attn = self.lrelu(self.spatial_attn3(attn)) + attn_level
        attn = self.lrelu(self.spatial_attn4(attn))
        attn = self.upsample(attn)
        attn = self.spatial_attn5(attn)
        attn_add = self.spatial_attn_add2(
            self.lrelu(self.spatial_attn_add1(attn)))
        attn = torch.sigmoid(attn)

        # after initialization, * 2 makes (attn * 2) to be close to 1.
        feat = feat * attn * 2 + attn_add
        return feat