 #
 # This file is part of ROA-Analysis
 #
 # Author: Samir Al-Sheikh (Freie Universitaet, Berlin)
 #         s.al-sheikh@fu-berlin.de
 #
 # MIT License
 #
 # Copyright (c) 2017 The ROA-Analysis authors
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
 # of this software and associated documentation files (the "Software"), to deal
 # in the Software without restriction, including without limitation the rights
 # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 # copies of the Software, and to permit persons to whom the Software is
 # furnished to do so, subject to the following conditions:
 # 
 # The above copyright notice and this permission notice shall be included in all
 # copies or substantial portions of the Software.
 # 
 # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 #

import sys, os, csv, operator, matplotlib, pandas, numpy, types
from utils.parse import parse_shell, parse_cfg, parse_window
from utils.util import cc_refactoring
from datetime import datetime
from matplotlib.ticker import NullFormatter
if parse_cfg(0).ssh_enabled: matplotlib.use('Agg')
import matplotlib as mpl
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable

options = 'i:p:d:t'
brk_max = 20 if not parse_cfg(0).ssh_enabled else 50
top_max = 100 if not parse_cfg(0).ssh_enabled else 500
stp_max = 10 if not parse_cfg(0).ssh_enabled else 25

def get_options():
  return options

def main():

  ## Parse shell arguments and RPKI Cache Server informations and build up file infrastructure
  args = parse_shell(options); cfg = parse_cfg(args)
  archive_delay = cfg.archive_an + 'Delay/'
  if not os.path.exists(archive_delay + 'Plots'): os.makedirs(archive_delay + '/Plots')

  if not args.collector:
    ccs = sorted([f for f in os.listdir(archive_delay) if f.startswith('delay.') \
                  and f.endswith('.csv')])

  elif isinstance(args.collector, types.StringTypes):
    d_file = 'delay.' + args.collector + '.csv'
    if not os.path.exists(archive_delay + d_file): exit('Delay file does not exist')
    ccs = [d_file]

  elif isinstance(args.collector, list):
    ccs = ['delay.' + c + '.csv' for c in args.collector \
            if os.path.exists(archive_delay + 'delay.' + c + '.csv')]

  result = pandas.DataFrame()
  fmt = '%Y%m%d.%H%M'

  ## Export the delay in seconds for every collector
  for cc in ccs:
    with open(archive_delay + cc, 'r') as delay_file:
      data = sorted(csv.reader(delay_file, delimiter=','), key=operator.itemgetter(0))
    data_inv = parse_window(args.interval, data, 0, 0)
    start = datetime.strptime(data_inv[0][0], fmt).strftime('%b %d ')
    end = datetime.strptime(data_inv[-1][0], fmt).strftime('until %b %d, %Y')
    df = pandas.DataFrame(data_inv); cc = cc.split('.')[1] + '(W)'
    df.columns = ['Timestamp', cc]
    df['Timestamp'] = pandas.to_datetime(df['Timestamp'], format=fmt)
    df[cc] = pandas.to_timedelta(df[cc]) / numpy.timedelta64(1, 's')
    result = pandas.concat([result, df.set_index('Timestamp')], axis=1)
  names = [cc.split('.')[1] + '(W)' for cc in ccs]
  start, end = result.index[0].strftime('%b %d '), result.index[-1].strftime('until %b %d, %Y')

  ## Get Acess data if available
  if isinstance(args.collector, list):
    acc_path = archive_delay + 'Access/access.'
    access = [acc_path + c for c in ccs if os.path.exists(acc_path + c)]
    access_df = pandas.DataFrame()
    
    for a_file in access:
      cc = a_file.split('.')[2] + '(W)'; export = []
      with open(a_file, 'r') as access_file:
        lines = sorted(csv.reader(access_file, delimiter='-'), key=operator.itemgetter(0))
        for line in lines:
          data = filter(None, [w.strip() for item in line for w in item.split('\"')])
          ex_fmt = '%d/%b/%Y:%H:%M:%S'
          export.append([datetime.strptime(data[1][1:-1].split()[0],ex_fmt), int(data[-1])])
      df = pandas.DataFrame(export)
      ac_col = cc + ' - Access-Delay'; df.columns = [ac_col, cc + ' - Latency']
      df['Timestamp'] = df[ac_col].map(lambda t: t.strftime(fmt))
      df['Timestamp'] = pandas.to_datetime(df['Timestamp'], format=fmt)
      df[ac_col] = (df[ac_col] - df['Timestamp']).astype('timedelta64[s]')
      df[cc + ' - Latency'] = df[ cc + ' - Latency'] / 1000
      access_df = pandas.concat([access_df, df.set_index('Timestamp')], axis=1)
    result = pandas.concat([result, access_df], axis=1)
    result = result.reindex_axis(sorted(result.columns), axis=1)

    ## Calculate the propagation delay = Total Delay - Access Delay - Latency
    for i, column in enumerate(result[result.columns[::3]]):
      result[column + ' - Propagation-Delay'] = result[result.columns[i * 3]] - \
      result[result.columns[i * 3 + 1]] - result[result.columns[i * 3 + 2]]
    result = result.reindex_axis(sorted(result.columns), axis=1)

    ## Except when the sum of single delays is greater than total delay (happens at c/s failure)
    for i, column in enumerate(result[result.columns[::4]]):
      prop = column + ' - Propagation-Delay'
      result.loc[result[prop] < 0, column:prop] = 0

  labels = [pandas.to_datetime(str(d)).strftime('%Y-%m-%d') for d in result.index.values]
  dates_lab = sorted(set(labels))[::3]

  ## Debug Print
  print result

  plot(result, cfg.archive_an, names, start, end, args.collector, args.dpi, dates_lab, args.tex)

def plot(result, archive_an, names, start, end, cc, dpi, dates_lab, tex):
  """ Plot the transmission delay for a collector """

  ## Set figure properties
  mpl.rcParams['figure.figsize'] = (20, 10); mpl.rcParams['figure.dpi'] = dpi
  mpl.rcParams['figure.facecolor'] = 'w'; mpl.rcParams['figure.edgecolor'] = 'k'
  if int(tex): mpl.rc('font', family = 'serif', serif = 'Latin Modern Roman', size = 24)

  ## Plotting - Access delay
  if isinstance(cc, list):
    if int(tex): mpl.rcParams['figure.figsize'] = (25, 20)
    names = list(result); ax_list = []; grid_pl = list(numpy.arange(0,len(names),1))

    ## Create the axes with the corresponding shares
    gs = mpl.gridspec.GridSpec(len(grid_pl), 1)
    for i in grid_pl:
      if i == 0: ax_list.append(plt.subplot(gs[0,0])); continue
      ax_list.append(plt.subplot(gs[i,0], sharex = ax_list[0], sharey = ax_list[0]))
   
    ## Plots labels and title
    colors = ['gray', 'sandybrown', 'lightseagreen', 'silver']
    for i, column in enumerate(result[result.columns]):
      ax = ax_list[i];
      result.plot(rot=0, y=names[i], kind='line', ax=ax, color=colors[i % len(colors)])

      ## If values of plot are greater than max -> broken plot
      d_x = -.03 if not int(tex) else -.05
      if(result[column].dropna().max() > brk_max):
        ax.set_ylim(-2, brk_max); ax.spines['top'].set_visible(False)
        ax.legend_.remove(); ax.get_yaxis().set_label_coords(d_x,1)
        ax2 = make_axes_locatable(ax).new_vertical(size="100%", pad=0.1)
        plt.gcf().add_axes(ax2)
        ax2 = result.plot(rot=0, y=names[i], kind='line', ax=ax2, color=colors[i % len(colors)])
        ax2.set_ylim(brk_max, top_max); ax2.tick_params(bottom="off", labelbottom='off')
        ax2.spines['bottom'].set_visible(False); ax2.xaxis.grid(alpha=0.75)
        if int(tex):
          ax2.tick_params(axis='x', pad=10);
          ax2.set_xticks(dates_lab); ax2.set_xticklabels([d[-5:] for d in dates_lab]);
        else:
          ax2.set_xticks(dates_lab); ax.set_xticklabels(dates_lab, ha='center')
        ax2.xaxis.label.set_visible(False); ax2.yaxis.label.set_visible(False)
        ax2.legend(loc='upper right'); d = .002 
        kwargs = dict(transform=ax2.transAxes, color='k', clip_on=False)
        ax2.plot((-d,+d), (-d,+d), **kwargs); ax2.plot((1-d,1+d), (-d,+d), **kwargs)
        kwargs.update(transform=ax.transAxes)
        ax.plot((-d, +d), (1-d,1+d), **kwargs); ax.plot((1-d,1+d), (1-d,1+d), **kwargs)

      ## Plots lower than max
      if(result[column].dropna().max() < brk_max):
        ax.set_ylim(-2, brk_max);
        ax.yaxis.set_ticks(numpy.arange(0, brk_max + 1, stp_max))
        ax.legend(loc='upper right'); ax.get_yaxis().set_label_coords(d_x,.5)

      ax.xaxis.grid(alpha=0.75)
      if int(tex):
        ax.tick_params(axis='x', pad=10);
        ax.set_xticks(dates_lab); ax.set_xticklabels([d[-5:] for d in dates_lab], ha='center');
      else:
        ax.xaxis.set_major_formatter(mpl.dates.DateFormatter('%Y-%m-%d'))
        ax.set_xticks(dates_lab); ax.set_xticklabels(dates_lab, ha='center')

    ax1 = plt.gcf().add_subplot(111, frameon=False)
    ax1_1 = ax1.twinx()
    y = 1.01 if int(tex) else 1; nl = '\n' if int(tex) else ''
    plt.title('Total-, Access-, Server- & Propagation-Delay for the ' + \
              'web-based collectors - ' +nl+ ', '.join(cc) + ' - ' + start + end, y=y)
    d_x = -.1 if int(tex) else -.05
    ax1.set_ylabel('Delay [s]'); ax1.get_yaxis().set_label_coords(d_x,.5)

  ## Plotting - Normal delay
  else:
    ax = result.plot(subplots=True, rot=0, y=names, kind='line')
    ax1 = plt.gcf().add_subplot(111, frameon=False)
    y = 1.03 if int(tex) else 1
    plt.title('Delay for all web-based collectors' + ' - ' + start + end, y=y)
    if int(tex): names = [n.replace('Web', 'W') for n in cc_refactoring(names)]
    for i, s in enumerate(ax):
      s.legend([names[i]], loc='upper right')
      s.xaxis.set_minor_formatter(NullFormatter())
      s.set_xticks(dates_lab); s.xaxis.grid(alpha=0.75);
      if int(tex):
        s.tick_params(axis='x', pad=10);
        s.set_xticks(dates_lab); s.set_xticklabels([d[-5:] for d in dates_lab]);
      else: s.set_xticklabels(dates_lab);
      for tick in s.xaxis.get_major_ticks(): 
        tick.label1.set_horizontalalignment('center')

  ## Hide the outline subplot
  ax1.spines['top'].set_color('none'); ax1.spines['bottom'].set_color('none')
  ax1.spines['left'].set_color('none'); ax1.spines['right'].set_color('none')
  ax1.tick_params(labelcolor='w', top='off', bottom='off', left='off', right='off')
  ax1.get_yaxis().set_label_coords(-.05 if int(tex) else -.03, 0.5)
  for xa in plt.gcf().get_axes(): xa.minorticks_off()
  for label in ax1.yaxis.get_ticklabels(): label.set_visible(False)
  for label in ax1.xaxis.get_ticklabels(): label.set_visible(False)

  ## Export Plot as file
  if isinstance(cc, list):
    ax1_1.spines['top'].set_color('none'); ax1_1.spines['bottom'].set_color('none')
    ax1_1.spines['left'].set_color('none'); ax1_1.spines['right'].set_color('none')
    ax1_1.tick_params(top='off', bottom='off', left='off', right='off')
    for label in ax1_1.yaxis.get_ticklabels(): label.set_visible(False)
    for label in ax1_1.xaxis.get_ticklabels(): label.set_visible(False)
    print '\nPlot: ' + archive_an + 'Delay/Plots/plots.delay.access.pdf'
    plt.savefig(archive_an + 'Delay/Plots/plots.delay.access.pdf', bbox_inches='tight') 
  else:
    plt.ylabel('Delay [seconds]')
    print '\nPlot: ' + archive_an + 'Delay/Plots/plots.delay.all.pdf'
    plt.savefig(archive_an + 'Delay/Plots/plots.delay.all.pdf', bbox_inches='tight') 
  return

if __name__ == '__main__':
    main()
