import torch
import torch.nn as nn

# --------------------------------------------
# GAN-RW
# --------------------------------------------
class _DCR_block(nn.Module):
    def __init__(self, channel_in):
        super(_DCR_block, self).__init__()

        self.conv_1 = nn.Conv2d(in_channels=channel_in, out_channels=int(channel_in/2.), kernel_size=3, stride=1, padding=1)
        self.ba1= nn.BatchNorm2d(num_features=int(channel_in/2.),momentum=0.9, eps=1e-04, affine=True)
        self.relu1 = nn.ReLU()
        
        self.conv_2 = nn.Conv2d(in_channels=int(channel_in*3/2.), out_channels=int(channel_in/2.), kernel_size=3, stride=1, padding=1)
        self.ba2= nn.BatchNorm2d(num_features=int(channel_in/2.),momentum=0.9, eps=1e-04, affine=True)
        self.relu2 = nn.ReLU()
        
        self.conv_3 = nn.Conv2d(in_channels=channel_in*2, out_channels=channel_in, kernel_size=3, stride=1, padding=1)
        self.ba3= nn.BatchNorm2d(num_features=int(channel_in),momentum=0.9, eps=1e-04, affine=True)
        self.relu3 = nn.ReLU()

    def forward(self, x):
        
        residual = x

        out=self.conv_1(x)
        out=self.ba1(out)
        out=self.relu1(out)

        conc = torch.cat([x, out], 1)

        out=self.conv_2(conc)
        out=self.ba2(out)
        out=self.relu2(out)

        conc = torch.cat([conc, out], 1)
 
        out=self.conv_3(conc)
        out=self.ba3(out)
        out=self.relu3(out)

        out = torch.add(out, residual)

        return out


class _down(nn.Module):
    def __init__(self, channel_in):
        super(_down, self).__init__()

        self.relu = nn.ReLU()
        self.maxpool = nn.MaxPool2d(2)
        self.conv = nn.Conv2d(in_channels=channel_in, out_channels=2*channel_in, kernel_size=1, stride=1, padding=0)
        self.ba=nn.BatchNorm2d(num_features=2*channel_in,momentum=0.9, eps=1e-04, affine=True)


    def forward(self, x):

        out = self.maxpool(x)
        out=self.conv(out)
        out=self.ba(out)
        out=self.relu(out)

        return out


class _up(nn.Module):
    def __init__(self, channel_in):
        super(_up, self).__init__()

        self.relu = nn.ReLU()
        self.ba= nn.BatchNorm2d(num_features=channel_in,momentum=0.9, eps=1e-04, affine=False)
        self.subpixel = nn.PixelShuffle(2)
        self.conv = nn.Conv2d(in_channels=channel_in, out_channels=channel_in, kernel_size=1, stride=1, padding=0)

    def forward(self, x):

        out=self.conv(x)
        out=self.ba(out)
        out=self.relu(out)
        out = self.subpixel(out)

        return out


class gan_rw(nn.Module):
    def __init__(self, in_nc=1, out_nc=1, nc=64, nb=17, act_mode='BR'):
        super(gan_rw, self).__init__()
        self.conv_i = nn.Conv2d(in_channels=1, out_channels=128, kernel_size=1, stride=1, padding=0)
        self.relu1 = nn.ReLU()
        self.DCR_block11 = self.make_layer(_DCR_block, 128)
        self.DCR_block12 = self.make_layer(_DCR_block, 128)
        self.down1 = self.make_layer(_down, 128)
        self.DCR_block21 = self.make_layer(_DCR_block, 256)
        self.DCR_block22 = self.make_layer(_DCR_block, 256)
        self.down2 = self.make_layer(_down, 256)
        self.DCR_block31 = self.make_layer(_DCR_block, 512)
        self.DCR_block32 = self.make_layer(_DCR_block, 512)
        self.down3 = self.make_layer(_down, 512)
        self.DCR_block41 = self.make_layer(_DCR_block, 1024)
        self.DCR_block42 = self.make_layer(_DCR_block, 1024)
        self.up3 = self.make_layer(_up, 2048)
        self.DCR_block33 = self.make_layer(_DCR_block, 1024)
        self.DCR_block34 = self.make_layer(_DCR_block, 1024)
        self.up2 = self.make_layer(_up, 1024)
        self.DCR_block23 = self.make_layer(_DCR_block, 512)
        self.DCR_block24 = self.make_layer(_DCR_block, 512)
        self.up1 = self.make_layer(_up, 512)
        self.DCR_block13 = self.make_layer(_DCR_block, 256)
        self.DCR_block14 = self.make_layer(_DCR_block, 256)
        self.conv_f = nn.Conv2d(in_channels=256, out_channels=1, kernel_size=1, stride=1, padding=0)
        self.relu2 = nn.ReLU()
    
    def make_layer(self, block, channel_in):
        layers = []
        layers.append(block(channel_in))
        return nn.Sequential(*layers)

    def forward(self, x):
        
        residual = x

        out = self.relu1(self.conv_i(x))

        out = self.DCR_block11(out)

        conc1 = self.DCR_block12(out)

        out = self.down1(conc1)

        out = self.DCR_block21(out)

        conc2 = self.DCR_block22(out)

        out = self.down2(conc2)

        out = self.DCR_block31(out)

        conc3 = self.DCR_block32(out)

        conc4 = self.down3(conc3)

        out = self.DCR_block41(conc4)

        out = self.DCR_block42(out)

        out = torch.cat([conc4, out], 1)

        out = self.up3(out)

        out = torch.cat([conc3, out], 1)

        out = self.DCR_block33(out)

        out = self.DCR_block34(out)

        out = self.up2(out)

        out = torch.cat([conc2, out], 1)

        out = self.DCR_block23(out)

        out = self.DCR_block24(out)

        out = self.up1(out)

        out = torch.cat([conc1, out], 1)

        out = self.DCR_block13(out)

        out = self.DCR_block14(out)

        out = self.relu2(self.conv_f(out))

        out = torch.add(residual, out)

        return out



        




