#!/usr/bin/env python3

import sys
sys.path.insert(1, '../lib')
import argparse
import gdp
import itertools
import lmdk_bgt
import lmdk_lib
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)
  # Correlation degree (higher values means weaker correlations)
  cor_deg = np.array([.01, .1, 1.0])
  cor_lbl = ['Strong correlation', 'Moderate correlation', 'Weak correlation']
  # 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)))
  # Width of bars
  bar_width = 1/(len(dist_type) + 1)
  # For each correlation degree
  for c_i, c in enumerate(cor_deg):
    # Logging
    title = cor_lbl[c_i]
    print('(%d/%d) %s' %(c_i + 1, len(cor_deg), title), end='', flush=True)
    # The transition matrix
    p = gdp.gen_trans_mt(2, c)
    # Bar offset
    x_offset = -(bar_width/2)*(len(dist_type) - 1)
    # Initialize plot
    lmdk_lib.plot_init()
    # The x axis
    x_i = np.arange(len(lmdk_n))
    plt.xticks(x_i, ((lmdk_n/len(seq))*100).astype(int))
    plt.xlabel('Landmarks (%)')  # Set x axis label.
    x_margin = bar_width*(len(dist_type)/2 + 1)
    plt.xlim(x_i.min() - x_margin, x_i.max() + x_margin)
    # The y axis
    plt.ylabel('Temporal privacy loss')  # Set y axis label.
    plt.yscale('log')
    plt.ylim(epsilon/10, 100*len(seq))
    # plt.ylim(0, 10000)
    for d_i, d in enumerate(dist_type):
      print('.', end='', flush=True)
      # Initialization
      e = np.zeros(len(lmdk_n))
      a = np.zeros(len(lmdk_n))
      for i, n in enumerate(lmdk_n):
        for r in range(args.iter):
          # Generate landmarks
          lmdks = lmdk_lib.get_lmdks(seq, n, d)
          # Uniform budget allocation
          e_cur = lmdk_bgt.uniform(seq, lmdks, epsilon)
          _, _, a_cur = gdp.tpl_lmdk_mem(e_cur, p, p, seq, lmdks)
          # Save privacy loss
          e[i] += np.sum(e_cur)/args.iter
          a[i] += np.sum(a_cur)/args.iter
      # Set label
      label = lmdk_lib.dist_type_to_str(d_i)
      if d_i == 1:
        label = 'Skewed'
      # Plot bar for current distribution
      plt.bar(
        x_i + x_offset,
        a,
        bar_width,
        label=label,
        linewidth=lmdk_lib.line_width
      )
      # Change offset for next bar
      x_offset += bar_width
    # Plot line for no correlation
    plt.plot(
      x_i,
      e,
      linewidth=lmdk_lib.line_width,
      color='#e0e0e0',
    )
    # Plot legend
    lmdk_lib.plot_legend()
    # Show plot
    # plt.show()
    # Save plot
    lmdk_lib.save_plot(str('../../rslt/dist_cor/' + title + '.pdf'))
    print(' [OK]', flush=True)


'''
  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:
    args = parse_args()
    start_time = time.time()
    main(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()