import torch.nn as nn
import torch.nn.functional as F
import torch
import functools


class ModuleBase(nn.Module):
    def __init__(self):
        super(ModuleBase, self).__init__()
        self._name = 'BaseNetwork'

    @property
    def name(self):
        return self._name

    def init_weights(self):
        self.apply(self._weights_init_fn)

    def _weights_init_fn(self, m):
        classname = m.__class__.__name__
        if classname.find('Conv') != -1:
            if hasattr(m, 'weight'):
                torch.nn.init.xavier_uniform(m.weight)
                #m.weight.data.normal_(0.0, 0.02)
            if hasattr(m, 'bias'):
                m.bias.data.fill_(0)
        elif classname.find('BatchNorm2d') != -1:
            m.weight.data.normal_(1.0, 0.02)
            if hasattr(m.bias, 'data'):
                m.bias.data.fill_(0)
        elif type(m) == nn.Linear:
            torch.nn.init.xavier_uniform(m.weight)
            if hasattr(m.bias, 'data'):
                m.bias.data.fill_(0.01)

    def _get_norm_layer(self, norm_type='batch'):
        if norm_type == 'batch':
            norm_layer = functools.partial(nn.BatchNorm2d, affine=True)
        elif norm_type == 'instance':
            norm_layer = functools.partial(nn.InstanceNorm2d, affine=False)
        elif norm_type =='batchnorm2d':
            norm_layer = nn.BatchNorm2d
        else:
            raise NotImplementedError('normalization layer [%s] is not found' % norm_type)

        return norm_layer
    
    
def train_D_wasserstein_gp(g, d, x_real, au, lambda_cl, lambda_cyc, data_loader,device,d_optimizer):
    for p in d.parameters():
        p.requires_grad = True
    batch_size = x_real.shape[0]
    #
    dtype = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor 
    des_au_1 = torch.tensor(data_loader.gen_rand_cond(batch_size=batch_size)).to(device).type(dtype)
    #alist = torch.tensor(data_loader.gen_rand_cond_for_binary_au(au.cpu().detach().numpy())).to(device).type(dtype)
    ##
    z = g.encode(x_real)
    ##
    fakes_1 = g.translate_decode(z,des_au_1)
    #img_rec = g.translate_decode(z,au)
    #
    d_adv_logits_true, d_reg_true = d(x_real)
    d_adv_logits_fake, d_reg_fake = d(fakes_1)
    #
    alpha = torch.rand(batch_size,1,1,1,device=device)
    x_gp = alpha*fakes_1+(1-alpha)*x_real
    d_gp,_ = d(x_gp)
    grad = torch.autograd.grad(d_gp.sum(), x_gp, create_graph=True)
    grad_norm = grad[0].reshape(batch_size, -1).norm(dim=1)
    #
    #d_cl_loss = F.l1_loss(d_reg_true, au)
    d_cl_loss = F.binary_cross_entropy_with_logits(d_reg_true, au)
    d_adv_loss = -d_adv_logits_true.mean() + d_adv_logits_fake.mean()
    d_loss = d_adv_loss+lambda_cl*d_cl_loss+10*((grad_norm - 1)**2).mean()
    d_loss_dict = {'d_adv_loss': d_adv_loss , "d_cl_loss": d_cl_loss}
    ## opt. discr. 
    d_optimizer.zero_grad()
    d_loss.backward(retain_graph=True)
    d_optimizer.step()
    #
    return d_loss_dict
    
def train_G_wasserstein_gp(g, d, x_real, au, lambda_cl, lambda_cyc, data_loader,device,g_optimizer):
    for p in d.parameters():
        p.requires_grad = False
    batch_size = x_real.shape[0]
    #
    dtype = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor 
    des_au_1 = torch.tensor(data_loader.gen_rand_cond(batch_size=batch_size)).to(device).type(dtype)
    #alist = torch.tensor(data_loader.gen_rand_cond_for_binary_au(au.cpu().detach().numpy())).to(device).type(dtype)
    ##
    z = g.encode(x_real)
    ##
    fakes_1 = g.translate_decode(z,des_au_1)
    #fakes_1 = g.translate_decode(z,des_au_1)
    img_rec = g.translate_decode(z,au)
    #
    #d_adv_logits_true, d_reg_true = d(x_real)
    d_adv_logits_fake, d_reg_fake = d(fakes_1)
    #
    g_adv_loss = -d_adv_logits_fake.mean()
    #g_cl_loss = F.l1_loss(d_reg_fake, des_au_1)
    g_cl_loss = F.binary_cross_entropy_with_logits(d_reg_fake, des_au_1)
    rec_loss = F.l1_loss(img_rec, x_real)
    g_loss = g_adv_loss + lambda_cl*g_cl_loss + lambda_cyc*rec_loss
    g_loss_dict = {'g_adv_loss': g_adv_loss , "g_cl_loss": g_cl_loss, "rec_loss": rec_loss}
    #
    g_optimizer.zero_grad()
    g_loss.backward()
    g_optimizer.step()
    #
    return g_loss_dict
    
    
def loss_nonsaturating(g, d, x_real, au, lambda_cl, lambda_cyc, data_loader,device,train_generator=True):
    batch_size = x_real.shape[0]
    ## TODO repeat 
    dtype = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor 
    des_au_1 = torch.tensor(data_loader.gen_rand_cond(batch_size=batch_size)).to(device).type(dtype)
    ##
    z = g.encode(x_real)
    ##
    fakes_1 = g.translate_decode(z,des_au_1)
    img_rec = g.translate_decode(z,au)
    ## d adv loss 
    d_adv_logits_true, d_reg_true = d(x_real)
    d_adv_logits_fake, d_reg_fake = d(fakes_1)
    d_adv_loss = -F.logsigmoid(d_adv_logits_true).mean() - F.logsigmoid(-d_adv_logits_fake).mean()
    ## d cond reg. loss 
    d_cl_loss = F.l1_loss(d_reg_true, au)
    
    ## d_loss 
    d_loss = d_adv_loss + lambda_cl*d_cl_loss
    d_loss_dict = {'d_adv_loss': d_adv_loss , "d_cl_loss": d_cl_loss}
    
    ## g_loss 
    if train_generator:
        g_adv_loss = -F.logsigmoid(d_adv_logits_fake).mean()
        g_cl_loss = F.l1_loss(d_reg_fake, des_au_1)
        ## rec. loss 
        rec_loss = F.l1_loss(img_rec, x_real)
        g_loss = g_adv_loss + lambda_cl*g_cl_loss + lambda_cyc*rec_loss
        g_loss_dict = {'g_adv_loss': g_adv_loss , "g_cl_loss": g_cl_loss, "rec_loss": rec_loss}
        ## 
        return d_loss , d_loss_dict , g_loss, g_loss_dict
    else:
        return d_loss, d_loss_dict, None, None 

class Generator(ModuleBase):
    def __init__(self, img_shape,gf,AU_num,num_layers=4,f_size=6,tranform_layer=False,res_blocks=1):
        super(Generator, self).__init__()
        
        channels, height, width = img_shape
        
        self.res_blocks = res_blocks
        
        
        class Conv2dBlock(ModuleBase):
            def __init__(self, in_filters, out_filters, f_size=6, normalize=True,stride=2,padding=0):
                super(Conv2dBlock, self).__init__()
                layers = [nn.Conv2d(in_filters, out_filters, f_size, stride=stride, padding=padding)]
                if normalize:
                    layers.append(nn.InstanceNorm2d(out_filters,affine=True))
                layers.append(nn.LeakyReLU(0.2, inplace=True))
                self.layers = nn.Sequential(*layers)
            
            def forward(self, x):
                return self.layers(x)
        
        def conv2d(in_filters, out_filters, f_size=6, normalize=True,stride=2,padding=0):
            return Conv2dBlock(in_filters, out_filters, f_size, normalize,stride,padding)
        
        class Deconv2dBlock(ModuleBase):
            def __init__(self, in_filters, out_filters, f_size=6, stride=2,padding=0,normalize=True,bias=True,output_padding=0):
                super(Deconv2dBlock, self).__init__()
                layers = [nn.ConvTranspose2d(in_filters, out_filters, kernel_size=f_size, stride=stride, padding=padding, output_padding=output_padding, bias=bias)]
                if normalize:
                    layers.append(nn.InstanceNorm2d(out_filters, affine=True))
                layers.append(nn.LeakyReLU(0.2, inplace=True))
                self.layers = nn.Sequential(*layers)
            
            def forward(self, x):
                return self.layers(x)
        
        def deconv2d(in_filters, out_filters, f_size=6, stride=2,padding=0,normalize=True,bias=True,output_padding=0):
             return Deconv2dBlock(in_filters, out_filters, f_size, stride,padding,normalize,bias,output_padding)
            
        ### ENCODER
        layers = [
        conv2d(channels, gf, f_size=4,stride=2),
            conv2d(gf, gf*2,f_size=4),
            conv2d(gf*2, gf*4,f_size=4),
            conv2d(gf*4, gf*8,f_size=4)
        ]
        self.enc_layers = nn.ModuleList(layers)
        
        ### 1x1 conv 
        self.conv_1_1 = nn.Conv2d(gf*8+AU_num, gf*8, kernel_size=1, stride=1, padding=0)
        
        ### DECODER
        layers = [
            deconv2d(gf*8, gf*4, f_size=4), ## residual block
            deconv2d(gf*8, gf*2, f_size=4), 
            deconv2d(gf*2, gf, f_size=4,output_padding=1),    
            deconv2d(gf, channels,f_size=4,stride=2,output_padding=0)
        ]
        self.dec_layers = nn.ModuleList(layers)
        
       
    def encode(self, x):
        z = x
        zs = []
        for layer in self.enc_layers:
            z = layer(z)
            zs.append(z)
        #print("last z:",z.shape)
        return zs
    
    def translate_decode_old(self, z,au_target):
        _au_target = au_target.view(au_target.size(0), -1, 1, 1)
        _z = z[-1]
        #print("_au_target",_au_target.shape)
        #print("_z",_z.shape,len(z))
        res_block = torch.cat([_z, _au_target], dim=1)
        res_block = self.conv_1_1(res_block)
        for i,layer in enumerate(self.dec_layers):
            res_block = layer(res_block)
            if len(z) > 0 and i < self.res_blocks:
                _z = z[-2-i]
                #print("res_block",res_block.shape)
                #print("_z",_z.shape,len(z))
                res_block = torch.cat([res_block, _z], dim=1)
        return res_block
    
    def translate_decode(self, z,au_target):
        _au_target = au_target.view(au_target.size(0), -1, 1, 1).repeat(1, 1, 5, 5)
        _z = z[-1]
        #print("_au_target",_au_target.shape)
        #print("_z",_z.shape,len(z))
        res_block = torch.cat([_z, _au_target], dim=1)
        #print("- before 1x1 res_block",res_block.shape)
        res_block = self.conv_1_1(res_block)
        #print("- after 1x1 res_block",res_block.shape)
        for i,layer in enumerate(self.dec_layers):
            res_block = layer(res_block)
            if len(z) > 0 and i < self.res_blocks:
                _z = z[-2-i]
                #print(i,"res_block",res_block.shape)
                #print(i,"_z",_z.shape,len(z))
                res_block = torch.cat([res_block, _z], dim=1)
            #print(i," -- final -- res_block",res_block.shape)
        return res_block
    
    def forward(self, img,au_target):
        return self.translate_decode(self.encode(img),au_target)
        
        
            
class Discriminator(ModuleBase):
    def __init__(self, img_shape,df,AU_num):
        super(Discriminator, self).__init__()

        channels, height, width = img_shape

        # Calculate output shape of image discriminator (PatchGAN)
        self.output_shape_PatchGAN = (1, height // 2 ** 4, width // 2 ** 4)

        def discriminator_block(in_filters, out_filters, f_size=4, normalize=True,stride=2):
            """Returns downsampling layers of each discriminator block"""
            layers = [nn.Conv2d(in_filters, out_filters, f_size, stride=stride, padding=0)]
            if normalize:
                layers.append(nn.InstanceNorm2d(out_filters))
            layers.append(nn.LeakyReLU(0.2, inplace=True))
            return layers

        self.net = nn.Sequential(
            *discriminator_block(channels, df, normalize=False,f_size=6,stride=3),
            *discriminator_block(64, 128,f_size=6),
            *discriminator_block(128, 256,f_size=6),
            *discriminator_block(256, 512,f_size=6),
            #nn.ZeroPad2d((1, 0, 1, 0)),
            ##nn.Conv2d(512, 1, 4, padding=1)
        )
        self.gan_task = nn.Sequential(
            torch.nn.Linear(512,512),
            #nn.LeakyReLU(0.2, inplace=True),
            nn.ReLU(),
            torch.nn.Linear(512,1)
        )
        self.au_reg_task = nn.Sequential(
            torch.nn.Linear(512,512),
            #nn.LeakyReLU(0.2, inplace=True),
            nn.ReLU(),
            torch.nn.Linear(512,AU_num)
        )

    def forward(self, img):
        #print("img:",img.shape)
        common_pass = self.net.forward(img).view(img.shape[0],-1)
        #print("common_pass::",common_pass.shape)
        gan_out = self.gan_task(common_pass)
        au_reg_out = self.au_reg_task(common_pass)
        return gan_out , au_reg_out

if __name__ == '__main__':
    d = Discriminator(img_shape=(3,112,112),df=64,AU_num=17)
    print("******** Discriminator/Classifier ********")
    print(d)
    b = torch.rand(5,3,112,112)
    gan_out , au_reg_out  = d.forward(b)
    print("gan_out::",gan_out.shape)
    print(gan_out)
    print("au_reg_out::",au_reg_out.shape)
    g = Generator(img_shape=(3,112,112),gf=64,AU_num=17)
    print("******** Generator ********")
    print(g)
    b = torch.rand(5,3,112,112)
    au_target = torch.rand(5,17)
    a = g(b,au_target)
    print("a",a.shape)
