#!/usr/bin/env python3

import sys
sys.path.insert(1, '../lib')
import argparse
import lmdk_lib
import lmdk_sel
import exp_mech
import numpy as np
import os
from matplotlib import pyplot as plt
import time


def main(args):
  # Privacy goal
  epsilon = [.01, .1, 1.0, 10.0, 100.0]
  # Number of timestamps
  seq = lmdk_lib.get_seq(1, args.time)
  # Distribution type
  dist_type = np.array(range(-1, 4))
  # Number of landmarks
  lmdk_n = np.array(range(int(.2*args.time), args.time, int(args.time/5)))
  # Width of bars
  bar_width = 1/(len(epsilon) + 1)
  # The x axis
  x_i = np.arange(len(lmdk_n))
  x_margin = bar_width*(len(epsilon)/2 + 1)
  for d_i, d in enumerate(dist_type):
    # Logging
    title =  lmdk_lib.dist_type_to_str(d) + ' landmark distribution'
    print('(%d/%d) %s... ' %(d_i + 1, len(dist_type), title), end='', flush=True)
    # Initialize plot
    lmdk_lib.plot_init()
    # The x axis
    plt.xticks(x_i, ((lmdk_n/len(seq))*100).astype(int))
    plt.xlabel('Landmarks (%)')  # Set x axis label.
    plt.xlim(x_i.min() - x_margin, x_i.max() + x_margin)
    # The y axis
    plt.ylabel('Mean absolute error')  # Set y axis label.
    # plt.ylim(0, len(seq)/3)
    # Bar offset
    x_offset = -(bar_width/2)*(len(epsilon) - 1)
    for e_i, e in enumerate(epsilon):
      mae = np.zeros(len(lmdk_n))
      for n_i, n in enumerate(lmdk_n):
        for r in range(args.reps):
          lmdks = lmdk_lib.get_lmdks(seq, n, d)
          hist, h = lmdk_lib.get_hist(seq, lmdks)
          res = np.zeros([len(hist)])
          # Split sequence in parts of size h 
          pt_idx = []
          for idx in range(h, len(seq), h):
            pt_idx.append(idx)
          seq_pt = np.split(seq, pt_idx)
          for pt_i, pt in enumerate(seq_pt):
            # Find this part's landmarks
            lmdks_pt = np.intersect1d(pt, lmdks)
            # Find possible options for this part
            opts = lmdk_sel.get_opts_from_top_h(pt, lmdks_pt)
            # Turn part to histogram
            hist_pt, _ = lmdk_lib.get_hist(pt, lmdks_pt)
            # Get an option for this part
            if len(opts) > 1:
              res_pt, _ = exp_mech.exponential(hist_pt, opts, exp_mech.score, 1.0, e)
            elif len(opts) > 0:
              res_pt = opts[0]
            # Merge options of all parts
            res[pt_i] = np.sum(res_pt)
          # Calculate MAE
          mae[n_i] += lmdk_lib.get_norm(hist, res)/args.reps
      # Plot bar for current epsilon
      plt.bar(
        x_i + x_offset,
        mae,
        bar_width,
        label=u'\u03B5 = ' + str("{:.0e}".format(e)),
        linewidth=lmdk_lib.line_width
      )
      # Change offset for next bar
      x_offset += bar_width
    path = str('../../rslt/lmdk_sel/' + title)
    # Plot legend
    lmdk_lib.plot_legend()
    # Show plot
    # plt.show()
    # Save plot
    lmdk_lib.save_plot(path + '.pdf')
    print('[OK]', flush=True)


'''
  Parse arguments.

  Optional:
    reps - The number of repetitions.
    time - The time limit of the sequence.
'''
def parse_args():
  # Create argument parser.
  parser = argparse.ArgumentParser()

  # Mandatory arguments.

  # Optional arguments.
  parser.add_argument('-r', '--reps', help='The number of repetitions.', type=int, default=1)
  parser.add_argument('-t', '--time', help='The time limit of the sequence.', type=int, default=100)

  # 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()