import torch

from deepdraken.image_generation.gans.dcgan import DCGAN

class WGAN(DCGAN):

    def __init__(self, net_G, net_D, device='cpu', gpu_ids=None):
        super().__init__(net_G, net_D, device, gpu_ids)

    def backward_G(self):
        
        self.loss_G = -torch.mean(self.net_D(self.fake_images))
        self.loss_G.backward()

    def backward_D(self):
        
        self.loss_D_real = -torch.mean(self.net_D(self.real_images))
        self.loss_D_fake = torch.mean(self.net_D(self.fake_images.detach()))
        self.loss_D = self.loss_D_real + self.loss_D_fake # TODO check if division by 2 is to be done or not
        self.loss_D.backward()
    
    def optimize_parameters(self) -> None:
         
        data = super().optimize_parameters()
        for p in self.net_D.parameters():
            p.data.clamp_(self.lower_clip, self.upper_clip)

        return data

    def compile(self, optim_G, optim_D):
        self.optimizer_G = optim_G
        self.optimizer_D = optim_D

    def fit(self, data, epochs, batch_size=64, shuffle=True, transform=None, clip_value=0.01):

        if isinstance(clip_value, float):
            self.upper_clip = abs(clip_value)
            self.lower_clip = -self.upper_clip
        if isinstance(clip_value, tuple):
            self.upper_clip, self.lower_clip == clip_value

        return super().fit(data, epochs, batch_size, shuffle, transform)

class WGANGP(WGAN):

    def __init__(self, net_G, net_D, device, is_train=True, gpu_ids=None):
        super().__init__(net_G, net_D, device, is_train, gpu_ids)

    def compute_gradient_penalty(self, real_samples, fake_samples):
        
        """
        Calculates the gradient penalty for WGAN GP
        
        :param real_samples: real images from the dataset
        :param fake_samples: images generated by the generator
        :return: gradient penalty for the model
        """
        
        # Random weight term for interpolation between real and fake samples
        # Random numbers from a uniform distribution on the interval [0, 1)[0,1)
        alpha = torch.rand((self.batch_size, 1, 1, 1), device=self.device)
        
        # Get random interpolation between real and fake samples
        interpolates = (alpha * real_samples + ((1 - alpha) * fake_samples)).requires_grad_(True)
        d_interpolates = self.disc(interpolates)
        fake = torch.ones((self.batch_size, 1)).to(self.device)
        
        # Get gradient w.r.t. interpolates
        gradients = torch.autograd.grad( outputs=d_interpolates,
                                         inputs=interpolates,
                                         grad_outputs=fake,
                                         create_graph=True,
                                         retain_graph=True,
                                         only_inputs=True, )[0]

        gradients = gradients.view(gradients.size(0), -1)
        gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean()
        return gradient_penalty

    def backward_D(self):
        
        self.loss_D_real = -torch.mean(self.net_D(self.real_images))
        self.loss_D_fake = torch.mean(self.net_D(self.fake_images.detach()))

        if self.lambda_gp == 0:
            self.loss_D = self.loss_D_real + self.loss_D_fake
        else:
            gradient_penalty = self.compute_gradient_penalty(self.real_images, self.fake_images)
            self.loss_D = self.loss_D_real + self.loss_D_fake + self.lambda_gp * gradient_penalty
        
        self.loss_D.backward()

    def optimize_parameters(self) -> None:
        return super(WGAN, self).optimize_parameters()

    def fit(self, data, epochs, batch_size=64, shuffle=True, transform=None, lambda_gp=10):

        self.lambda_gp = lambda_gp
        return super(WGAN, self).fit(data, epochs, batch_size, shuffle, transform)
        
        

# TODO loss vaue is comming negative, check about it
# TODO loss must be 1L continuous and thus weight clip or gradient penalty is must
# TODO weight cliping is applied only on the critic