# coding: utf-8
###
 # @file   empire.py
 # @author Sébastien Rouault <sebastien.rouault@alumni.epfl.ch>
 #
 # @section LICENSE
 #
 # Copyright © 2019-2020 École Polytechnique Fédérale de Lausanne (EPFL).
 # All rights reserved.
 #
 # @section DESCRIPTION
 #
 # Strict implementation of the attack introduced in the following paper:
 #   Cong Xie, Oluwasanmi Koyejo, Indranil Gupta.
 #   Fall of Empires: Breaking Byzantine-tolerant SGD by Inner Product Manipulation.
 #   UAI 2019. URL: http://auai.org/uai2019/proceedings/papers/83.pdf
###

import tools

import math
import torch

# ---------------------------------------------------------------------------- #
# Attack implementation

def fall_empires_attack(grad_honests, f_real, f_decl, defense, model, epsilon=1, **kwargs):
    """ Generate the attack gradients.
    Args:
      grad_honests Non-empty list of honest gradients
      f_decl       Number of declared Byzantine gradients
      f_real       Number of Byzantine gradients to generate
      defense      Aggregation rule in use to defeat
      model        Model with valid default dataset and loss set
      epsilon      Fixed epsilon if positive, number of evaluations for best attack epsilon if negative
      ...          Ignored keyword-arguments
    Returns:
      Generated Byzantine gradients
    """
    # Fast path
    if f_real == 0:
        return list()
    # Compute the honest average gradient
    grad_avg = grad_honests[0].clone().detach()
    for grad in grad_honests[1:]:
        grad_avg.add_(grad)
    grad_avg.div_(len(grad_honests))
    # Evaluate the best attack epsilon (if required)
    if epsilon < 0:
        def eval_epsilon(epsilon):
          # Apply the given epsilon
          grad_attack = grad_avg.mul(-epsilon)
          # Measure effective squared distance
          aggregated = defense(gradients=(grad_honests + [grad_attack] * f_real), f=f_decl, model=model)
          aggregated.sub_(grad_avg)
          return aggregated.dot(aggregated).item()
        epsilon = tools.line_maximize(eval_epsilon, evals=math.ceil(-epsilon))
    # Generate the Byzantine gradient
    byz_grad = grad_avg
    byz_grad.mul_(-epsilon)
    # Return this Byzantine gradient 'f_real' times
    return [byz_grad] * f_real

