"""Tool which runs all attacks against all defenses and computes results."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import argparse
import csv
import json
import os
import subprocess
import numpy as np
from PIL import Image
import pickle

def parse_args():
  """Parses command line arguments."""
  parser = argparse.ArgumentParser(
      description='Tool to run attacks and defenses.')
  parser.add_argument('--attacks_dir', required=True,
                      help='Location of all attacks.')
  parser.add_argument('--targeted_attacks_dir', required=True,
                      help='Location of all targeted attacks.')
  parser.add_argument('--defenses_dir', required=True,
                      help='Location of all defenses.')
  parser.add_argument('--dataset_dir', required=True,
                      help='Location of the dataset.')
  parser.add_argument('--dataset_metadata', required=True,
                      help='Location of the dataset metadata.')
  parser.add_argument('--intermediate_results_dir', required=True,
                      help='Directory to store intermediate results.')
  parser.add_argument('--output_dir', required=True,
                      help=('Output directory.'))
  parser.add_argument('--batch_size', required=False,type=int,default=10,
                      help=('batch size.'))


  parser.add_argument('--epsilon', required=False, type=int, default=16,
                      help='Maximum allowed size of adversarial perturbation')
  parser.add_argument('--gpu', dest='use_gpu', action='store_true')
  parser.add_argument('--nogpu', dest='use_gpu', action='store_false')
  parser.set_defaults(use_gpu=False)
  parser.add_argument('--gpuid', default=0, type=str)
  parser.add_argument('--use_docker',action='store_true')
  parser.add_argument('--use_existing', type=str, default='0')
  parser.add_argument('--models', default='all', type=str)
  parser.add_argument('--save_all_classification',
                      dest='save_all_classification', action='store_true')
  parser.add_argument('--nosave_all_classification',
                      dest='save_all_classification', action='store_false')
  parser.set_defaults(save_all_classification=False)
  return parser.parse_args()


class Submission(object):
  """Base class for all submissions."""

  def __init__(self, directory, container, entry_point, use_gpu):
    """Initializes instance of Submission class.

    Args:
      directory: location of the submission.
      container: URL of Docker container which should be used to run submission.
      entry_point: entry point script, which invokes submission.
      use_gpu: whether to use Docker with GPU or not.
    """
    self.name = os.path.basename(directory)
    self.directory = directory
    self.container = container
    self.entry_point = entry_point
    self.use_gpu = use_gpu

  def docker_binary(self):
    """Returns appropriate Docker binary to use."""
    return 'nvidia-docker' if self.use_gpu else 'docker'


class Attack(Submission):
  """Class which stores and runs attack."""

  def __init__(self, directory, container, entry_point, use_gpu):
    """Initializes instance of Attack class."""
    super(Attack, self).__init__(directory, container, entry_point, use_gpu)

  def run(self, input_dir, output_dir, epsilon):
    """Runs attack inside Docker.

    Args:
      input_dir: directory with input (dataset).
      output_dir: directory where output (adversarial images) should be written.
      epsilon: maximum allowed size of adversarial perturbation,
        should be in range [0, 255].
    """
    print('Running attack ', self.name)
    if not args.use_docker:
        cwd = os.getcwd()
        os.chdir(self.directory)
        print(os.listdir('./'))
        cmd = ['CUDA_VISIBLE_DEVICES='+str(args.gpuid),
            'bash',
            self.entry_point,
            input_dir,
            output_dir,
            str(epsilon),
            str(args.gpuid),
            str(args.use_existing),
            str(args.batch_size)]

    else:
        cmd = ['NV_GPU='+str(args.gpuid),
           self.docker_binary(), 'run',
           '-v', '{0}:/input_images'.format(input_dir),
           '-v', '{0}:/output_images'.format(output_dir),
           '-v', '{0}:/code'.format(self.directory),
           '-w', '/code',
           self.container,
           './' + self.entry_point,
           '/input_images',
           '/output_images',
           str(epsilon),
           str(args.gpuid),
           str(args.use_existing),
           str(args.batch_size)]
    print(' '.join(cmd))
    subprocess.call(' '.join(cmd),shell = True)    
    os.chdir(cwd)

class Defense(Submission):
  """Class which stores and runs defense."""

  def __init__(self, directory, container, entry_point, use_gpu):
    """Initializes instance of Defense class."""
    super(Defense, self).__init__(directory, container, entry_point, use_gpu)

  def run(self, input_dir, output_dir):
    """Runs defense inside Docker.

    Args:
      input_dir: directory with input (adversarial images).
      output_dir: directory to write output (classification result).
    """
    print('Running defense ', self.name)
    cmd = [self.docker_binary(), 'run',
           '-v', '{0}:/input_images'.format(input_dir),
           '-v', '{0}:/output_data'.format(output_dir),
           '-v', '{0}:/code'.format(self.directory),
           '-w', '/code',
           self.container,
           './' + self.entry_point,
           '/input_images',
           '/output_data/result.csv']
    print(' '.join(cmd))
    subprocess.call(cmd)


def read_submissions_from_directory(dirname, use_gpu, models='all'):
  """Scans directory and read all submissions.

  Args:
    dirname: directory to scan.
    use_gpu: whether submissions should use GPU. This argument is
      used to pick proper Docker container for each submission and create
      instance of Attack or Defense class.

  Returns:
    List with submissions (subclasses of Submission class).
  """
  result = []
  if models == 'all':
      sub_dirs = os.listdir(dirname)
  else:
      sub_dirs = models.split(',')

  for sub_dir in sub_dirs:
    submission_path = os.path.join(dirname, sub_dir)
    try:
      if not os.path.isdir(submission_path):
        continue
      if not os.path.exists(os.path.join(submission_path, 'metadata.json')):
        continue
      with open(os.path.join(submission_path, 'metadata.json')) as f:
        metadata = json.load(f)
      if use_gpu and ('container_gpu' in metadata):
        container = metadata['container_gpu']
      else:
        container = metadata['container']
      entry_point = metadata['entry_point']
      submission_type = metadata['type']
      if submission_type == 'attack' or submission_type == 'targeted_attack':
        submission = Attack(submission_path, container, entry_point, use_gpu)
      elif submission_type == 'defense':
        submission = Defense(submission_path, container, entry_point, use_gpu)
      else:
        raise ValueError('Invalid type of submission: %s', submission_type)
      result.append(submission)
    except (IOError, KeyError, ValueError):
      print('Failed to read submission from directory ', submission_path)
  return result


class AttacksOutput(object):
  """Helper class to store data about images generated by attacks."""

  def __init__(self,
               dataset_dir,
               attacks_output_dir,
               targeted_attacks_output_dir,
               all_adv_examples_dir,
               epsilon):
    """Initializes instance of AttacksOutput class.

    Args:
      dataset_dir: location of the dataset.
      attacks_output_dir: where to write results of attacks.
      targeted_attacks_output_dir: where to write results of targeted attacks.
      all_adv_examples_dir: directory to copy all adversarial examples from
        all attacks.
      epsilon: maximum allowed size of adversarial perturbation.
    """
    self.attacks_output_dir = attacks_output_dir
    self.targeted_attacks_output_dir = targeted_attacks_output_dir
    self.all_adv_examples_dir = all_adv_examples_dir
    #self._load_dataset_clipping(dataset_dir, epsilon)
    self._output_image_idx = 0
    self._output_to_attack_mapping = {}
    self._attack_image_count = 0
    self._targeted_attack_image_count = 0
    self._attack_names = set()
    self._targeted_attack_names = set()

  def _load_dataset_clipping(self, dataset_dir, epsilon):
    """Helper method which loads dataset and determines clipping range.

    Args:
      dataset_dir: location of the dataset.
      epsilon: maximum allowed size of adversarial perturbation.
    """
    self.dataset_max_clip = {}
    self.dataset_min_clip = {}
    self._dataset_image_count = 0
    for fname in os.listdir(dataset_dir):
      if not fname.endswith('.png'):
        continue
      image_id = fname[:-4]
      image = np.array(
          Image.open(os.path.join(dataset_dir, fname)).convert('RGB'))
      image = image.astype('int32')
      self._dataset_image_count += 1
      self.dataset_max_clip[image_id] = np.clip(image + epsilon,
                                                0,
                                                255).astype('uint8')
      self.dataset_min_clip[image_id] = np.clip(image - epsilon,
                                                0,
                                                255).astype('uint8')

  def clip_and_copy_attack_outputs(self, attack_name, is_targeted):
    """Clips results of attack and copy it to directory with all images.

    Args:
      attack_name: name of the attack.
      is_targeted: if True then attack is targeted, otherwise non-targeted.
    """
    if is_targeted:
      self._targeted_attack_names.add(attack_name)
    else:
      self._attack_names.add(attack_name)
    attack_dir = os.path.join(self.targeted_attacks_output_dir
                              if is_targeted
                              else self.attacks_output_dir,
                              attack_name)
    for fname in os.listdir(attack_dir):
      if not (fname.endswith('.png') or fname.endswith('.jpg')):
        continue
      image_id = fname[:-4]
      if image_id not in self.dataset_max_clip:
        continue
      image_max_clip = self.dataset_max_clip[image_id]
      image_min_clip = self.dataset_min_clip[image_id]
      adversarial_image = np.array(
          Image.open(os.path.join(attack_dir, fname)).convert('RGB'))
      clipped_adv_image = np.clip(adversarial_image,
                                  image_min_clip,
                                  image_max_clip)
      output_basename = '{0:08d}'.format(self._output_image_idx)
      self._output_image_idx += 1
      self._output_to_attack_mapping[output_basename] = (attack_name,
                                                         is_targeted,
                                                         image_id)
      if is_targeted:
        self._targeted_attack_image_count += 1
      else:
        self._attack_image_count += 1
      Image.fromarray(clipped_adv_image).save(
          os.path.join(self.all_adv_examples_dir, output_basename + '.png'))

  @property
  def attack_names(self):
    """Returns list of all non-targeted attacks."""
    return self._attack_names

  @property
  def targeted_attack_names(self):
    """Returns list of all targeted attacks."""
    return self._targeted_attack_names

  @property
  def attack_image_count(self):
    """Returns number of all images generated by non-targeted attacks."""
    return self._attack_image_count

  @property
  def dataset_image_count(self):
    """Returns number of all images in the dataset."""
    return self._dataset_image_count

  @property
  def targeted_attack_image_count(self):
    """Returns number of all images generated by targeted attacks."""
    return self._targeted_attack_image_count

  def image_by_base_filename(self, filename):
    """Returns information about image based on it's filename."""
    return self._output_to_attack_mapping[filename]


class DatasetMetadata(object):
  """Helper class which loads and stores dataset metadata."""

  def __init__(self, filename):
    """Initializes instance of DatasetMetadata."""
    self._true_labels = {}
    self._target_classes = {}
    with open(filename) as f:
      reader = csv.reader(f)
      header_row = next(reader)
      try:
        row_idx_image_id = header_row.index('ImageId')
        row_idx_true_label = header_row.index('TrueLabel')
        row_idx_target_class = header_row.index('TargetClass')
      except ValueError:
        raise IOError('Invalid format of dataset metadata.')
      for row in reader:
        if len(row) < len(header_row):
          # skip partial or empty lines
          continue
        try:
          image_id = row[row_idx_image_id]
          self._true_labels[image_id] = int(row[row_idx_true_label])
          self._target_classes[image_id] = int(row[row_idx_target_class])
        except (IndexError, ValueError):
          raise IOError('Invalid format of dataset metadata')

  def get_true_label(self, image_id):
    """Returns true label for image with given ID."""
    return self._true_labels[image_id]

  def get_target_class(self, image_id):
    """Returns target class for image with given ID."""
    return self._target_classes[image_id]

  def save_target_classes(self, filename):
    """Saves target classed for all dataset images into given file."""
    with open(filename, 'w') as f:
      for k, v in self._target_classes.iteritems():
        f.write('{0}.png,{1}\n'.format(k, v))


def load_defense_output(filename):
  """Loads output of defense from given file."""
  result = {}
  with open(filename) as f:
    for row in csv.reader(f):
      try:
        image_filename = row[0]
        if image_filename.endswith('.png') or image_filename.endswith('.jpg'):
          image_filename = image_filename[:image_filename.rfind('.')]
        label = int(row[1])
      except (IndexError, ValueError):
        continue
      result[image_filename] = label
  return result



def main():
  global args
  args = parse_args()
  
  attacks_output_dir = os.path.join(args.intermediate_results_dir,
                                    'attacks_output')
  targeted_attacks_output_dir = os.path.join(args.intermediate_results_dir,
                                             'targeted_attacks_output')
  defenses_output_dir = os.path.join(args.intermediate_results_dir,
                                     'defenses_output')
  all_adv_examples_dir = os.path.join(args.intermediate_results_dir,
                                      'all_adv_examples')

  # Load dataset metadata.
  dataset_meta = DatasetMetadata(args.dataset_metadata)

  # Load attacks and defenses.
  attacks = [
      a for a in read_submissions_from_directory(args.attacks_dir,
                                                 args.use_gpu,args.models)
      if isinstance(a, Attack)
  ]
  print('Found attacks: ', [a.name for a in attacks])

  # Prepare subdirectories for intermediate results.
  try:
    os.mkdir(attacks_output_dir)
    os.mkdir(all_adv_examples_dir)                                                        
  except:
    None

  for a in attacks:
    try:
      os.mkdir(os.path.join(attacks_output_dir, a.name))
    except:
      None
  # Run all non-targeted attacks.
  attacks_output = AttacksOutput(args.dataset_dir,
                                 attacks_output_dir,
                                 targeted_attacks_output_dir,
                                 all_adv_examples_dir,
                                 args.epsilon)
  for a in attacks:
    a.run(args.dataset_dir,
          os.path.join(attacks_output_dir, a.name),
          args.epsilon)
    #attacks_output.clip_and_copy_attack_outputs(a.name, False)
    
    
  #attacks_output.dataset_max_clip = None
  #attacks_output.dataset_min_clip = None
 # saving_file = os.path.join(args.intermediate_results_dir,'attacks_output.save')
  #pickle.dump(attacks_output, open(saving_file,'w'))

if __name__ == '__main__':
  main()
