#!/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 = 1.0
  # Number of timestamps
  seq = lmdk_lib.get_seq(1, args.time)
  # Distribution type
  dist_type = np.array(range(0, 4))
  # Number of landmarks
  lmdk_n = np.array(range(0, args.time + 1, int(args.time/5)))

  markers = [
    '^', # Symmetric
    'v', # Skewed
    'D', # Bimodal
    's'  # Uniform
  ]

  # Initialize plot
  lmdk_lib.plot_init()
  # Width of bars
  bar_width = 1/(len(dist_type) + 1)
  # The x axis
  x_i = np.arange(len(lmdk_n))
  x_margin = bar_width*(len(dist_type)/2 + 1)
  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)
  plt.xlim(x_i.min(), x_i.max())
  # The y axis
  # plt.yscale('log')
  plt.ylim(0, 1)
  plt.ylabel('Normalized Euclidean distance')  # Set y axis label.
  # plt.ylabel('Normalized Wasserstein distance')  # Set y axis label.
  # Bar offset
  x_offset = -(bar_width/2)*(len(dist_type) - 1)
  for d_i, d in enumerate(dist_type):
    # Set label
    label = lmdk_lib.dist_type_to_str(d)
    if d_i == 1:
      label = 'Skewed'
    # Logging
    title =  label + ' landmark distribution'
    print('(%d/%d) %s... ' %(d_i + 1, len(dist_type), title), end='', flush=True)
    mae = np.zeros(len(lmdk_n))
    for n_i, n in enumerate(lmdk_n):
      if n == lmdk_n[-1]:
        break
      for r in range(args.iter):
        lmdks = lmdk_lib.get_lmdks(seq, n, d)
        hist, h = lmdk_lib.get_hist(seq, lmdks)
        opts = lmdk_sel.get_opts_from_top_h(seq, lmdks)
        delta = 1.0
        res, _ = exp_mech.exponential(hist, opts, exp_mech.score, delta, epsilon)
        mae[n_i] += lmdk_lib.get_norm(hist, res)/args.iter  # Euclidean
        # mae[n_i] += lmdk_lib.get_emd(hist, res)/args.iter  # Wasserstein
    mae = mae/21  # Euclidean
    # mae = mae/11.75  # Wasserstein
    print('[OK]', flush=True)
    # # Plot bar for current distribution
    # plt.bar(
    #   x_i + x_offset,
    #   mae,
    #   bar_width,
    #   label=label,
    #   linewidth=lmdk_lib.line_width
    # )
    # # Change offset for next bar
    # x_offset += bar_width
    # Plot line
    plt.plot(
      x_i,
      mae,
      label=label,
      marker=markers[d_i],
      markersize=lmdk_lib.marker_size,
      markeredgewidth=0,
      linewidth=lmdk_lib.line_width
    )
  path = str('../../rslt/lmdk_sel_cmp/' + 'lmdk_sel_cmp-norm-l')
  # path = str('../../rslt/lmdk_sel_cmp/' + 'lmdk_sel_cmp-emd-l')
  # Plot legend
  lmdk_lib.plot_legend()
  # Show plot
  # plt.show()
  # Save plot
  lmdk_lib.save_plot(path + '.pdf')



'''
  Parse arguments.

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

  # Mandatory arguments.

  # Optional arguments.
  parser.add_argument('-i', '--iter', help='The number of iterations.', 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()