import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
import torch.nn.functional as F


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

    def init_weights(self, init_type='normal', gain=0.02):
        '''
        initialize network's weights
        init_type: normal | xavier | kaiming | orthogonal
        https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/9451e70673400885567d08a9e97ade2524c700d0/models/networks.py#L39
        '''

        def init_func(m):
            classname = m.__class__.__name__
            if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
                if init_type == 'normal':
                    nn.init.normal_(m.weight.data, 0.0, gain)
                elif init_type == 'xavier':
                    nn.init.xavier_normal_(m.weight.data, gain=gain)
                elif init_type == 'kaiming':
                    nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
                elif init_type == 'orthogonal':
                    nn.init.orthogonal_(m.weight.data, gain=gain)

                if hasattr(m, 'bias') and m.bias is not None:
                    nn.init.constant_(m.bias.data, 0.0)

            elif classname.find('BatchNorm2d') != -1:
                nn.init.normal_(m.weight.data, 1.0, gain)
                nn.init.constant_(m.bias.data, 0.0)

        self.apply(init_func)

def weights_init(init_type='gaussian'):
    def init_fun(m):
        classname = m.__class__.__name__
        if (classname.find('Conv') == 0 or classname.find(
                'Linear') == 0) and hasattr(m, 'weight'):
            if init_type == 'gaussian':
                nn.init.normal_(m.weight, 0.0, 0.02)
            elif init_type == 'xavier':
                nn.init.xavier_normal_(m.weight, gain=math.sqrt(2))
            elif init_type == 'kaiming':
                nn.init.kaiming_normal_(m.weight, a=0, mode='fan_in')
            elif init_type == 'orthogonal':
                nn.init.orthogonal_(m.weight, gain=math.sqrt(2))
            elif init_type == 'default':
                pass
            else:
                assert 0, "Unsupported initialization: {}".format(init_type)
            if hasattr(m, 'bias') and m.bias is not None:
                nn.init.constant_(m.bias, 0.0)

    return init_fun

class PartialConv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
                 padding=0, dilation=1, groups=1, bias=True):
        super().__init__()
        self.input_conv = nn.Conv2d(in_channels, out_channels, kernel_size,
                                    stride, padding, dilation, groups, bias)
        self.mask_conv = nn.Conv2d(in_channels, out_channels, kernel_size,
                                   stride, padding, dilation, groups, False)
        self.input_conv.apply(weights_init('kaiming'))
        self.slide_winsize = in_channels * kernel_size * kernel_size

        torch.nn.init.constant_(self.mask_conv.weight, 1.0)

        # mask is not updated
        for param in self.mask_conv.parameters():
            param.requires_grad = False

    def forward(self, input, mask):
        # http://masc.cs.gmu.edu/wiki/partialconv
        # C(X) = W^T * X + b, C(0) = b, D(M) = 1 * M + 0 = sum(M)
        # W^T* (M .* X) / sum(M) + b = [C(M .* X) – C(0)] / D(M) + C(0)
        output = self.input_conv(input * mask)
        if self.input_conv.bias is not None:
            output_bias = self.input_conv.bias.view(1, -1, 1, 1).expand_as(
                output)
        else:
            output_bias = torch.zeros_like(output)

        with torch.no_grad():
            output_mask = self.mask_conv(mask)

        no_update_holes = output_mask == 0

        mask_sum = output_mask.masked_fill_(no_update_holes, 1.0)

        output_pre = ((output - output_bias) * self.slide_winsize) / mask_sum + output_bias
        output = output_pre.masked_fill_(no_update_holes, 0.0)

        new_mask = torch.ones_like(output)
        new_mask = new_mask.masked_fill_(no_update_holes, 0.0)

        return output, new_mask


class PCBActiv(nn.Module):
    def __init__(self, in_ch, out_ch, bn=True, sample='none-3', activ='relu',
                 conv_bias=False):
        super().__init__()
        if sample == 'down-5':
            self.conv = PartialConv(in_ch, out_ch, 5, 2, 2, bias=conv_bias)
        elif sample == 'down-7':
            self.conv = PartialConv(in_ch, out_ch, 7, 2, 3, bias=conv_bias)
        elif sample == 'down-3':
            self.conv = PartialConv(in_ch, out_ch, 3, 2, 1, bias=conv_bias)
        else:
            self.conv = PartialConv(in_ch, out_ch, 3, 1, 1, bias=conv_bias)

        if bn:
            self.bn = nn.BatchNorm2d(out_ch)
        if activ == 'relu':
            self.activation = nn.ReLU()
        elif activ == 'leaky':
            self.activation = nn.LeakyReLU(negative_slope=0.2)

    def forward(self, input, input_mask):
        h, h_mask = self.conv(input, input_mask)
        if hasattr(self, 'bn'):
            h = self.bn(h)
        if hasattr(self, 'activation'):
            h = self.activation(h)
        return h, h_mask

class Inpaint_Depth_Net(nn.Module):
    def __init__(self, layer_size=7, upsampling_mode='nearest'):
        super().__init__()
        in_channels = 4
        out_channels = 1
        self.freeze_enc_bn = False
        self.upsampling_mode = upsampling_mode
        self.layer_size = layer_size
        self.enc_1 = PCBActiv(in_channels, 64, bn=False, sample='down-7', conv_bias=True)
        self.enc_2 = PCBActiv(64, 128, sample='down-5', conv_bias=True)
        self.enc_3 = PCBActiv(128, 256, sample='down-5')
        self.enc_4 = PCBActiv(256, 512, sample='down-3')
        for i in range(4, self.layer_size):
            name = 'enc_{:d}'.format(i + 1)
            setattr(self, name, PCBActiv(512, 512, sample='down-3'))

        for i in range(4, self.layer_size):
            name = 'dec_{:d}'.format(i + 1)
            setattr(self, name, PCBActiv(512 + 512, 512, activ='leaky'))
        self.dec_4 = PCBActiv(512 + 256, 256, activ='leaky')
        self.dec_3 = PCBActiv(256 + 128, 128, activ='leaky')
        self.dec_2 = PCBActiv(128 + 64, 64, activ='leaky')
        self.dec_1 = PCBActiv(64 + in_channels, out_channels,
                              bn=False, activ=None, conv_bias=True)
    def add_border(self, input, mask_flag, PCONV=True):
        with torch.no_grad():
            h = input.shape[-2]
            w = input.shape[-1]
            require_len_unit = 2 ** self.layer_size
            residual_h = int(np.ceil(h / float(require_len_unit)) * require_len_unit - h) # + 2*require_len_unit
            residual_w = int(np.ceil(w / float(require_len_unit)) * require_len_unit - w) # + 2*require_len_unit
            enlarge_input = torch.zeros((input.shape[0], input.shape[1], h + residual_h, w + residual_w)).to(input.device)
            if mask_flag:
                if PCONV is False:
                    enlarge_input += 1.0
                enlarge_input = enlarge_input.clamp(0.0, 1.0)
            else:
                enlarge_input[:, 2, ...] = 0.0
            anchor_h = residual_h//2
            anchor_w = residual_w//2
            enlarge_input[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w] = input

        return enlarge_input, [anchor_h, anchor_h+h, anchor_w, anchor_w+w]

    def forward_3P(self, mask, context, depth, edge, unit_length=128, cuda=None):
        with torch.no_grad():
            input = torch.cat((depth, edge, context, mask), dim=1)
            n, c, h, w = input.shape
            residual_h = int(np.ceil(h / float(unit_length)) * unit_length - h)
            residual_w = int(np.ceil(w / float(unit_length)) * unit_length - w)
            anchor_h = residual_h//2
            anchor_w = residual_w//2
            enlarge_input = torch.zeros((n, c, h + residual_h, w + residual_w)).to(cuda)
            enlarge_input[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w] = input
            # enlarge_input[:, 3] = 1. - enlarge_input[:, 3]
            depth_output = self.forward(enlarge_input)
            depth_output = depth_output[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w]
            # import pdb; pdb.set_trace()

        return depth_output

    def forward(self, input_feat, refine_border=False, sample=False, PCONV=True):
        input = input_feat
        input_mask = (input_feat[:, -2:-1] + input_feat[:, -1:]).clamp(0, 1).repeat(1, input.shape[1], 1, 1)

        vis_input = input.cpu().data.numpy()
        vis_input_mask = input_mask.cpu().data.numpy()
        H, W = input.shape[-2:]
        if refine_border is True:
            input, anchor = self.add_border(input, mask_flag=False)
            input_mask, anchor = self.add_border(input_mask, mask_flag=True, PCONV=PCONV)
        h_dict = {}  # for the output of enc_N
        h_mask_dict = {}  # for the output of enc_N
        h_dict['h_0'], h_mask_dict['h_0'] = input, input_mask

        h_key_prev = 'h_0'
        for i in range(1, self.layer_size + 1):
            l_key = 'enc_{:d}'.format(i)
            h_key = 'h_{:d}'.format(i)
            h_dict[h_key], h_mask_dict[h_key] = getattr(self, l_key)(
                h_dict[h_key_prev], h_mask_dict[h_key_prev])
            h_key_prev = h_key

        h_key = 'h_{:d}'.format(self.layer_size)
        h, h_mask = h_dict[h_key], h_mask_dict[h_key]

        for i in range(self.layer_size, 0, -1):
            enc_h_key = 'h_{:d}'.format(i - 1)
            dec_l_key = 'dec_{:d}'.format(i)

            h = F.interpolate(h, scale_factor=2, mode=self.upsampling_mode)
            h_mask = F.interpolate(h_mask, scale_factor=2, mode='nearest')

            h = torch.cat([h, h_dict[enc_h_key]], dim=1)
            h_mask = torch.cat([h_mask, h_mask_dict[enc_h_key]], dim=1)
            h, h_mask = getattr(self, dec_l_key)(h, h_mask)
        output = h
        if refine_border is True:
            h_mask = h_mask[..., anchor[0]:anchor[1], anchor[2]:anchor[3]]
            output = output[..., anchor[0]:anchor[1], anchor[2]:anchor[3]]

        return output

class Inpaint_Edge_Net(BaseNetwork):
    def __init__(self, residual_blocks=8, init_weights=True):
        super(Inpaint_Edge_Net, self).__init__()
        in_channels = 7
        out_channels = 1
        self.encoder = []
        # 0
        self.encoder_0 = nn.Sequential(
                            nn.ReflectionPad2d(3),
                            spectral_norm(nn.Conv2d(in_channels=in_channels, out_channels=64, kernel_size=7, padding=0), True),
                            nn.InstanceNorm2d(64, track_running_stats=False),
                            nn.ReLU(True))
        # 1
        self.encoder_1 = nn.Sequential(
                            spectral_norm(nn.Conv2d(in_channels=64, out_channels=128, kernel_size=4, stride=2, padding=1), True),
                            nn.InstanceNorm2d(128, track_running_stats=False),
                            nn.ReLU(True))
        # 2
        self.encoder_2 = nn.Sequential(
                            spectral_norm(nn.Conv2d(in_channels=128, out_channels=256, kernel_size=4, stride=2, padding=1), True),
                            nn.InstanceNorm2d(256, track_running_stats=False),
                            nn.ReLU(True))
        # 3
        blocks = []
        for _ in range(residual_blocks):
            block = ResnetBlock(256, 2)
            blocks.append(block)

        self.middle = nn.Sequential(*blocks)
        # + 3
        self.decoder_0 = nn.Sequential(
                            spectral_norm(nn.ConvTranspose2d(in_channels=256+256, out_channels=128, kernel_size=4, stride=2, padding=1), True),
                            nn.InstanceNorm2d(128, track_running_stats=False),
                            nn.ReLU(True))
        # + 2
        self.decoder_1 = nn.Sequential(
                            spectral_norm(nn.ConvTranspose2d(in_channels=128+128, out_channels=64, kernel_size=4, stride=2, padding=1), True),
                            nn.InstanceNorm2d(64, track_running_stats=False),
                            nn.ReLU(True))
        # + 1
        self.decoder_2 = nn.Sequential(
                            nn.ReflectionPad2d(3),
                            nn.Conv2d(in_channels=64+64, out_channels=out_channels, kernel_size=7, padding=0),
                            )

        if init_weights:
            self.init_weights()

    def add_border(self, input, channel_pad_1=None):
        h = input.shape[-2]
        w = input.shape[-1]
        require_len_unit = 16
        residual_h = int(np.ceil(h / float(require_len_unit)) * require_len_unit - h) # + 2*require_len_unit
        residual_w = int(np.ceil(w / float(require_len_unit)) * require_len_unit - w) # + 2*require_len_unit
        enlarge_input = torch.zeros((input.shape[0], input.shape[1], h + residual_h, w + residual_w)).to(input.device)
        if channel_pad_1 is not None:
            for channel in channel_pad_1:
                enlarge_input[:, channel] = 1
        anchor_h = residual_h//2
        anchor_w = residual_w//2
        enlarge_input[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w] = input

        return enlarge_input, [anchor_h, anchor_h+h, anchor_w, anchor_w+w]

    def forward_3P(self, mask, context, rgb, disp, edge, unit_length=128, cuda=None):
        with torch.no_grad():
            input = torch.cat((rgb, disp/disp.max(), edge, context, mask), dim=1)
            n, c, h, w = input.shape
            residual_h = int(np.ceil(h / float(unit_length)) * unit_length - h)
            residual_w = int(np.ceil(w / float(unit_length)) * unit_length - w)
            anchor_h = residual_h//2
            anchor_w = residual_w//2
            enlarge_input = torch.zeros((n, c, h + residual_h, w + residual_w)).to(cuda)
            enlarge_input[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w] = input
            edge_output = self.forward(enlarge_input)
            edge_output = edge_output[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w]

        return edge_output

    def forward(self, x, refine_border=False):
        if refine_border:
            x, anchor = self.add_border(x, [5])
        x1 = self.encoder_0(x)
        x2 = self.encoder_1(x1)
        x3 = self.encoder_2(x2)
        x4 = self.middle(x3)
        x5 = self.decoder_0(torch.cat((x4, x3), dim=1))
        x6 = self.decoder_1(torch.cat((x5, x2), dim=1))
        x7 = self.decoder_2(torch.cat((x6, x1), dim=1))
        x = torch.sigmoid(x7)
        if refine_border:
            x = x[..., anchor[0]:anchor[1], anchor[2]:anchor[3]]

        return x

class Inpaint_Color_Net(nn.Module):
    def __init__(self, layer_size=7, upsampling_mode='nearest', add_hole_mask=False, add_two_layer=False, add_border=False):
        super().__init__()
        self.freeze_enc_bn = False
        self.upsampling_mode = upsampling_mode
        self.layer_size = layer_size
        in_channels = 6
        self.enc_1 = PCBActiv(in_channels, 64, bn=False, sample='down-7')
        self.enc_2 = PCBActiv(64, 128, sample='down-5')
        self.enc_3 = PCBActiv(128, 256, sample='down-5')
        self.enc_4 = PCBActiv(256, 512, sample='down-3')
        self.enc_5 = PCBActiv(512, 512, sample='down-3')
        self.enc_6 = PCBActiv(512, 512, sample='down-3')
        self.enc_7 = PCBActiv(512, 512, sample='down-3')

        self.dec_7 = PCBActiv(512+512, 512, activ='leaky')
        self.dec_6 = PCBActiv(512+512, 512, activ='leaky')

        self.dec_5A = PCBActiv(512 + 512, 512, activ='leaky')
        self.dec_4A = PCBActiv(512 + 256, 256, activ='leaky')
        self.dec_3A = PCBActiv(256 + 128, 128, activ='leaky')
        self.dec_2A = PCBActiv(128 + 64, 64, activ='leaky')
        self.dec_1A = PCBActiv(64 + in_channels, 3, bn=False, activ=None, conv_bias=True)
        '''
        self.dec_5B = PCBActiv(512 + 512, 512, activ='leaky')
        self.dec_4B = PCBActiv(512 + 256, 256, activ='leaky')
        self.dec_3B = PCBActiv(256 + 128, 128, activ='leaky')
        self.dec_2B = PCBActiv(128 + 64, 64, activ='leaky')
        self.dec_1B = PCBActiv(64 + 4, 1, bn=False, activ=None, conv_bias=True)
        '''
    def cat(self, A, B):
        return torch.cat((A, B), dim=1)

    def upsample(self, feat, mask):
        feat = F.interpolate(feat, scale_factor=2, mode=self.upsampling_mode)
        mask = F.interpolate(mask, scale_factor=2, mode='nearest')

        return feat, mask

    def forward_3P(self, mask, context, rgb, edge, unit_length=128, cuda=None):
        with torch.no_grad():
            input = torch.cat((rgb, edge, context, mask), dim=1)
            n, c, h, w = input.shape
            residual_h = int(np.ceil(h / float(unit_length)) * unit_length - h) # + 128
            residual_w = int(np.ceil(w / float(unit_length)) * unit_length - w) # + 256
            anchor_h = residual_h//2
            anchor_w = residual_w//2
            enlarge_input = torch.zeros((n, c, h + residual_h, w + residual_w)).to(cuda)
            enlarge_input[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w] = input
            # enlarge_input[:, 3] = 1. - enlarge_input[:, 3]
            enlarge_input = enlarge_input.to(cuda)
            rgb_output = self.forward(enlarge_input)
            rgb_output = rgb_output[..., anchor_h:anchor_h+h, anchor_w:anchor_w+w]

        return rgb_output

    def forward(self, input, add_border=False):
        input_mask = (input[:, -2:-1] + input[:, -1:]).clamp(0, 1)
        H, W = input.shape[-2:]
        f_0, h_0 = input, input_mask.repeat((1,input.shape[1],1,1))
        f_1, h_1 = self.enc_1(f_0, h_0)
        f_2, h_2 = self.enc_2(f_1, h_1)
        f_3, h_3 = self.enc_3(f_2, h_2)
        f_4, h_4 = self.enc_4(f_3, h_3)
        f_5, h_5 = self.enc_5(f_4, h_4)
        f_6, h_6 = self.enc_6(f_5, h_5)
        f_7, h_7 = self.enc_7(f_6, h_6)

        o_7, k_7 = self.upsample(f_7, h_7)
        o_6, k_6 = self.dec_7(self.cat(o_7, f_6), self.cat(k_7, h_6))
        o_6, k_6 = self.upsample(o_6, k_6)
        o_5, k_5 = self.dec_6(self.cat(o_6, f_5), self.cat(k_6, h_5))
        o_5, k_5 = self.upsample(o_5, k_5)
        o_5A, k_5A = o_5, k_5
        o_5B, k_5B = o_5, k_5
        ###############
        o_4A, k_4A = self.dec_5A(self.cat(o_5A, f_4), self.cat(k_5A, h_4))
        o_4A, k_4A = self.upsample(o_4A, k_4A)
        o_3A, k_3A = self.dec_4A(self.cat(o_4A, f_3), self.cat(k_4A, h_3))
        o_3A, k_3A = self.upsample(o_3A, k_3A)
        o_2A, k_2A = self.dec_3A(self.cat(o_3A, f_2), self.cat(k_3A, h_2))
        o_2A, k_2A = self.upsample(o_2A, k_2A)
        o_1A, k_1A = self.dec_2A(self.cat(o_2A, f_1), self.cat(k_2A, h_1))
        o_1A, k_1A = self.upsample(o_1A, k_1A)
        o_0A, k_0A = self.dec_1A(self.cat(o_1A, f_0), self.cat(k_1A, h_0))

        return torch.sigmoid(o_0A)

    def train(self, mode=True):
        """
        Override the default train() to freeze the BN parameters
        """
        super().train(mode)
        if self.freeze_enc_bn:
            for name, module in self.named_modules():
                if isinstance(module, nn.BatchNorm2d) and 'enc' in name:
                    module.eval()

class Discriminator(BaseNetwork):
    def __init__(self, use_sigmoid=True, use_spectral_norm=True, init_weights=True, in_channels=None):
        super(Discriminator, self).__init__()
        self.use_sigmoid = use_sigmoid
        self.conv1 = self.features = nn.Sequential(
            spectral_norm(nn.Conv2d(in_channels=in_channels, out_channels=64, kernel_size=4, stride=2, padding=1, bias=not use_spectral_norm), use_spectral_norm),
            nn.LeakyReLU(0.2, inplace=True),
        )

        self.conv2 = nn.Sequential(
            spectral_norm(nn.Conv2d(in_channels=64, out_channels=128, kernel_size=4, stride=2, padding=1, bias=not use_spectral_norm), use_spectral_norm),
            nn.LeakyReLU(0.2, inplace=True),
        )

        self.conv3 = nn.Sequential(
            spectral_norm(nn.Conv2d(in_channels=128, out_channels=256, kernel_size=4, stride=2, padding=1, bias=not use_spectral_norm), use_spectral_norm),
            nn.LeakyReLU(0.2, inplace=True),
        )

        self.conv4 = nn.Sequential(
            spectral_norm(nn.Conv2d(in_channels=256, out_channels=512, kernel_size=4, stride=1, padding=1, bias=not use_spectral_norm), use_spectral_norm),
            nn.LeakyReLU(0.2, inplace=True),
        )

        self.conv5 = nn.Sequential(
            spectral_norm(nn.Conv2d(in_channels=512, out_channels=1, kernel_size=4, stride=1, padding=1, bias=not use_spectral_norm), use_spectral_norm),
        )

        if init_weights:
            self.init_weights()

    def forward(self, x):
        conv1 = self.conv1(x)
        conv2 = self.conv2(conv1)
        conv3 = self.conv3(conv2)
        conv4 = self.conv4(conv3)
        conv5 = self.conv5(conv4)

        outputs = conv5
        if self.use_sigmoid:
            outputs = torch.sigmoid(conv5)

        return outputs, [conv1, conv2, conv3, conv4, conv5]

class ResnetBlock(nn.Module):
    def __init__(self, dim, dilation=1):
        super(ResnetBlock, self).__init__()
        self.conv_block = nn.Sequential(
            nn.ReflectionPad2d(dilation),
            spectral_norm(nn.Conv2d(in_channels=dim, out_channels=dim, kernel_size=3, padding=0, dilation=dilation, bias=not True), True),
            nn.InstanceNorm2d(dim, track_running_stats=False),
            nn.LeakyReLU(negative_slope=0.2),

            nn.ReflectionPad2d(1),
            spectral_norm(nn.Conv2d(in_channels=dim, out_channels=dim, kernel_size=3, padding=0, dilation=1, bias=not True), True),
            nn.InstanceNorm2d(dim, track_running_stats=False),
        )

    def forward(self, x):
        out = x + self.conv_block(x)

        # Remove ReLU at the end of the residual block
        # http://torch.ch/blog/2016/02/04/resnets.html

        return out


def spectral_norm(module, mode=True):
    if mode:
        return nn.utils.spectral_norm(module)

    return module