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": "side",
	"record": "record",
}
# Physical units by object
UNIT = {
	"cpu": "s",
	"energy": "J",
	"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"

# Energy per byte on the network, in W/s
# Estimated at 0.06 kWh/GB by Aslan 2018
#WS_PER_BYTE = 0.06 * 1000 * 3600 / 1024**3
WS_PER_BYTE = 0.001875 * 1000 * 3600 / 1024**3

def gnuplot_histogram(**kwargs):
	if "machine" in kwargs and kwargs["machine"] != None:
		kwargs["machine"] = ", " + kwargs["machine"]
	else:
		kwargs["machine"] = ""
	titleline = ""
	if kwargs["maketitle"]:
		titleline = 'set title font "CMU Sans Serif,12" "{object_title} by {criterion_title} ({record}, {side}{machine}) ({unit})"'.format(**kwargs)
	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
{titleline}
#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, titleline=titleline, **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"] = ""
	titleline = ""
	if kwargs["maketitle"]:
		titleline = 'set title font "CMU Sans Serif,12" "{object_title} by {criterion_title} ({record}, {side}{machine}) ({unit})"'.format(**kwargs)
	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
{titleline}
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, titleline=titleline, **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 gnuplot_stacked_histogram_grouped(**kwargs):
	if "machine" in kwargs and kwargs["machine"] != None:
		kwargs["machine"] = ", " + kwargs["machine"]
	else:
		kwargs["machine"] = ""
	titleline = ""
	if kwargs["maketitle"]:
		titleline = 'set title font "CMU Sans Serif,12" "{object_title} by {criterion_title} ({record}, {side}{machine}) ({unit})"'.format(**kwargs)
	cluster = ""
	cluster_i = 0
	for group in kwargs["groups"]:
		if cluster_i > 0:
			cluster += ",\\\n\t"
		cluster += f' newhistogram "{group}"'
		for i in range(3, kwargs["nb_functions"]+2):
			if i == 3:
				cluster += ',\\\n\t "{plots_dir}/{object}_by_{criterion}_{side}_{record}_{group}.dat"'.format(group=group, plots_dir=PLOTS_DIR, **kwargs)
			else:
				cluster += ',\\\n\t ""'
			if cluster_i == 0:
				cluster += f' using {i}:xticlabels(2) title col'
			else:
				cluster += f' using {i}:xticlabels(2) notitle'
		cluster_i += 1
	f = open("{plots_dir}/{object}_by_{criterion}_{side}_{record}_{group_by}.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}_{group_by}.png"
set boxwidth 0.9 absolute
set style fill solid 1.0 border lt -1
set style histogram rowstacked gap 0
set style data histograms
{titleline}
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 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{cluster}
#set term cairolatex pdf
set term pict2e font ",10"
set output "{plots_dir}/{object}_by_{criterion}_{side}_{record}_{group_by}.tex"
set key font ",10" spacing 0.8
replot
""".format(plots_dir=PLOTS_DIR, cluster=cluster, titleline=titleline, **kwargs).replace("aws_lc", "aws-lc"))
	f.close()
	os.system("gnuplot {plots_dir}/{object}_by_{criterion}_{side}_{record}_{group_by}.gnuplot".format(plots_dir=PLOTS_DIR, **kwargs))

def make_log_plot(logs, exp, criterion, side, obj, record, machine=None, version=None, maketitle=True):
	f = open(f"/dev/shm/plots/{obj}_by_{criterion}_{side}_{record}.dat", "w")
	ciphers = {}
	impls = []
	plain_line = None
	idle_val = None
	conv_factor = 1.0
	if obj == "energy":
		# Convert Wh to Ws
		conv_factor = 3600.0
	
	for log in logs:
		if log["exp"] == "idle":
			idle_val = float(log[COL[obj]]) / float(log["time"]) * conv_factor
		if log["exp"] != exp or log["record"] != record:
			continue
		if log["side"] == side and log["tls"] == "0":
			n = float(log.get("n", "1000"))
			plain_line = "plain {}".format((float(log[COL[obj]]) * conv_factor - idle_val * float(log["time"])) / n)
	
	if plain_line == None:
		return
	
	for log in logs:
		if log["exp"] == exp and log["record"] == record and log["side"] == side:
			if version != None and VER_LABEL[log["cipher"]] != version:
				continue
			n = float(log.get("n", "1000"))
			if version == None:
				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]]) * conv_factor - idle_val * float(log["time"])) / n
			else:
				if log[COL[criterion]] not in ciphers:
					ciphers[log[COL[criterion]]] = {}
				ciphers[log[COL[criterion]]][log["impl"]] = (float(log[COL[obj]]) * conv_factor - idle_val * float(log["time"])) / n
			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
		if version == None:
			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
				]),
			))
		else:
			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,
		maketitle=maketitle
	)

def make_log_plot_points(logs_by_target, exp, criterion, side, obj, record):
	f = open(f"/dev/shm/plots/{obj}_by_{criterion}_{side}_{record}.dat", "w")
	ciphers = {target: {} for target in logs_by_target}
	impls = []
	plain = {target: None for target in logs_by_target}
	idle_val = {target: None for target in logs_by_target}
	conv_factor = 1.0
	if obj == "energy":
		# Convert Wh to Ws
		conv_factor = 3600.0
	
	for target in logs_by_target:
		logs = logs_by_target[target]
		for log in logs:
			if log["exp"] == "idle":
				idle_val[target] = float(log[COL[obj]]) / float(log["time"]) * conv_factor
			if log["exp"] != exp or log["record"] != record:
				continue
			if log["side"] == side and log["tls"] == "0":
				n = float(log.get("n", "1000"))
				plain[target] = (float(log[COL[obj]]) * conv_factor - idle_val[target] * float(log["time"])) / n
	
	for target in plain:
		if plain[target] == None:
			return

	for target in logs_by_target:
		logs = logs_by_target[target]
		for log in logs:
			if log["exp"] == exp and log["record"] == record and log["side"] == side:
				if version != None and VER_LABEL[log["cipher"]] != version:
					continue
				n = float(log.get("n", "1000"))
				if version == None:
					ver = VER_LABEL[log["cipher"]]
					if log[COL[criterion]]+"/"+ver not in ciphers[target]:
						ciphers[target][log[COL[criterion]]+"/"+ver] = {}
					ciphers[target][log[COL[criterion]]+"/"+ver][log["impl"]] = (float(log[COL[obj]]) * conv_factor - idle_val[target] * float(log["time"])) / n
				else:
					if log[COL[criterion]] not in ciphers[target]:
						ciphers[target][log[COL[criterion]]] = {}
					ciphers[target][log[COL[criterion]]][log["impl"]] = (float(log[COL[obj]]) * conv_factor - idle_val[target] * float(log["time"])) / n
				if log["impl"] not in impls:
					impls.append(log["impl"])
	impls.sort()
	f.write("target {} none {}\n".format(criterion, " ".join([impl_title(impl) for impl in impls])))
	for target in plain:
		for cipher in ciphers[target]:
			for impl in impls:
				if impl not in ciphers[target][cipher]:
					ciphers[target][cipher][impl] = 0
			cipher_parts = cipher.split("/")
			f.write("{} {}({}) {} {}\n".format(
				target,
				ALG_LABEL[cipher_parts[0]],
				cipher_parts[1],
				plain[target],
				" ".join([
					str(ciphers[target][cipher][impl])
					for impl in impls
				]),
			))
	f.close()

def make_profile_plot(logs, exp, criterion, side, record, no_flamegraph=False, machine=None, version=None, maketitle=False):
	f = open(f"/dev/shm/plots/profile_by_{criterion}_{side}_{record}.dat", "w")
	runs = []
	functions = []
	
	for log in logs:
		if version != None and VER_LABEL[log["cipher"]] != version:
			continue
		if log["exp"] == exp and log["record"] == record and log["side"] == side and log["tls"] == "1":
			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,
		maketitle=maketitle
	)

def make_profile_plot_grouped(logs, exp, criterion, side, record, group_by, no_flamegraph=False, machine=None, version=None, maketitle=False):
	runs = []
	functions = []
	files = {}
	
	for log in logs:
		if version != None and VER_LABEL[log["cipher"]] != version:
			continue
		if log["exp"] == exp and log["record"] == record and log["side"] == side and log["tls"] == "1":
			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)
			log_group = log[group_by]
			runs.append({
				criterion: log[COL[criterion]],
				"impl": log["impl"],
				"group": log_group,
				"functions": profile_results,
			})
			if log_group not in files:
				files[log_group] = open(f"/dev/shm/plots/profile_by_{criterion}_{side}_{record}_{log_group}.dat", "w")
	for group in files:
		files[group].write("{} {} {}\n".format(group_by, criterion, " ".join(functions)))
	for run in runs:
		files[run["group"]].write("{} {} {}\n".format(
			run["group"],
			ALG_LABEL[run[criterion]],
			" ".join([
				str(run["functions"][function][0])
				for function in functions
			]),
		))
	for group in files:
		files[group].close()
	gnuplot_stacked_histogram_grouped(
		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,
		maketitle=maketitle,
		group_by=group_by,
		groups=[group for group in files]
	)

# Are CPU and energy proportional
def make_linear_regression(logs, cat="impl"):
	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[cat] not in samples_cpu:
			samples_cpu[log[cat]] = []
			samples_energy[log[cat]] = []
		samples_cpu[log[cat]].append(sample_cpu)
		samples_energy[log[cat]].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["tls"] == "0":
			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 tabulate(lines, separator=" ", endline="", ruler=True):
	widths = [0] * len(lines[0])
	for line in lines:
		for col in range(len(line)):
			if len(line[col]) > widths[col]:
				widths[col] = len(line[col])
	table = ""
	ruler_state = False
	for line in lines:
		for col in range(len(line)):
			if col > 0:
				table += separator
			table += line[col]
			rem = widths[col] - len(line[col])
			if ruler_state:
				table += " " * (rem % 4) + "  · " * (rem // 4)
			else:
				table += " " * rem
		table += endline+"\n"
		ruler_state ^= ruler
	return table

def make_summary(logs):
	plain_result_key = lambda log: (log["exp"], log["side"], log["record"])
	result_key = lambda log: (log["exp"], log["side"], log["record"], log["impl"], log["alg"], log["kex"], log["cipher"], log["ed"])

	plain_results = {}
	results = {}
	idle_val = None
	
	for log in logs:
		if log["exp"] == "idle":
			idle_val = {
				"cpu": float(log["cpu"]) / float(log["time"]),
				"energy": float(log["Wh"]) / float(log["time"]) * 3600,
				"in": float(log["bytes_in"]) / float(log["time"]),
				"out": float(log["bytes_out"]) / float(log["time"]),
			}
		if log["tls"] == "0":
			n = float(log.get("n", "1000"))
			plain_results[plain_result_key(log)] = {
				"cpu": (float(log["cpu"]) - idle_val["cpu"] * float(log["time"])) / n,
				"total_energy": float(log["Wh"]) * 3600 / n,
				"energy": (float(log["Wh"]) * 3600 - idle_val["energy"] * float(log["time"])) / n,
				"in": (float(log["bytes_in"]) - idle_val["in"] * float(log["time"])) / n,
				"out": (float(log["bytes_out"]) - idle_val["out"] * float(log["time"])) / n,
			}
		if log["exp"] != "idle" and log["tls"] == "1":
			n = float(log.get("n", "1000"))
			results[result_key(log)] = {
				"cpu": (float(log["cpu"]) - idle_val["cpu"] * float(log["time"])) / n,
				"total_energy": float(log["Wh"]) * 3600 / n,
				"energy": (float(log["Wh"]) * 3600 - idle_val["energy"] * float(log["time"])) / n,
				"in": (float(log["bytes_in"]) - idle_val["in"] * float(log["time"])) / n,
				"out": (float(log["bytes_out"]) - idle_val["out"] * float(log["time"])) / n,
			}
	
	lines = [[
		"key",
		"idle (W)",
		"no_tls (Ws/S)",
		"tls (Ws/S)",
		"tls_min (Ws/S)",
		#"tls_max (Ws/S)",
		"tls_in (MB/S)",
		"tls_out (MB/S)",
		"tls_io (MB/S)",
		"tls_net (Ws/S)",
	]]
	for k in results:
		r = results[k]
		p = plain_results[k[:3]]
		lines.append([
			"/".join([str(i) for i in k]),
			str(round(idle_val["energy"], 1)),
			str(round(p["energy"], 1)),
			str(round(r["energy"], 1)),
			str(round(r["energy"] - p["energy"], 1)) + " (" + str(round((r["energy"] - p["energy"])/r["energy"]*100, 1)) + "%)",
			#str(round(r["total_energy"] - p["total_energy"], 1)),
			str(round((r["in"] - p["in"])/1024**2, 2)) + " (" + str(round((r["in"] - p["in"])/r["in"]*100, 1)) + "%)",
			str(round((r["out"] - p["out"])/1024**2, 2)) + " (" + str(round((r["out"] - p["out"])/r["out"]*100, 1)) + "%)",
			str(round((r["in"] + r["out"] - p["in"] - p["out"])/1024**2, 2)) + " (" + str(round((r["in"] + r["out"] - p["in"] - p["out"])/(r["in"] + r["out"])*100, 1)) + "%)",
			str(round(r["energy"] - p["energy"] + (r["out"] + r["in"] - p["out"] - p["in"]) * WS_PER_BYTE, 2)),
		])
	print(tabulate(lines))

def analyze_logs(logs):
	plain_result_key = lambda log: (log["exp"], log["side"], log["record"])
	result_key = lambda log: (log["exp"], log["side"], log["record"], log["impl"], log["alg"], log["kex"], log["cipher"], log["ed"])

	plain_results = {}
	results = {}
	idle_val = None
	
	for log in logs:
		if log["exp"] == "idle":
			idle_val = {
				"cpu": float(log["cpu"]) / float(log["time"]),
				"energy": float(log["Wh"]) / float(log["time"]) * 3600,
				"in": float(log["bytes_in"]) / float(log["time"]),
				"out": float(log["bytes_out"]) / float(log["time"]),
			}
		if log["tls"] == "0":
			n = float(log.get("n", "1000"))
			plain_results[plain_result_key(log)] = {
				"cpu": (float(log["cpu"]) - idle_val["cpu"] * float(log["time"])) / n,
				"total_energy": float(log["Wh"]) * 3600 / n,
				"energy": (float(log["Wh"]) * 3600 - idle_val["energy"] * float(log["time"])) / n,
				"in": (float(log["bytes_in"]) - idle_val["in"] * float(log["time"])) / n,
				"out": (float(log["bytes_out"]) - idle_val["out"] * float(log["time"])) / n,
			}
		if log["exp"] != "idle" and log["tls"] == "1":
			n = float(log.get("n", "1000"))
			results[result_key(log)] = {
				"exp": log["exp"],
				"side": log["side"],
				"record": log["record"],
				"impl": log["impl"],
				"alg": log["alg"],
				"kex": log["kex"],
				"cipher": log["cipher"],
				"ed": log["ed"],
				"cpu": (float(log["cpu"]) - idle_val["cpu"] * float(log["time"])) / n,
				"total_energy": float(log["Wh"]) * 3600 / n,
				"energy": (float(log["Wh"]) * 3600 - idle_val["energy"] * float(log["time"])) / n,
				"in": (float(log["bytes_in"]) - idle_val["in"] * float(log["time"])) / n,
				"out": (float(log["bytes_out"]) - idle_val["out"] * float(log["time"])) / n,
			}
	
	lines = []
	for k in results:
		r = results[k]
		p = plain_results[k[:3]]
		lines.append({
			"exp": r["exp"],
			"side": r["side"],
			"record": r["record"],
			"impl": r["impl"],
			"alg": r["alg"],
			"kex": r["kex"],
			"cipher": r["cipher"],
			"ed": r["ed"],
			"tls": r,
			"plain": p,
		})
	return lines

def aggregate(items, unique, constraints=[]):
	results = {}
	for item in items:
		satisfied = True
		for constraint in constraints:
			if type(constraints[constraint]) == list:
				if item[constraint] not in constraints[constraint]:
					satisfied = False
					break
			else:
				if item[constraint] != constraints[constraint]:
					satisfied = False
					break
		if not satisfied:
			continue
		
		unique_key = tuple(item[u] for u in unique)
		if unique_key not in results:
			results[unique_key] = {u: item[u] for u in unique}
			results[unique_key]["items"] = [item]
		else:
			results[unique_key]["items"].append(item)
	
	return results

def fzfill(x, n):
	if type(x) != str:
		x = str(round(x, n))
	return x + "0" * max(0, n - len(x) + 1 + x.find("."))

def make_tx_summary(client_logs, server_logs):
	client_results = analyze_logs(client_logs)
	server_results = analyze_logs(server_logs)
	txs = []
	lines = [["exp", "record", "alg", "kex", "cipher", "ed", "client", "server", "plain (J/S)", "TLS (J/S)", "TLS only (J/S)", "TLS rel"]]
	for client in client_results:
		#print("client", client)
		if client["side"] != "client":
			continue
		for server in server_results:
			#print("server", server)
			if server["side"] != "server":
				continue
			#print(client, server)
			if client["exp"] == server["exp"] \
				and client["record"] == server["record"] \
				and client["alg"] == server["alg"] \
				and client["kex"] == server["kex"] \
				and client["cipher"] == server["cipher"] \
				and client["ed"] == server["ed"]:
				plain_io_avg = (client["plain"]["in"] + client["plain"]["out"] + server["plain"]["in"] + server["plain"]["out"]) / 2
				tls_io_avg = (client["tls"]["in"] + client["tls"]["out"] + server["tls"]["in"] + server["tls"]["out"]) / 2
				tls_only_io_avg = tls_io_avg - plain_io_avg
				plain_energy = client["plain"]["energy"] + server["plain"]["energy"] + plain_io_avg * WS_PER_BYTE
				tls_energy = client["tls"]["energy"] + server["tls"]["energy"] + tls_io_avg * WS_PER_BYTE
				tls_only_energy = tls_energy - plain_energy
				tls_rel_energy = tls_only_energy / plain_energy
				txs.append({
					"exp": client["exp"],
					"record": client["record"],
					"alg": client["alg"],
					"kex": client["kex"],
					"cipher": client["cipher"],
					"ed": client["ed"],
					"version": VER_LABEL[client["cipher"]],
					"client_impl": client["impl"],
					"server_impl": server["impl"],
					"plain_energy": plain_energy,
					"tls_energy": tls_energy,
					"tls_only_energy": tls_only_energy,
					"tls_rel_energy": tls_rel_energy,
				})
				lines.append([
					client["exp"],
					client["record"],
					client["alg"],
					client["kex"],
					client["cipher"],
					client["ed"],
					client["impl"],
					server["impl"],
					str(round(plain_energy, 2)),
					str(round(tls_energy, 2)),
					str(round(tls_only_energy, 2)),
					str(round(tls_rel_energy*100, 2))+"%",
				])
	print(tabulate(lines))

	aggregated = aggregate(txs, ["exp", "record", "alg"], {"server_impl": ["aws-lc"], "client_impl": ["aws-lc"]})
	latex = []
	for cluster_k in aggregated:
		cluster = aggregated[cluster_k]
		items_to_find = {"kex": ["X25519", "X25519MLKEM768"], "cipher": []}
		items_to_find = {
			"1.2": {"kex": "X25519", "cipher": "ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,ECDHE_RSA_WITH_AES_256_GCM_SHA384"},
			"ec": {"kex": "X25519", "cipher": "AES_256_GCM_SHA384", "ed": "0"},
			"pq": {"kex": "X25519MLKEM768", "ed": "0"},
			"ec+ed": {"kex": "X25519", "cipher": "AES_256_GCM_SHA384", "ed": "1"},
			"pq+ed": {"kex": "X25519MLKEM768", "ed": "1"},
		}
		items = {}
		for k in items_to_find:
			for item in cluster["items"]:
				ok = True
				for kk in items_to_find[k]:
					if item[kk] != items_to_find[k][kk]:
						ok = False
						break
				# named breaks please!
				if ok:
					items[k] = item
		latex.append([
			"                          ",
			cluster["record"],
			"$"+str(round(items["ec"]["plain_energy"], 2))+"$",
			("$" + str(fzfill(items["1.2"]["tls_only_energy"], 2)) + "$ ($" + str(fzfill(items["1.2"]["tls_rel_energy"]*100, 2)) + "\\%$)") if "1.2" in items else "",
			"$" + str(fzfill(items["ec"]["tls_only_energy"], 2)) + "$ ($" + str(fzfill(items["ec"]["tls_rel_energy"]*100, 2)) + "\\%$)",
			("$" + str(fzfill(items["pq"]["tls_only_energy"], 2)) + "$ ($" + str(fzfill(items["pq"]["tls_rel_energy"]*100, 2)) + "\\%$)") if "pq" in items else "",
			("$" + str(fzfill(items["ec+ed"]["tls_only_energy"], 2)) + "$ ($" + str(fzfill(items["ec+ed"]["tls_rel_energy"]*100, 2)) + "\\%$)") if "ec+ed" in items else "",
			("$" + str(fzfill(items["pq+ed"]["tls_only_energy"], 2)) + "$ ($" + str(fzfill(items["pq+ed"]["tls_rel_energy"]*100, 2)) + "\\%$)") if "pq+ed" in items else "",
		])
	print(tabulate(latex, separator=" & ", endline="\\\\ \\cline{2-8}", ruler=False))

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

def parse_logs(logfile_name):
	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)
	
	return (logs, records)

if __name__ == "__main__":
	cmd = sys.argv[1]
	logfile_name = sys.argv[2]

	logs, records = parse_logs(logfile_name)

	os.makedirs("/dev/shm/plots", exist_ok=True)

	no_flamegraph = "-f" in sys.argv
	machine = getargv("-m", None)
	maketitle = "-t" not in sys.argv
	version = getargv("-v", None)

	if cmd == "log":
		for side in ["client", "server"]:
			for record in records:
				make_log_plot(logs, "impl-cipher-ver", "cipher", side, "cpu", record, machine=machine, maketitle=maketitle, version=version)
				make_log_plot(logs, "impl-cipher-ver", "cipher", side, "energy", record, machine=machine, maketitle=maketitle, version=version)
				make_log_plot(logs, "impl-cert-ver", "cert", side, "cpu", record, machine=machine, maketitle=maketitle, version=version)
				make_log_plot(logs, "impl-cert-ver", "cert", side, "energy", record, machine=machine, maketitle=maketitle, version=version)
				make_log_plot(logs, "impl-kex-ver", "kex", side, "cpu", record, machine=machine, maketitle=maketitle, version=version)
				make_log_plot(logs, "impl-kex-ver", "kex", side, "energy", record, machine=machine, maketitle=maketitle, version=version)
				make_log_plot(logs, "zrtt", "ed", side, "cpu", record, machine=machine, maketitle=maketitle, version="1.3")
				make_log_plot(logs, "zrtt", "ed", side, "energy", record, machine=machine, maketitle=maketitle, version="1.3")
		cmp_versions(logs, ["impl-cipher-ver", "impl-cert-ver", "impl-kex-ver"], ["side", "cipher", "cert", "kex", "record"], ["cpu", "energy"])
	elif cmd == "log2":
		# Args: log2 <file1> <target1> [<file2> <target2>] [...]
		# Only use the intersection of the records
		logs_by_target = {sys.argv[3]: logs}
		records_to_remove = {}
		for i in range(4, len(sys.argv), 2):
			logs_i, records_i = parse_logs(sys.argv[i])
			for record in records:
				if record not in records_i:
					records_to_remove[record] = ()
			target_i = sys.argv[i+1]
			logs_by_target[target_i] = logs_i
		for record in records_to_remove:
			records.pop(record)
		for side in ["client", "server"]:
			for record in records:
				make_log_plot_points(logs_by_target, "impl-cipher-ver", "cipher", side, "energy", record)
				make_log_plot_points(logs_by_target, "impl-cert-ver", "cert", side, "energy", record)
				make_log_plot_points(logs_by_target, "impl-kex-ver", "kex", side, "energy", record)
		cmp_versions(logs, ["impl-cipher-ver", "impl-cert-ver", "impl-kex-ver"], ["side", "cipher", "cert", "kex", "record"], ["cpu", "energy"])
	elif cmd == "prof":
		for side in ["client", "server"]:
			for record in records:
				#make_profile_plot(logs, "impl-cipher-ver", "cipher", side, record, no_flamegraph=no_flamegraph, machine=machine, maketitle=maketitle, version="1.3")
				#make_profile_plot(logs, "impl-cert-ver", "cert", side, record, no_flamegraph=no_flamegraph, machine=machine, maketitle=maketitle, version="1.3")
				#make_profile_plot(logs, "impl-kex-ver", "kex", side, record, no_flamegraph=no_flamegraph, machine=machine, maketitle=maketitle, version="1.3")
				make_profile_plot_grouped(logs, "impl-cipher-ver", "cipher", side, record, "impl", no_flamegraph=no_flamegraph, machine=machine, maketitle=maketitle, version="1.3")
				make_profile_plot_grouped(logs, "impl-cert-ver", "cert", side, record, "impl", no_flamegraph=no_flamegraph, machine=machine, maketitle=maketitle, version="1.3")
				make_profile_plot_grouped(logs, "impl-kex-ver", "kex", side, record, "impl", no_flamegraph=no_flamegraph, machine=machine, maketitle=maketitle, version="1.3")
	elif cmd == "summary":
		make_summary(logs)
	elif cmd == "tx":
		logfile_name = sys.argv[3]
		logfile = open(logfile_name, "r")
		lines = logfile.readlines()
		logfile.close()
		colnames = lines[0].removesuffix("\n").split(" ")
		logs2 = []
		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"]] = ()
			logs2.append(log)
		make_tx_summary(logs, logs2)
	elif cmd == "correl":
		from scipy import stats
		import matplotlib.pyplot as plt
		import numpy as np
		make_linear_regression(logs, "alg")