# Copyright 2022 The Authors
# 
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# 
#     https://www.apache.org/licenses/LICENSE-2.0
# 
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import tensorflow as tf
import numpy as np
import torch

class ModelAdapter:
  def __init__(self, logits, x, y, sess, num_classes=10, device="cuda"):
    self.logits = logits
    self.sess = sess
    self.x_input = x
    self.y_input = y
    self.num_classes = num_classes
    self.device = device

    # gradients of logits
    if num_classes <= 10:
      self.grads = [None] * num_classes
      for cl in range(num_classes):
        self.grads[cl] = tf.gradients(self.logits[:, cl], self.x_input)[0]

    # cross-entropy loss
    self.xent = tf.nn.sparse_softmax_cross_entropy_with_logits(
        logits=self.logits, labels=self.y_input)
    self.grad_xent = tf.gradients(self.xent, self.x_input)[0]

    # dlr loss
    self.dlr = dlr_loss(self.logits, self.y_input, num_classes=self.num_classes)
    self.grad_dlr = tf.gradients(self.dlr, self.x_input)[0]

    # targeted dlr loss
    self.y_target = tf.placeholder(tf.int64, shape=[None])
    self.dlr_target = dlr_loss_targeted(self.logits, self.y_input, self.y_target, num_classes=self.num_classes)
    self.grad_target = tf.gradients(self.dlr_target, self.x_input)[0]

    self.la = tf.placeholder(tf.int64, shape=[None])
    self.la_target = tf.placeholder(tf.int64, shape=[None])
    la_mask = tf.one_hot(self.la, self.num_classes)
    la_target_mask = tf.one_hot(self.la_target, self.num_classes)
    la_logit = tf.reduce_sum(la_mask * self.logits, axis=1)
    la_target_logit = tf.reduce_sum(la_target_mask * self.logits, axis=1)
    self.diff_logits = la_target_logit - la_logit
    self.grad_diff_logits = tf.gradients(self.diff_logits, self.x_input)[0]

  def predict(self, x):
    x = x.detach()
    x2 = np.moveaxis(x.cpu().numpy(), 1, 3)
    y = self.sess.run(self.logits, {self.x_input: x2})

    return torch.from_numpy(y).to(self.device)

  def grad_logits(self, x):
    x = x.detach()
    x2 = np.moveaxis(x.cpu().numpy(), 1, 3)
    logits, g2 = self.sess.run([self.logits, self.grads], {self.x_input: x2})
    g2 = np.moveaxis(np.array(g2), 0, 1)
    g2 = np.transpose(g2, (0, 1, 4, 2, 3))

    return torch.from_numpy(logits).cuda(), torch.from_numpy(g2).cuda()

  def get_grad_diff_logits_target(self, x, y=None, y_target=None):
    x = x.detach()
    la = y.cpu().numpy()
    la_target = y_target.cpu().numpy()
    x2 = np.moveaxis(x.cpu().numpy(), 1, 3)
    dl, g2 = self.sess.run([self.diff_logits, self.grad_diff_logits], {self.x_input: x2, self.la: la, self.la_target: la_target})
    g2 = np.transpose(np.array(g2), (0, 3, 1, 2))

    return torch.from_numpy(dl).to(self.device), torch.from_numpy(g2).to(self.device)

  def get_logits_loss_grad_xent(self, x, y):
    x = x.detach()
    x2 = np.moveaxis(x.cpu().numpy(), 1, 3)
    y2 = y.clone().cpu().numpy()
    logits_val, loss_indiv_val, grad_val = self.sess.run([self.logits, self.xent, self.grad_xent], {self.x_input: x2, self.y_input: y2})
    grad_val = np.moveaxis(grad_val, 3, 1)

    return torch.from_numpy(logits_val).to(self.device), torch.from_numpy(loss_indiv_val).to(self.device), torch.from_numpy(grad_val).to(self.device)

  def get_logits_loss_grad_dlr(self, x, y):
    x = x.detach()
    x2 = np.moveaxis(x.cpu().numpy(), 1, 3)
    y2 = y.clone().cpu().numpy()
    logits_val, loss_indiv_val, grad_val = self.sess.run([self.logits, self.dlr, self.grad_dlr], {self.x_input: x2, self.y_input: y2})
    grad_val = np.moveaxis(grad_val, 3, 1)

    return torch.from_numpy(logits_val).to(self.device), torch.from_numpy(loss_indiv_val).to(self.device), torch.from_numpy(grad_val).to(self.device)

  def get_logits_loss_grad_target(self, x, y, y_target):
    x = x.detach()
    x2 = np.moveaxis(x.cpu().numpy(), 1, 3)
    y2 = y.clone().cpu().numpy()
    y_targ = y_target.clone().cpu().numpy()
    logits_val, loss_indiv_val, grad_val = self.sess.run([self.logits, self.dlr_target, self.grad_target], {self.x_input: x2, self.y_input: y2, self.y_target: y_targ})
    grad_val = np.moveaxis(grad_val, 3, 1)

    return torch.from_numpy(logits_val).to(self.device), torch.from_numpy(loss_indiv_val).to(self.device), torch.from_numpy(grad_val).to(self.device)


def dlr_loss(x, y, num_classes=10):
  x_sort = tf.contrib.framework.sort(x, axis=1)
  y_onehot = tf.one_hot(y, num_classes)
  if num_classes > 2:
    ### TODO: adapt to the case when the point is already misclassified
    loss = -(x_sort[:, -1] - x_sort[:, -2]) / (x_sort[:, -1] - x_sort[:, -3] + 1e-12)
  else:
    loss = (tf.reduce_max(x - y_onehot * 1e9) - tf.gather(x, y, axis=-1)) / tf.reduce_max(x - y_onehot * 1e9)

  return loss

def dlr_loss_targeted(x, y, y_target, num_classes=10):
  x_sort = tf.contrib.framework.sort(x, axis=1)
  y_onehot = tf.one_hot(y, num_classes)
  y_target_onehot = tf.one_hot(y_target, num_classes)
  loss = -(tf.reduce_sum(x * y_onehot, axis=1) - tf.reduce_sum(x * y_target_onehot, axis=1)) / (x_sort[:, -1] - .5 * x_sort[:, -3] - .5 * x_sort[:, -4] + 1e-12)

  return loss