#!/usr/bin/env python3

import sys
sys.path.insert(1, '../lib')
import argparse
import ast
from datetime import datetime
from geopy.distance import distance
import lmdk_bgt
import lmdk_lib
import lmdk_sel
import exp_mech
import math
import numpy as np
from matplotlib import pyplot as plt
import time


def main(args):
  # User's consumption
  seq = lmdk_lib.load_data(args, 'cons')
  # The name of the dataset
  d = 'HUE'
  # The landmarks percentages
  lmdks_pct = [0, 20, 40, 60, 80, 100]
  # Landmarks' thresholds
  lmdks_th = [0, .54, .68, .88, 1.12, 10]
  # The privacy budget
  epsilon = 1.0
  eps_pct = [.01, .1, .25, .5]

  markers = [
    '^',
    'v',
    'D',
    's'
  ]

  print('\n##############################', d, '\n')

  # Initialize plot
  lmdk_lib.plot_init()
  # The x axis
  x_i = np.arange(len(lmdks_pct))
  plt.xticks(x_i, np.array(lmdks_pct, int))
  plt.xlabel('Landmarks (%)')  # Set x axis label.
  plt.xlim(x_i.min(), x_i.max())
  # The y axis
  plt.ylabel('Mean absolute error (kWh)')  # Set y axis label.
  # plt.yscale('log')
  # plt.ylim(.1, 100000)

  mae_evt = 0
  mae_usr = 0

  for i_e, e in enumerate(eps_pct):
    mae = np.zeros(len(lmdks_pct))

    for i, pct in enumerate(lmdks_pct):
      # Find landmarks
      lmdks = seq[seq[:, 1] < lmdks_th[i]]

      for _ in range(args.iter):
        lmdks_sel = lmdk_sel.find_lmdks_eps(seq, lmdks, epsilon*e)

        # Uniform
        rls_data, _ = lmdk_bgt.uniform_cons(seq, lmdks_sel, epsilon*(1 - e))
        mae[i] += lmdk_bgt.mae_cons(seq, rls_data)/args.iter

        # Calculate once
        if e == eps_pct[0] and pct == lmdks_pct[0]:
          # Event
          rls_data_evt, _ = lmdk_bgt.uniform_cons(seq, lmdks, epsilon)
          mae_evt += lmdk_bgt.mae_cons(seq, rls_data_evt)/args.iter
        elif e == eps_pct[-1] and pct == lmdks_pct[-1]:
          # User
          rls_data_usr, _ = lmdk_bgt.uniform_cons(seq, lmdks, epsilon)
          mae_usr += lmdk_bgt.mae_cons(seq, rls_data_usr)/args.iter

    # Plot line
    plt.plot(
      x_i,
      mae,
      label='{0:.2f}'.format(e) + 'ε',
      marker=markers[i_e],
      markersize=lmdk_lib.marker_size,
      markeredgewidth=0,
      linewidth=lmdk_lib.line_width
    )

  plt.axhline(
    y = mae_evt,
    color = '#212121',
    linewidth=lmdk_lib.line_width
  )
  plt.text(x_i[-1] + x_i[-1]*.01, mae_evt - mae_evt*.04, 'event')

  plt.axhline(
    y = mae_usr,
    color = '#616161',
    linewidth=lmdk_lib.line_width
  )
  plt.text(x_i[-1] + x_i[-1]*.01, mae_usr - mae_usr*.04, 'user')

  path = str('../../rslt/lmdk_sel_eps/' + d)
  # Plot legend
  lmdk_lib.plot_legend()
  # Show plot
  # plt.show()
  # Save plot
  lmdk_lib.save_plot(path + '-sel-eps.pdf')
  print('[OK]', flush=True)


def parse_args():
  '''
    Parse arguments.

    Optional:
      res  - The results archive file.
      iter - The total iterations.
  '''
  # Create argument parser.
  parser = argparse.ArgumentParser()

  # Mandatory arguments.

  # Optional arguments.
  parser.add_argument('-r', '--res', help='The results archive file.', type=str, default='/home/manos/Cloud/Data/HUE/Results.zip')
  parser.add_argument('-i', '--iter', help='The total iterations.', type=int, default=1)

  # Parse arguments.
  args = parser.parse_args()

  return args


if __name__ == '__main__':
  try:
    start_time = time.time()
    main(parse_args())
    end_time = time.time()
    print('##############################')
    print('Time elapsed: %s' % (time.strftime('%H:%M:%S', time.gmtime(end_time - start_time))))
    print('##############################')
  except KeyboardInterrupt:
    print('Interrupted by user.')
    exit()
