# Copyright (c) Facebook, Inc. and its affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of https://github.com/facebookresearch/AdversarialAndDimensionality
#


import math
import time
import numpy as np
import scipy.stats as st

import torch
from torch.autograd import grad
from torch.linalg import norm

from torch.autograd.functional import hessian as hess_fn

def classify(net, x, x_adv=None, means=None, stds=None, top_k=None):
    device = next(net.parameters()).device

    if means is not None and stds is not None:
        x = x.to(device)
        means = torch.FloatTensor(means).to(device)
        stds = torch.FloatTensor(stds).to(device)
        x = ((x - means) / stds).detach()
    else:
        x = x.detach()
        if x_adv is not None:
            x_adv = x_adv.detach()
    x.requires_grad = True
    y = net(x)
    if x_adv is not None:
        x_adv.requires_grad = True
        g = [(a, norm(grad(y[a].sum(), x, retain_graph=True)[0].view(x.size(0), -1)-
                      grad(y[a].sum(), x_adv, retain_graph=True)[0].view(x_adv.size(0), -1)).item()) for a in y.keys() if a != 'rep']
    else:
        g1 = [(a, norm(grad(y[a].sum(), x, retain_graph=True)[0].view(x.size(0), -1), ord=1).item()) for a in y.keys()  if a != 'rep']
        g2 = [(a, norm(grad(y[a].sum(), x, retain_graph=True)[0].view(x.size(0), -1)).item()) for a in y.keys()  if a != 'rep']
        ginf = [(a, norm(grad(y[a].sum(), x, retain_graph=True)[0].view(x.size(0), -1), ord=float("inf")).item()) for a in y.keys()  if a != 'rep']

        g = [grad(y[a].sum(), x, retain_graph=True)[0].view(x.size(0), -1) for a in y.keys() if a != 'rep']
        g1_all = norm(torch.stack(g, dim=0).sum(dim=0), ord=1).item()
        g2_all = norm(torch.stack(g, dim=0).sum(dim=0)).item()
        ginf_all = norm(torch.stack(g, dim=0).sum(dim=0), ord=float("inf")).item()

        #g = [(a, grad(y[a].sum(), x, retain_graph=True)[0].view(x.size(0), -1).norm().item()) for a in y.keys()]
    if top_k is None:
        return y , dict(g1),dict(g2),dict(ginf), g1_all, g2_all, ginf_all
    else:
        _, top_indices = y.data.cpu().view(-1).topk(2)
        return top_indices[0].item(), dict(g1),dict(g2),dict(ginf), g1_all, g2_all, ginf_all


def conf95(a):
    return st.t.interval(
        0.95, len(a) - 1, loc=np.nanmean(a),
        scale=st.sem(a, nan_policy='omit'))


def compute_vulnerability(clean_images, dirty_images, net, loss_criterion=None, y=None, mask=None):
    """
    Computes vulnerability using foolbox package of net

    Parameters
    ----------
    net : :class:`torch.nn.Module`
        The network whose vulnerability is computed.
        Default: -1.

    """

    device = next(net.parameters()).device

    results = {}

    t0 = time.time()

    l2_norm = norm((clean_images.to(device) - dirty_images.to(device)).view(clean_images.size(0), -1), dim=1)
        #.mean().item()
    linf_norm = norm((clean_images.to(device) - dirty_images.to(device)).view(clean_images.size(0), -1),
                                ord=float("inf"),dim=1)

    l1_norm = norm((clean_images.to(device) - dirty_images.to(device)).view(clean_images.size(0), -1),
                     ord=1, dim=1)

    l2_snr = 20. * torch.log10(norm(clean_images.to(device).view(clean_images.size(0), -1),dim=1) / (l2_norm + 1e-6))

    _, clean_grad1, clean_grad2, clean_gradinf,clean_all_grad1, clean_all_grad2, clean_all_gradinf  = \
        classify(net, clean_images)

    #_, dirty_grad = classify(net, dirty_images)
    _, dirty_grad1, dirty_grad2, dirty_gradinf, dirty_all_grad1, dirty_all_grad2, dirty_all_gradinf = \
        classify(net, dirty_images)

    h = net(clean_images)
    h_adv = net(dirty_images)

    vulns = 0
    nb_vulns = 0

    for (task,grad) in dirty_grad1.items():

        if task == 'rep':
            pass

        results['norm1 clean grad {}'.format(task)] = clean_grad1[task]
        results['norm1 adv grad {}'.format(task)] = dirty_grad1[task]

        results['norm2 clean grad {}'.format(task)] = clean_grad2[task]
        results['norm2 adv grad {}'.format(task)] = dirty_grad2[task]

        results['normInf clean grad {}'.format(task)] = clean_gradinf[task]
        results['normInf adv grad {}'.format(task)] = dirty_gradinf[task]


        if loss_criterion is not None:
            loss_clean = loss_criterion[task](h[task], y[task], mask[task])
            loss_adv = loss_criterion[task](h_adv[task], y[task], mask[task])
            vuln =  loss_adv - loss_clean
            vulns = vulns+ vuln
            nb_vulns = nb_vulns+1
            results['vulnerability {}'.format(task)] = np.abs(vuln.item())

    results['vulnerability'] = np.abs(vulns.item())
    results['vulnerability weighted'] = np.abs(vulns.item()/nb_vulns)

    results['norm1 clean grad'] = clean_all_grad1
    results['norm2 clean grad'] = clean_all_grad2
    results['normInf clean grad'] = clean_all_gradinf

    results['norm1 dirty grad'] = dirty_all_grad1
    results['norm2 dirty grad'] = dirty_all_grad2
    results['normInf dirty grad'] = dirty_all_gradinf

    results['l2 snr'] = l2_snr.mean().item()
    results['l2 norm'] = l2_norm.mean().item()
    results['linf norm'] = linf_norm.mean().item()
    results['l1 norm'] = l1_norm.mean().item()

    t1 = time.time()
    # Printing summary
    """
    summary = {}
    for key, value in results.items():
        low95, high95 = conf95(value)
        print("{:>10} mean:{:>10.5f} std:{:>10.5f} conf95:({:>10.5f}, "
              "{:>10.5f}) minmax:({:>10.5f}, {:>10.5f})".format(
                  key, np.nanmean(value), np.nanstd(value), low95, high95,
                  np.nanmin(value), np.nanmax(value)))
        summary[key] = [np.nanmean(value), np.nanstd(value), low95, high95]

    print("{:>10} {:10d}s".format("Time", int(time.time() - t0)))
    """

    return results


def jacobian(y, x, create_graph=False):
    jac = []
    flat_y = y.reshape(-1)
    grad_y = torch.zeros_like(flat_y)
    for i in range(len(flat_y)):
        grad_y[i] = 1.
        grad_x, = torch.autograd.grad(flat_y, x, grad_y, retain_graph=True, create_graph=create_graph, allow_unused=True)
        jac.append(grad_x.reshape(x.shape))
        grad_y[i] = 0.
    return torch.stack(jac).reshape(y.shape + x.shape)


def hessian(y, x):
    return jacobian(jacobian(y, x, create_graph=True), x)



def get_second_order_grad(xs,net):
    xs = xs.detach()
    xs.requires_grad = True
    ys = net(xs)

    def task(task):
        def run(x):
            y = net(x)
            print("forward", task)
            return y[task].sum()

        return run

    print("computing hessian matrix")
    h = [(a, norm(hess_fn(task(a),xs)[0].view(xs.size(0), -1)).item()) for a in ys.keys() if a!="rep"]

    return h

    """
    for t,y in ys.items():
        if t!="null":
            break;
        for i,x in enumerate(xs):
            hess = hessian(y[i],x)
            print(hess.shape)

    device = next(net.parameters()).device

#g = [(a, norm(grad(y[a].sum(), x, retain_graph=True)[0].view(x.size(0), -1)).item()) for a in y.keys()]

    grad_tasks = [grad(ys[a].sum(),xs,retain_graph=True,create_graph=True)[0]  for a in ys.keys() if a!="rep"]
    start = time.time()
    second_order_grads = []
    for grads in grad_tasks:
        grads2 = []
        for j, (grad_, x) in enumerate(zip(grads, xs)):
            print('2nd order on ', j, 'th layer')
            print(x.size())
            grad_ = torch.reshape(grad_, [-1])
            grads2_tmp = []
            for count, g in enumerate(grad_):
                g2 = torch.autograd.grad(g, x, retain_graph=True,create_graph=True,allow_unused=True)[0]
                if g2 is not None:
                    g2 = torch.reshape(g2, [-1])
                    grads2_tmp.append(g2[count].data.cpu().numpy())
                else:
                    grads2_tmp.append(0)
            grads2.append(torch.from_numpy(np.reshape(grads2_tmp, x.size())).to(device))
            print('Time used is ', time.time() - start)
            
    """

# Copyright (c) Facebook, Inc. and its affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of https://github.com/facebookresearch/AdversarialAndDimensionality
#


import math
import time
import numpy as np
import scipy.stats as st

import torch
from torch.autograd import grad
from torch.linalg import norm

from torch.autograd.functional import hessian as hess_fn

def classify(net, x, x_adv=None, means=None, stds=None, top_k=None):
    device = next(net.parameters()).device

    if means is not None and stds is not None:
        x = x.to(device)
        means = torch.FloatTensor(means).to(device)
        stds = torch.FloatTensor(stds).to(device)
        x = ((x - means) / stds).detach()
    else:
        x = x.detach()
        if x_adv is not None:
            x_adv = x_adv.detach()
    x.requires_grad = True
    y = net(x)
    if x_adv is not None:
        x_adv.requires_grad = True
        g = [(a, norm(grad(y[a].sum(), x, retain_graph=True)[0].view(x.size(0), -1)-
                      grad(y[a].sum(), x_adv, retain_graph=True)[0].view(x_adv.size(0), -1)).item()) for a in y.keys() if a != 'rep']
    else:
        g1 = [(a, norm(grad(y[a].sum(), x, retain_graph=True)[0].view(x.size(0), -1), ord=1).item()) for a in y.keys()  if a != 'rep']
        g2 = [(a, norm(grad(y[a].sum(), x, retain_graph=True)[0].view(x.size(0), -1)).item()) for a in y.keys()  if a != 'rep']
        ginf = [(a, norm(grad(y[a].sum(), x, retain_graph=True)[0].view(x.size(0), -1), ord=float("inf")).item()) for a in y.keys()  if a != 'rep']

        g = [grad(y[a].sum(), x, retain_graph=True)[0].view(x.size(0), -1) for a in y.keys() if a != 'rep']
        g1_all = norm(torch.stack(g, dim=0).sum(dim=0), ord=1).item()
        g2_all = norm(torch.stack(g, dim=0).sum(dim=0)).item()
        ginf_all = norm(torch.stack(g, dim=0).sum(dim=0), ord=float("inf")).item()

        #g = [(a, grad(y[a].sum(), x, retain_graph=True)[0].view(x.size(0), -1).norm().item()) for a in y.keys()]
    if top_k is None:
        return y , dict(g1),dict(g2),dict(ginf), g1_all, g2_all, ginf_all
    else:
        _, top_indices = y.data.cpu().view(-1).topk(2)
        return top_indices[0].item(), dict(g1),dict(g2),dict(ginf), g1_all, g2_all, ginf_all


def conf95(a):
    return st.t.interval(
        0.95, len(a) - 1, loc=np.nanmean(a),
        scale=st.sem(a, nan_policy='omit'))


def compute_vulnerability(clean_images, dirty_images, net, loss_criterion=None, y=None, mask=None):
    """
    Computes vulnerability using foolbox package of net

    Parameters
    ----------
    net : :class:`torch.nn.Module`
        The network whose vulnerability is computed.
        Default: -1.

    """

    device = next(net.parameters()).device

    results = {}

    t0 = time.time()

    l2_norm = norm((clean_images.to(device) - dirty_images.to(device)).view(clean_images.size(0), -1), dim=1)
        #.mean().item()
    linf_norm = norm((clean_images.to(device) - dirty_images.to(device)).view(clean_images.size(0), -1),
                                ord=float("inf"),dim=1)

    l1_norm = norm((clean_images.to(device) - dirty_images.to(device)).view(clean_images.size(0), -1),
                     ord=1, dim=1)

    l2_snr = 20. * torch.log10(norm(clean_images.to(device).view(clean_images.size(0), -1),dim=1) / (l2_norm + 1e-6))

    _, clean_grad1, clean_grad2, clean_gradinf,clean_all_grad1, clean_all_grad2, clean_all_gradinf  = \
        classify(net, clean_images)

    #_, dirty_grad = classify(net, dirty_images)
    _, dirty_grad1, dirty_grad2, dirty_gradinf, dirty_all_grad1, dirty_all_grad2, dirty_all_gradinf = \
        classify(net, dirty_images)

    h = net(clean_images)
    h_adv = net(dirty_images)

    vulns = 0
    nb_vulns = 0

    for (task,grad) in dirty_grad1.items():

        if task == 'rep':
            pass

        results['norm1 clean grad {}'.format(task)] = clean_grad1[task]
        results['norm1 adv grad {}'.format(task)] = dirty_grad1[task]

        results['norm2 clean grad {}'.format(task)] = clean_grad2[task]
        results['norm2 adv grad {}'.format(task)] = dirty_grad2[task]

        results['normInf clean grad {}'.format(task)] = clean_gradinf[task]
        results['normInf adv grad {}'.format(task)] = dirty_gradinf[task]


        if loss_criterion is not None:
            loss_clean = loss_criterion[task](h[task], y[task], mask[task])
            loss_adv = loss_criterion[task](h_adv[task], y[task], mask[task])
            vuln =  loss_adv - loss_clean
            vulns = vulns+ vuln
            nb_vulns = nb_vulns+1
            results['vulnerability {}'.format(task)] = np.abs(vuln.item())

    results['vulnerability'] = np.abs(vulns.item())
    results['vulnerability weighted'] = np.abs(vulns.item()/nb_vulns)

    results['norm1 clean grad'] = clean_all_grad1
    results['norm2 clean grad'] = clean_all_grad2
    results['normInf clean grad'] = clean_all_gradinf

    results['norm1 dirty grad'] = dirty_all_grad1
    results['norm2 dirty grad'] = dirty_all_grad2
    results['normInf dirty grad'] = dirty_all_gradinf

    results['l2 snr'] = l2_snr.mean().item()
    results['l2 norm'] = l2_norm.mean().item()
    results['linf norm'] = linf_norm.mean().item()
    results['l1 norm'] = l1_norm.mean().item()

    t1 = time.time()
    # Printing summary
    """
    summary = {}
    for key, value in results.items():
        low95, high95 = conf95(value)
        print("{:>10} mean:{:>10.5f} std:{:>10.5f} conf95:({:>10.5f}, "
              "{:>10.5f}) minmax:({:>10.5f}, {:>10.5f})".format(
                  key, np.nanmean(value), np.nanstd(value), low95, high95,
                  np.nanmin(value), np.nanmax(value)))
        summary[key] = [np.nanmean(value), np.nanstd(value), low95, high95]

    print("{:>10} {:10d}s".format("Time", int(time.time() - t0)))
    """

    return results


def jacobian(y, x, create_graph=False):
    jac = []
    flat_y = y.reshape(-1)
    grad_y = torch.zeros_like(flat_y)
    for i in range(len(flat_y)):
        grad_y[i] = 1.
        grad_x, = torch.autograd.grad(flat_y, x, grad_y, retain_graph=True, create_graph=create_graph, allow_unused=True)
        jac.append(grad_x.reshape(x.shape))
        grad_y[i] = 0.
    return torch.stack(jac).reshape(y.shape + x.shape)


def hessian(y, x):
    return jacobian(jacobian(y, x, create_graph=True), x)



def get_second_order_grad(xs,net):
    xs = xs.detach()
    xs.requires_grad = True
    ys = net(xs)

    def task(task):
        def run(x):
            y = net(x)
            print("forward", task)
            return y[task].sum()

        return run

    print("computing hessian matrix")
    h = [(a, norm(hess_fn(task(a),xs)[0].view(xs.size(0), -1)).item()) for a in ys.keys() if a!="rep"]

    return h

    """
    for t,y in ys.items():
        if t!="null":
            break;
        for i,x in enumerate(xs):
            hess = hessian(y[i],x)
            print(hess.shape)

    device = next(net.parameters()).device

#g = [(a, norm(grad(y[a].sum(), x, retain_graph=True)[0].view(x.size(0), -1)).item()) for a in y.keys()]

    grad_tasks = [grad(ys[a].sum(),xs,retain_graph=True,create_graph=True)[0]  for a in ys.keys() if a!="rep"]
    start = time.time()
    second_order_grads = []
    for grads in grad_tasks:
        grads2 = []
        for j, (grad_, x) in enumerate(zip(grads, xs)):
            print('2nd order on ', j, 'th layer')
            print(x.size())
            grad_ = torch.reshape(grad_, [-1])
            grads2_tmp = []
            for count, g in enumerate(grad_):
                g2 = torch.autograd.grad(g, x, retain_graph=True,create_graph=True,allow_unused=True)[0]
                if g2 is not None:
                    g2 = torch.reshape(g2, [-1])
                    grads2_tmp.append(g2[count].data.cpu().numpy())
                else:
                    grads2_tmp.append(0)
            grads2.append(torch.from_numpy(np.reshape(grads2_tmp, x.size())).to(device))
            print('Time used is ', time.time() - start)
            
    """

# Copyright (c) Facebook, Inc. and its affiliates. # All rights reserved. # # This source code is licensed under the license found in the # LICENSE file in the root directory of https://github.com/facebookresearch/AdversarialAndDimensionality # #g = [(a, grad(y[a].sum(), x, retain_graph=True)[0].view(x.size(0), -1).norm().item()) for a in y.keys()] Computes vulnerability using foolbox package of net      Parameters     ----------     net : :class:`torch.nn.Module`         The network whose vulnerability is computed.         Default: -1. #.mean().item() #_, dirty_grad = classify(net, dirty_images) # Printing summary summary = {}     for key, value in results.items():         low95, high95 = conf95(value)         print("{:>10} mean:{:>10.5f} std:{:>10.5f} conf95:({:>10.5f}, "               "{:>10.5f}) minmax:({:>10.5f}, {:>10.5f})".format(                   key, np.nanmean(value), np.nanstd(value), low95, high95,                   np.nanmin(value), np.nanmax(value)))         summary[key] = [np.nanmean(value), np.nanstd(value), low95, high95]      print("{:>10} {:10d}s".format("Time", int(time.time() - t0))) for t,y in ys.items():         if t!="null":             break;         for i,x in enumerate(xs):             hess = hessian(y[i],x)             print(hess.shape)      device = next(net.parameters()).device  #g = [(a, norm(grad(y[a].sum(), x, retain_graph=True)[0].view(x.size(0), -1)).item()) for a in y.keys()]      grad_tasks = [grad(ys[a].sum(),xs,retain_graph=True,create_graph=True)[0]  for a in ys.keys() if a!="rep"]     start = time.time()     second_order_grads = []     for grads in grad_tasks:         grads2 = []         for j, (grad_, x) in enumerate(zip(grads, xs)):             print('2nd order on ', j, 'th layer')             print(x.size())             grad_ = torch.reshape(grad_, [-1])             grads2_tmp = []             for count, g in enumerate(grad_):                 g2 = torch.autograd.grad(g, x, retain_graph=True,create_graph=True,allow_unused=True)[0]                 if g2 is not None:                     g2 = torch.reshape(g2, [-1])                     grads2_tmp.append(g2[count].data.cpu().numpy())                 else:                     grads2_tmp.append(0)             grads2.append(torch.from_numpy(np.reshape(grads2_tmp, x.size())).to(device))             print('Time used is ', time.time() - start)