import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable

# trade loss proposed at ICML2019

def div_for_torch(a , b):
    """
    division for torch >=1.0.1
    :param a: tensor, shape[a.shape[0],a.shape[1],a.shape[2]] eg.[3,224,224]
    :param b: tensor, shape[b.shape[0]] eg.[3]
    :return:a/b
    """
    assert a.shape[0] == b.shape[0]
    a_trans = a.permute(1, 2, 0)
    c = a_trans / b
    return c.permute(2, 0, 1)


def squared_l2_norm(x):
    flattened = x.view(x.shape[0], -1)
    #print('flattened:',flattened.shape)
    return (flattened ** 2).sum(1)


def l2_norm(x):
    return squared_l2_norm(x).sqrt()


def trades_loss(model,
                x_natural,
                y,
                optimizer,
                step_size=0.003,
                epsilon=0.031,
                perturb_steps=10,
                beta=1.0,
                distance='l_2',
                device='cpu'):
    # define KL-loss
    criterion_kl = nn.KLDivLoss(size_average=False)
    model.eval()
    batch_size = len(x_natural)
    # generate adversarial example
    x_adv = x_natural.detach() + 0.001 * torch.randn(x_natural.shape).to(device).detach()
    if distance == 'l_inf':
        for _ in range(perturb_steps):
            x_adv.requires_grad_()
            with torch.enable_grad():
                loss_kl = criterion_kl(F.log_softmax(model(x_adv), dim=1),
                                       F.softmax(model(x_natural), dim=1))
            grad = torch.autograd.grad(loss_kl, [x_adv])[0]
            x_adv = x_adv.detach() + step_size * torch.sign(grad.detach())
            x_adv = torch.min(torch.max(x_adv, x_natural - epsilon), x_natural + epsilon)
            x_adv = torch.clamp(x_adv, 0.0, 1.0)
    elif distance == 'l_2':
        for _ in range(perturb_steps):
            x_adv.requires_grad_()
            with torch.enable_grad():
                loss_kl = criterion_kl(F.log_softmax(model(x_adv), dim=1),
                                       F.softmax(model(x_natural), dim=1))
            grad = torch.autograd.grad(loss_kl, [x_adv])[0]
            for idx_batch in range(batch_size):
                grad_idx = grad[idx_batch]
                grad_idx_norm = l2_norm(grad_idx)
                #print('grad_idx_norm:', grad_idx_norm.shape)
                #grad_idx /= (grad_idx_norm + 1e-8)
                grad_idx = div_for_torch(grad_idx,(grad_idx_norm + 1e-8))
                #print('grad_idx:', grad_idx.shape)
                x_adv[idx_batch] = x_adv[idx_batch].detach() + step_size * grad_idx
                eta_x_adv = x_adv[idx_batch] - x_natural[idx_batch]
                norm_eta = l2_norm(eta_x_adv)
                #print(norm_eta, torch.mean(norm_eta))
                if torch.mean(norm_eta) > epsilon:
                    #eta_x_adv = eta_x_adv * epsilon / l2_norm(eta_x_adv)
                    eta_x_adv = div_for_torch((eta_x_adv * epsilon), l2_norm(eta_x_adv))
                x_adv[idx_batch] = x_natural[idx_batch] + eta_x_adv
            x_adv = torch.clamp(x_adv, 0.0, 1.0)
    else:
        x_adv = torch.clamp(x_adv, 0.0, 1.0)
    model.train()

    x_adv = Variable(torch.clamp(x_adv, 0.0, 1.0), requires_grad=False)
    # zero gradient
    optimizer.zero_grad()
    # calculate robust loss
    logits = model(x_natural)
    loss_natural = F.cross_entropy(logits, y)
    loss_robust = (1.0 / batch_size) * criterion_kl(F.log_softmax(model(x_adv), dim=1),
                                                    F.softmax(model(x_natural), dim=1))
    loss = loss_natural + beta * loss_robust
    return loss
