import torch


def compute_gradient_penalty(network, samples_a, samples_b, device):
        # https://github.com/eriklindernoren/PyTorch-GAN/blob/master/implementations/wgan_gp/wgan_gp.py

        """Calculates the gradient penalty loss for WGAN GP"""
        # Random weight term for interpolation between real and fake samples
        batch_size = samples_a.size(0)
        alpha = torch.rand_like(samples_a)
        # Get random interpolation between real and fake samples
        interpolates = (alpha * samples_a + ((1 - alpha) * samples_b))
        interpolated_obs = torch.autograd.Variable(interpolates, requires_grad=True)

        d_interpolates = network(interpolated_obs)
        grad = torch.ones(d_interpolates.size(), requires_grad=False).to(device)

        # Get gradient w.r.t. interpolates
        gradients = torch.autograd.grad(
            outputs=d_interpolates,
            inputs=interpolated_obs,
            grad_outputs=grad,
            create_graph=True,
            retain_graph=True,
            only_inputs=True,
        )[0]
        gradients = gradients.view(int(batch_size), -1)
        gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean()
        
        return gradient_penalty