import os, sys
import profile

# Nice labels for algorithms of all kinds
ALG_LABEL = {
	"AES_128_GCM_SHA256": "aes128",
	"AES_256_GCM_SHA384": "aes256",
	"CHACHA20_POLY1305_SHA256": "chacha20",
	"ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,ECDHE_RSA_WITH_AES_128_GCM_SHA256": "aes128",
	"ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,ECDHE_RSA_WITH_AES_256_GCM_SHA384": "aes256",
	"ECDHE_ECDSA_WITH_CHACHA20_POLY1305_SHA256,ECDHE_RSA_WITH_CHACHA20_POLY1305_SHA256": "chacha20",
	"prime256v1": "p256",
	"secp384r1": "p384",
	"rsa2048": "rsa2048",
	"rsa3072": "rsa3072",
	"rsa4096": "rsa4096",
	"X25519": "x25519",
	"SECP256R1": "p256",
	"SECP384R1": "p384",
	"X25519MLKEM768": "x25519mlkem",
	"SECP256R1MLKEM768": "p256mlkem",
	"MLKEM768": "mlkem",
	"0": "Off",
	"1": "On",
}

# Nice labels for TLS versions using ciphers
VER_LABEL = {
	"AES_128_GCM_SHA256": "1.3",
	"AES_256_GCM_SHA384": "1.3",
	"CHACHA20_POLY1305_SHA256": "1.3",
	"ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,ECDHE_RSA_WITH_AES_128_GCM_SHA256": "1.2",
	"ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,ECDHE_RSA_WITH_AES_256_GCM_SHA384": "1.2",
	"ECDHE_ECDSA_WITH_CHACHA20_POLY1305_SHA256,ECDHE_RSA_WITH_CHACHA20_POLY1305_SHA256": "1.2"
}

# Titles for measured quantities
OBJ_TITLE = {
	"cpu": "CPU time",
	"energy": "energy consumption",
	"profile": "time profile",
}
# Logfile column names
COL = {
	"cpu": "cpu",
	"energy": "Wh",
	"cipher": "cipher",
	"cert": "alg",
	"kex": "kex",
	"ed": "ed",
	"side": "setup",
	"record": "record",
}
# Physical units by object
UNIT = {
	"cpu": "s",
	"energy": "W",
	"profile": "samples",
}
# Titles for criteria
CRITERION_TITLE = {
	"cipher": "cipher",
	"cert": "signature algorithm",
	"kex": "key exchange",
	"ed": "0-RTT",
}

def impl_title(impl):
	# Gnuplot does not escape underscores when generating tex
	return impl.replace("_", "-")

# Where gnuplot files, data files and images are output
PLOTS_DIR = "/dev/shm/plots"

def gnuplot_histogram(**kwargs):
	if "machine" in kwargs and kwargs["machine"] != None:
		kwargs["machine"] = ", " + kwargs["machine"]
	else:
		kwargs["machine"] = ""
	cluster = ""
	for i in range(kwargs["nb_impls"]-1):
		cluster += """, "" using {}:xticlabels(1) title col""".format(i+4)
	f = open("{plots_dir}/{object}_by_{criterion}_{side}_{record}.gnuplot".format(plots_dir=PLOTS_DIR, **kwargs), "w")
	f.write("""\
set terminal pngcairo enhanced font "CMU Sans Serif,11" fontscale 1.0 size 800, 600
set output "{plots_dir}/{object}_by_{criterion}_{side}_{record}.png"
set boxwidth 0.9 absolute
set style fill solid 1.0 border lt -1
set style histogram clustered gap 1 title textcolor lt -1
set style data histograms
set title font "CMU Sans Serif,12" "{object_title} by {criterion_title} ({record}, {side}{machine}) ({unit})"
#set xtics border in scale 0,0 nomirror rotate by -45 autojustify
set xtics border in scale 0,0 nomirror autojustify
#set key fixed right top vertical Right noreverse noenhanced autotitle nobox
set key fixed left top vertical Left reverse noenhanced autotitle nobox
set colorbox vertical origin screen 0.9, 0.2 size screen 0.05, 0.6 front noinvert bdefault
set xrange [ * : * ] noreverse writeback
set yrange [ 0 : * ]
set grid y lt 1 lw .75 lc "gray"
plot \
	newhistogram "", "{plots_dir}/{object}_by_{criterion}_{side}_{record}.dat" using 2:xticlabels(1) notitle col, \
	newhistogram "", "{plots_dir}/{object}_by_{criterion}_{side}_{record}.dat" using 3:xticlabels(1) title col{cluster}
set term cairolatex pdf
set output "{plots_dir}/{object}_by_{criterion}_{side}_{record}.tex"
replot
""".format(plots_dir=PLOTS_DIR, cluster=cluster, **kwargs))
	f.close()
	os.system("gnuplot {plots_dir}/{object}_by_{criterion}_{side}_{record}.gnuplot".format(plots_dir=PLOTS_DIR, **kwargs))

def gnuplot_stacked_histogram(**kwargs):
	if "machine" in kwargs and kwargs["machine"] != None:
		kwargs["machine"] = ", " + kwargs["machine"]
	else:
		kwargs["machine"] = ""
	cluster = ""
	#for i in range(kwargs["nb_impls"]-1):
	#	cluster += """, "" using {}:xticlabels(1) title col""".format(i+4)
	f = open("{plots_dir}/{object}_by_{criterion}_{side}_{record}.gnuplot".format(plots_dir=PLOTS_DIR, **kwargs), "w")
	f.write("""\
set terminal pngcairo enhanced font "CMU Sans Serif,11" fontscale 1.0 size 800, 600
set output "{plots_dir}/{object}_by_{criterion}_{side}_{record}.png"
set boxwidth 0.9 absolute
set style fill solid 1.0 border lt -1
set style histogram rowstacked
set style data histograms
set title font "CMU Sans Serif,12" "{object_title} by {criterion_title} ({record}, {side}{machine}) ({unit})"
set xtics border in scale 0,0 nomirror noenhanced rotate by 30 right
set lmargin 9
set rmargin 1
set bmargin 5
set tmargin 2.5
set key fixed left top vertical Left noenhanced autotitle nobox invert reverse opaque
set colorbox vertical origin screen 0.9, 0.2 size screen 0.05, 0.6 front noinvert bdefault
set xrange [ * : * ] noreverse writeback
set yrange [ 0 : * ]
set grid y lt 1 lw .75 lc "gray"
plot for [i=2:{nb_functions}] "{plots_dir}/{object}_by_{criterion}_{side}_{record}.dat" using i:xticlabels(1) title col
#set term cairolatex pdf
set term pict2e font ",10"
set output "{plots_dir}/{object}_by_{criterion}_{side}_{record}.tex"
set key font ",10" spacing 0.8
replot
""".format(plots_dir=PLOTS_DIR, cluster=cluster, **kwargs).replace("aws_lc", "aws-lc"))
	f.close()
	os.system("gnuplot {plots_dir}/{object}_by_{criterion}_{side}_{record}.gnuplot".format(plots_dir=PLOTS_DIR, **kwargs))

def make_log_plot(logs, exp, criterion, side, obj, record, machine=None, version=None):
	f = open(f"/dev/shm/plots/{obj}_by_{criterion}_{side}_{record}.dat", "w")
	ciphers = {}
	impls = []
	plain_line = None
	idle_val = None
	
	for log in logs:
		if log["exp"] == "idle":
			idle_val = float(log[COL[obj]]) / float(log["time"])
		if log["exp"] != exp or log["record"] != record:
			continue
		if log["setup"] == "none":
			plain_line = "plain {}".format(float(log[COL[obj]]) - idle_val * float(log["time"]))
	
	if plain_line == None:
		return
	
	for log in logs:
		if log["exp"] == exp and log["record"] == record and log["setup"] == side:
			#ver = VER_LABEL[log["cipher"]]
			#if log[COL[criterion]]+"/"+ver not in ciphers:
			#	ciphers[log[COL[criterion]]+"/"+ver] = {}
			#ciphers[log[COL[criterion]]+"/"+ver][log["impl"]] = float(log[COL[obj]]) - idle_val * float(log["time"])
			if version != None and VER_LABEL[log["cipher"]] != version:
				continue
			if log[COL[criterion]] not in ciphers:
				ciphers[log[COL[criterion]]] = {}
			ciphers[log[COL[criterion]]][log["impl"]] = float(log[COL[obj]]) - idle_val * float(log["time"])
			if log["impl"] not in impls:
				impls.append(log["impl"])
	impls.sort()
	f.write("{} none {}\n".format(criterion, " ".join([impl_title(impl) for impl in impls])))
	f.write(plain_line+" -"*len(impls)+"\n")
	for cipher in ciphers:
		for impl in impls:
			if impl not in ciphers[cipher]:
				ciphers[cipher][impl] = 0
		#cipher_parts = cipher.split("/")
		#f.write("{}({}) - {}\n".format(
		#	ALG_LABEL[cipher_parts[0]],
		#	cipher_parts[1],
		#	" ".join([
		#		str(ciphers[cipher][impl])
		#		for impl in impls
		#	]),
		#))
		f.write("{} - {}\n".format(
			ALG_LABEL[cipher],
			" ".join([
				str(ciphers[cipher][impl])
				for impl in impls
			]),
		))
	f.close()
	gnuplot_histogram(
		object=obj,
		criterion=criterion,
		side=side,
		object_title=OBJ_TITLE[obj],
		criterion_title=CRITERION_TITLE[criterion],
		unit=UNIT[obj],
		nb_impls=len(impls),
		record=record,
		machine=machine
	)

def make_profile_plot(logs, exp, criterion, side, record, no_flamegraph=False, machine=None):
	f = open(f"/dev/shm/plots/profile_by_{criterion}_{side}_{record}.dat", "w")
	runs = []
	functions = []
	
	for log in logs:
		if log["exp"] == exp and log["record"] == record and log["setup"] == side:
			svg_filename = log["prof"] + ".svg"
			if not no_flamegraph:
				os.system("flamegraph --perfdata {} -o {}".format(log["prof"], svg_filename))
			try:
				profile_results = profile.extract_from_file(svg_filename)
			except FileNotFoundError:
				print(f"Cannot read {svg_filename}")
				continue
			print(profile_results)
			for function in profile_results:
				if function not in functions:
					functions.append(function)
			runs.append({
				criterion: log[COL[criterion]],
				"impl": log["impl"],
				"functions": profile_results,
			})
	f.write("{} {}\n".format(criterion, " ".join(functions)))
	for run in runs:
		"""f.write("\"{} {}({})\" {}\n".format(
			run["impl"],
			ALG_LABEL[run[criterion]],
			VER_LABEL[log["cipher"]],
			" ".join([
				str(run["functions"][function][0])
				for function in functions
			]),
		))"""
		f.write("\"{} {}\" {}\n".format(
			impl_title(run["impl"]),
			ALG_LABEL[run[criterion]],
			" ".join([
				str(run["functions"][function][0])
				for function in functions
			]),
		))
	f.close()
	gnuplot_stacked_histogram(
		object="profile",
		criterion=criterion,
		side=side,
		object_title=OBJ_TITLE["profile"],
		criterion_title=CRITERION_TITLE[criterion],
		unit=UNIT["profile"],
		record=record,
		nb_functions=len(functions)+1,
		machine=machine
	)

# Are CPU and energy proportional
def make_linear_regression(logs):
	idle_cpu = None
	idle_energy = None
	for log in logs:
		if log["exp"] == "idle":
			idle_cpu = float(log["cpu"]) / float(log["time"])
			idle_energy = float(log["Wh"]) / float(log["time"])
			break

	samples_cpu = {"global":[]}
	samples_energy = {"global":[]}
	for log in logs:
		if log["impl"] == "-":
			continue
		sample_cpu = float(log["cpu"]) - idle_cpu * float(log["time"])
		sample_energy = float(log["Wh"]) - idle_energy * float(log["time"])
		samples_cpu["global"].append(sample_cpu)
		samples_energy["global"].append(sample_energy)
		if log["impl"] not in samples_cpu:
			samples_cpu[log["impl"]] = []
			samples_energy[log["impl"]] = []
		samples_cpu[log["impl"]].append(sample_cpu)
		samples_energy[log["impl"]].append(sample_energy)
	print("Pearson correlation coefficients (energy/cpu)")
	results = {}
	for impl in samples_cpu:
		res = stats.linregress(samples_cpu[impl], samples_energy[impl])
		print(impl, "\t", res.rvalue)
		results[impl] = res
		if impl != "global":
			plt.plot(samples_cpu[impl], samples_energy[impl], 'o', label=impl)
	#plt.plot(samples_cpu["global"], samples_energy["global"], 'o', label='samples')
	plt.plot(samples_cpu["global"], res.intercept + res.slope*np.array(samples_cpu["global"]), 'r', label='fitted line')
	plt.xlabel("CPU (s)")
	plt.ylabel("Energy (Wh)")
	plt.legend()
	#plt.show()
	plt.savefig(f"{PLOTS_DIR}/correlation_energy_cpu.png")

# Measure relative difference between TLS versions
def cmp_versions(logs, exps, criteria, objs):
	ciphers = {}
	idle_val = None
	
	for log in logs:
		if log["exp"] == "idle":
			idle_val = {obj:float(log[COL[obj]]) / float(log["time"]) for obj in objs}
	
	for log in logs:
		if log["exp"] not in exps or log["setup"] == "none":
			continue
		ver = VER_LABEL[log["cipher"]]
		if ver not in ciphers:
			ciphers[ver] = {}
		key = []
		for criterion in criteria:
			key.append(ALG_LABEL.get(log[COL[criterion]], log[COL[criterion]]))
		key = "/".join(key)
		if key not in ciphers[ver]:
			ciphers[ver][key] = log
	
	diff_rel_max = {obj:0.0 for obj in objs}
	diff_rel_sum = {obj:0.0 for obj in objs}
	diff_rel_num = {obj:0 for obj in objs}
	for key in ciphers["1.2"]:
		if key not in ciphers["1.3"]:
			continue
		log12 = ciphers["1.2"][key]
		log13 = ciphers["1.3"][key]
		for obj in objs:
			val12 = float(log12[COL[obj]]) - idle_val[obj] * float(log12["time"])
			val13 = float(log13[COL[obj]]) - idle_val[obj] * float(log13["time"])
			# Difference relative to the mean of the two values
			try:
				diff_rel = abs(val12 - val13) / ((val12 + val13) / 2)
			except ZeroDivisionError:
				continue
			diff_rel_max[obj] = max(diff_rel_max[obj], diff_rel)
			diff_rel_sum[obj] += diff_rel
			diff_rel_num[obj] += 1
	print("Diff rel max: ", diff_rel_max)
	try:
		diff_rel_avg = {obj:diff_rel_sum[obj]/diff_rel_num[obj] for obj in objs}
		print("Diff rel avg: ", diff_rel_avg)
	except ZeroDivisionError:
		pass

def getargv(arg:str, default="", n:int=1, args:list=sys.argv):
	if arg in args and len(args) > args.index(arg)+n:
		return args[args.index(arg)+n]
	else:
		return default

if __name__ == "__main__":
	cmd = sys.argv[1]
	logfile_name = sys.argv[2]
	logfile = open(logfile_name, "r")
	lines = logfile.readlines()
	logfile.close()

	colnames = lines[0].removesuffix("\n").split(" ")

	logs = []
	records = {}
	for line in lines[1:]:
		cols = line.removesuffix("\n").split(" ")
		log = {}
		for col in range(len(cols)):
			log[colnames[col]] = cols[col]
		if log["record"] != "-":
			records[log["record"]] = ()
		logs.append(log)
	
	os.makedirs("/dev/shm/plots", exist_ok=True)

	no_flamegraph = "-f" in sys.argv

	machine = getargv("-m", None)

	if cmd == "log":
		cmp_versions(logs, ["impl-cipher-ver", "impl-cert-ver", "impl-kex-ver"], ["side", "cipher", "cert", "kex", "record"], ["cpu", "energy"])
		for side in ["client", "server"]:
			for record in records:
				make_log_plot(logs, "impl-cipher-ver", "cipher", side, "cpu", record, machine=machine, version="1.3")
				make_log_plot(logs, "impl-cipher-ver", "cipher", side, "energy", record, machine=machine, version="1.3")
				make_log_plot(logs, "impl-cert-ver", "cert", side, "cpu", record, machine=machine, version="1.3")
				make_log_plot(logs, "impl-cert-ver", "cert", side, "energy", record, machine=machine, version="1.3")
				make_log_plot(logs, "impl-kex-ver", "kex", side, "cpu", record, machine=machine, version="1.3")
				make_log_plot(logs, "impl-kex-ver", "kex", side, "energy", record, machine=machine, version="1.3")
				make_log_plot(logs, "zrtt", "ed", side, "cpu", record, machine=machine, version="1.3")
				make_log_plot(logs, "zrtt", "ed", side, "energy", record, machine=machine, version="1.3")
	elif cmd == "prof":
		for side in ["client-local", "server-local"]:
			for record in records:
				make_profile_plot(logs, "impl-cipher-ver", "cipher", side, record, no_flamegraph=no_flamegraph, machine=machine)
				make_profile_plot(logs, "impl-cert-ver", "cert", side, record, no_flamegraph=no_flamegraph, machine=machine)
				make_profile_plot(logs, "impl-kex-ver", "kex", side, record, no_flamegraph=no_flamegraph, machine=machine)
	elif cmd == "correl":
		from scipy import stats
		import matplotlib.pyplot as plt
		import numpy as np
		make_linear_regression(logs)