import torch


def calc_gradient_penalty(netD, real_data, fake_data, center=0):
    """Computes ell_2 gradient penalty at interpolates between real and fake data.
    
    Adapted from https://github.com/caogang/wgan-gp

    Args:
        netD: The discriminator network whose gradient will be regularized.
        real_data: Real input data.
        fake_data: Fake input data.
        center: Desired norm of discriminator gradients.

    Returns:
        Gradient penalty loss.
    """

    dev = real_data.device

    length = min(len(real_data),len(fake_data))
    real_data = real_data[:length]
    fake_data = fake_data[:length]

    alpha = torch.rand(length, 1)
    alpha = alpha.expand(real_data.size())
    alpha = alpha.to(dev)

    interpolates = alpha * real_data + ((1 - alpha) * fake_data)

    interpolates = interpolates.to(dev)
    interpolates = torch.autograd.Variable(interpolates, requires_grad=True)

    disc_interpolates = netD(interpolates)

    gradients = torch.autograd.grad(outputs=disc_interpolates, inputs=interpolates,
                            grad_outputs=torch.ones(disc_interpolates.size()).to(dev),
                            create_graph=True, retain_graph=True, only_inputs=True)[0]

    gradient_penalty = ((gradients.norm(2, dim=1) - center) ** 2).mean()
    return gradient_penalty