#!/usr/bin/env python3
#
import gzip
import sys
import pickle
import matplotlib.pylab as plt
import numpy as np
#
import log_system_usage
#
def plot_cells(stream, ax, time_offset=None, clickable=False, label=None, color=None):
	num_bytes = np.cumsum(stream['length'])
	timestamps = np.asarray(stream['timestamp'])
	#
	if time_offset is not None:
		timestamps = timestamps-time_offset
	#
	return ax.step(timestamps, num_bytes/(1024**2), where='post', label=label, color=color, picker=(5 if clickable else None))[0]
#
def onresize(event):
	# matplotlib axes size only scales based on fractions (even with tight_layout), so we manually calculate a fixed padding size
	w = event.width/event.canvas.figure.dpi
	h = event.height/event.canvas.figure.dpi
	w_padding = 0.8 # in inches
	h_padding = 0.6 # in inches
	#
	for ax in event.canvas.figure.axes:
		for item in ([ax.title, ax.xaxis.label, ax.yaxis.label] + ax.get_xticklabels() + ax.get_yticklabels()):
			item.set_fontsize(11+0.5*(w*h/(5**2)))
		#
	#
	event.canvas.figure.subplots_adjust(left=w_padding/w+0.01, right=1-((w_padding/2)/w), top=1-(h_padding/h), bottom=h_padding/h+0.01)
#
def zoom_axes_cb(event):
	scale = 2**(-event.step)
	ax = event.inaxes
	mouse_x = event.xdata
	mouse_y = event.ydata
	old_xlim = ax.get_xlim()
	old_ylim = ax.get_ylim()
	#
	ax.set_xlim([old_xlim[0]*scale+mouse_x*(1-scale), old_xlim[1]*scale+mouse_x*(1-scale)])
	ax.set_ylim([old_ylim[0]*scale+mouse_y*(1-scale), old_ylim[1]*scale+mouse_y*(1-scale)])
	#
	event.canvas.draw()
#
def pick_measureme_line_cb(event, lines, target_axes=None):
	if target_axes != None and event.mouseevent.inaxes not in target_axes:
		return
	#
	this_line = event.artist
	if event.mouseevent.button == 1 and not event.mouseevent.dblclick and event.mouseevent.key is None:
		# if the mouse is single-clicked and no keyboard key is held down
		if this_line.get_visible() and this_line in lines.keys():
			this_info = [x[1] for x in lines.items() if x[0] == this_line][0]
			#
			for_legend = []
			for (line, info) in lines.items():
				if info['measureme_id'] != this_info['measureme_id']:
					line.set_visible(False)
				else:
					line.set_visible(True)
					for_legend.append(line)
				#
			#
			event.mouseevent.inaxes.legend(for_legend, [x.get_label() for x in for_legend], loc='lower right')
			event.canvas.draw_idle()
		#
	#
#
def reset_measureme_lines_cb(event, lines, target_axes=None):
	if target_axes != None and event.inaxes not in target_axes:
		return
	#
	if event.button == 1 and event.dblclick and event.key is None:
		# if the mouse is double-clicked and no keyboard key is held down
		for (line, info) in lines.items():
			if info['is_main_plot']:
				line.set_visible(True)
			else:
				line.set_visible(False)
			#
		#
		legend = event.inaxes.get_legend()
		if legend != None:
			legend.remove()
		#
		event.canvas.draw_idle()
	#
#
def mouse_pan_cb(event, pan_settings):
	if event.inaxes is not None and event.key == 'control' and event.button == 1:
		ax = event.inaxes
		pixel_to_data = ax.transData.inverted()
		current_pos = pixel_to_data.transform_point((event.x, event.y))
		last_pos = pixel_to_data.transform_point((pan_settings['start_x'], pan_settings['start_y']))
		#
		old_xlim = ax.get_xlim()
		old_ylim = ax.get_ylim()
		#
		ax.set_xlim([old_xlim[0]+(last_pos[0]-current_pos[0]), old_xlim[1]+(last_pos[0]-current_pos[0])])
		ax.set_ylim([old_ylim[0]+(last_pos[1]-current_pos[1]), old_ylim[1]+(last_pos[1]-current_pos[1])])
		#
		event.canvas.draw_idle()
	#
	pan_settings['start_x'] = event.x
	pan_settings['start_y'] = event.y
#
def subsample_plot_cb(event, lines, target_axes=None):
	if target_axes != None and event.inaxes not in target_axes:
		return
	#
	if event.key == 'control':
		num_points = 0
		range_x = event.inaxes.xaxis.get_view_interval()
		range_y = event.inaxes.yaxis.get_view_interval()
		#
		for line in [l for l in lines if l.get_visible()]:
			data_x = line.get_xdata(orig=True)
			data_y = line.get_ydata(orig=True)
			#
			num_points += ((data_x>=range_x[0]) & (data_x<=range_x[1]) & (data_y>=range_y[0]) & (data_y<=range_y[1])).sum()
			# how many points are being rendered
		#
		for (line, info) in lines.items():
			data_x = line.get_xdata(orig=True)
			data_y = line.get_ydata(orig=True)
			line.orig_x = data_x
			line.orig_y = data_y
			line.orig_drawstyle = line.get_drawstyle()
			#
			subsample_spacing = max(1, int(num_points/10000))
			# the constant can be decreased to speed up plotting on slower computers
			#
			mask = np.ones(len(data_x))
			mask[::subsample_spacing] = 0
			line.set_xdata(data_x[mask==0])
			line.set_ydata(data_y[mask==0])
			line.set_drawstyle('default')
		#
		event.canvas.draw_idle()
	#
#
def undo_subsampling_cb(event, lines, target_axes=None):
	if target_axes != None and event.inaxes not in target_axes:
		return
	#
	if event.key == 'control':
		for (line, info) in lines.items():
			line.set_xdata(line.orig_x)
			line.set_ydata(line.orig_y)
			line.set_drawstyle(line.orig_drawstyle)
		#
		event.canvas.draw_idle()
	#
#
def get_complimentary_color(color_index):
	return (color_index+1 if color_index%2==0 else color_index-1)
#
if __name__ == '__main__':
	with gzip.GzipFile(sys.argv[1], 'rb') as f:
		streams = pickle.load(f)
	#
	with gzip.GzipFile(sys.argv[2], 'rb') as f:
		system_usage = pickle.load(f)
		system_usage['timestamps'] = np.array(system_usage['timestamps'])
	#
	with gzip.GzipFile(sys.argv[3], 'rb') as f:
		numa_data = pickle.load(f)
	#
	with gzip.GzipFile(sys.argv[4], 'rb') as f:
		fingerprints = pickle.load(f)
	#
	#fig, ax = plt.subplots()#constrained_layout=True
	fig, (ax, ax_cpu_usage) = plt.subplots(2, 1, sharex=True)
	#
	start_time = min([hop[t]['timestamp'][0] for m in streams for s in streams[m] for d in streams[m][s]
	                  for hop in streams[m][s][d] for t in ('received','sent')])
	#
	lines = {}
	assigned_colors = []
	colormap = plt.get_cmap('tab20') #'tab10'
	direction_shortforms = {'forward':'fwd', 'backward':'bwd'}
	transmission_shortforms = {'received':'recv', 'sent':'sent'}
	#
	for measureme_id in streams:
		# for each circuit
		for stream_id in streams[measureme_id]:
			# for each stream
			direction = 'forward'
			for hop_index in range(len(streams[measureme_id][stream_id][direction])):
				# for each hop in the circuit (including the OP)
				for transmission in ('received', 'sent'):
					data = streams[measureme_id][stream_id][direction][hop_index][transmission]
					#
					if hop_index == len(streams[measureme_id][stream_id][direction])-1:
						guard_fingerprint = streams[measureme_id][stream_id][direction][1]['fingerprint']
						if guard_fingerprint not in assigned_colors:
							assigned_colors.append(guard_fingerprint)
						#
						if transmission == 'sent':
							color_index = assigned_colors.index(guard_fingerprint)
						else:
							color_index = get_complimentary_color(assigned_colors.index(guard_fingerprint))
						#
					else:
						color_index = hop_index*2 + ('received', 'sent').index(transmission)
					#
					is_main_plot = (hop_index == len(streams[measureme_id][stream_id][direction])-1 and transmission == 'sent')
					direction_label = direction_shortforms[direction]
					transmission_label = transmission_shortforms[transmission]
					fingerprint = str(streams[measureme_id][stream_id][direction][hop_index]['fingerprint'])
					label = 'hop={}, {:.4}, {:.4}, mid={}, sid={}, fprnt={:.6s}'.format(hop_index, direction_label, transmission_label,
					                                                               measureme_id, stream_id, fingerprint)
					line = plot_cells(data, ax, time_offset=start_time, clickable=is_main_plot, label=label, color=colormap(color_index))
					if not is_main_plot:
						line.set_visible(False)
					#
					lines[line] = {'measureme_id':measureme_id, 'stream_id':stream_id, 'direction':direction,
					               'hop_index':hop_index, 'transmission':transmission, 'is_main_plot':is_main_plot}
				#
			#
		#
	#
	guard_counter = {}
	for measureme_id in streams:
		for stream_id in streams[measureme_id]:
			direction = 'forward'
			fingerprint = streams[measureme_id][stream_id][direction][1]['fingerprint']
			if fingerprint not in guard_counter:
				guard_counter[fingerprint] = 0
			#
			guard_counter[fingerprint] += 1
		#
	#
	system_usage_timestamps = (system_usage['timestamps'][1:]+system_usage['timestamps'][:-1])/2 - start_time
	cpu_usages = {int(cpu): np.array(log_system_usage.calculate_cpu_usage_continuous(system_usage['stats']['cpus'][cpu])) for cpu in system_usage['stats']['cpus']}
	tor_cpu_usages = {nick: np.sum([cpu_usages[cpu] for cpu in numa_data[nick][1]], axis=0) for nick in numa_data}
	#
	mask = system_usage_timestamps > 0
	system_usage_timestamps = system_usage_timestamps[mask]
	tor_cpu_usages = {nick: tor_cpu_usages[nick][mask] for nick in tor_cpu_usages}
	#
	for nick in sorted(fingerprints.keys()):
		#if fingerprints[nick] != None:
		# is not an OP
		guard_circuit_count = 0
		if fingerprints[nick] in guard_counter:
			guard_circuit_count = guard_counter[fingerprints[nick]]
		elif fingerprints[nick] is None:
			guard_circuit_count = '?'
		#
		ax_cpu_usage.plot(system_usage_timestamps, tor_cpu_usages[nick]*100, label='{} / {:.6s} (guard for {} circuits)'.format(nick, str(fingerprints[nick]),
		                                                                                                              guard_circuit_count))
		#
	#
	ax_cpu_usage.set_xlabel('Time (s)')
	ax_cpu_usage.set_ylabel('CPU Usage (%)')
	ax_cpu_usage.legend()
	ax_cpu_usage.grid(linestyle=':')
	#
	ax.set_xlabel('Time (s)')
	ax.set_ylabel('Data (MiB)')
	ax.set_title(sys.argv[1])
	ax.grid(linestyle=':')
	fig.tight_layout(pad=0)
	#ax.set_ylim(0, None)
	#
	fig.canvas.mpl_connect('resize_event', onresize)
	fig.canvas.mpl_connect('scroll_event', zoom_axes_cb)
	fig.canvas.mpl_connect('pick_event', lambda event,lines=lines,axes=[ax]: pick_measureme_line_cb(event, lines, target_axes=axes))
	fig.canvas.mpl_connect('button_press_event', lambda event,lines=lines,axes=[ax]: reset_measureme_lines_cb(event, lines, target_axes=axes))
	fig.canvas.mpl_connect('motion_notify_event', lambda event,pan_settings={'start_x':0,'start_y':0}: mouse_pan_cb(event, pan_settings))
	fig.canvas.mpl_connect('key_press_event', lambda event,lines=lines,axes=[ax]: subsample_plot_cb(event, lines, target_axes=axes))
	fig.canvas.mpl_connect('key_release_event', lambda event,lines=lines,axes=[ax]: undo_subsampling_cb(event, lines, target_axes=axes))
	#
	plt.show(fig)
#