import numpy as np
import scipy
import csv
import json
from math import log10, isinf, floor
from copy import deepcopy

import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from matplotlib.colors import LogNorm
from mpl_toolkits.axes_grid1 import make_axes_locatable
from matplotlib.figure import Figure
from matplotlib.ticker import MaxNLocator
import matplotlib

# Constants
output_data_path = "../../../Manuscripts/atc-paper/data/"
figure_path = "../../../Manuscripts/atc-paper/figures/"
table_path = "../../../Manuscripts/atc-paper/tables/"
text_width_in_inches = 5.95114 # Used to scale PGF plots
small_plot_size = 300
color_bar_width = 50
dpi = 1000
very_small_font_size = 5
small_font_size = 7
medium_font_size = 9
thin_line_width = 0.75
small_marker_size = 3.75

# Preliminaries

IMAGE = 0
LATEX = 1
def set_output_mode(output_mode):
	if output_mode == IMAGE:
		matplotlib.use("TkAgg") # Default value on this system, can be checked using "matplotlib.rcParams['backend']"
		matplotlib.rcParams.update(matplotlib.rcParamsDefault)
		"""matplotlib.rcParams.update({
			'font.family': 'serif',
			'pgf.rcfonts': False,
			'axes.labelsize': small_font_size,
			'axes.titlesize': medium_font_size,
			'xtick.labelsize': small_font_size,
			'ytick.labelsize': small_font_size,
			'legend.fontsize': small_font_size,
			'lines.linewidth': thin_line_width,
			'lines.markersize': small_marker_size,
		})"""
	elif output_mode == LATEX:
		# Set up Latex figure integration
		matplotlib.use("pgf")
		matplotlib.rcParams.update({
			"pgf.texsystem": "pdflatex",
			'font.family': 'serif',
			'text.usetex': True,
			'pgf.rcfonts': False,
			'axes.labelsize': small_font_size,
			'axes.titlesize': medium_font_size,
			'xtick.labelsize': small_font_size,
			'ytick.labelsize': small_font_size,
			'legend.fontsize': small_font_size,
			'lines.linewidth': thin_line_width,
			'lines.markersize': small_marker_size,
		})

def set_output_mode_image():
	set_output_mode(IMAGE)
def set_output_mode_latex():
	set_output_mode(LATEX)

def load_matrix(path):
	# Loads matrix stored in the format:
	# uint64: rows
	# uint64: cols
	# float64[]: data (column-major)
	with open(path, "rb") as f:
		rows, cols = tuple(np.fromfile(f, np.uint64, 2))
		return np.fromfile(f, np.float64, rows*cols).reshape((cols, rows)).transpose()

def load_csv(path, field_mappings):
	
	# field_mappings is a list of mappings for the identifying fields for an entry, for example [lambda x: x, lambda x: x == "1"] if the fields are Dataset and Unified quantization
	
	data = {}
	with open(path, "r") as f:
		reader = csv.reader(f, delimiter=",", quotechar="\"")
		for i, row in enumerate(reader):
			if i == 1:
				data["target_rel_errors"] = np.array(list(map(float, filter(lambda x: x != "", row))))
			elif i > 2:
				fields = []
				inner_data_dict = data
				for j, field_mapping in enumerate(field_mappings):
					fields.append(field_mapping(row[j]))
					if fields[j] not in inner_data_dict:
						inner_data_dict[fields[j]] = {}
					inner_data_dict = inner_data_dict[fields[j]]
				while row[-1] == "":
					row = row[:-1]
				inner_data_dict["actual_rel_errors"] = np.array(list(map(float, row[len(field_mappings)::2])))
				inner_data_dict["compression_factors"] = np.array(list(map(float, row[len(field_mappings) + 1::2])))
	
	return data

# Sets the width of the plot about to be saved relative to the text width, e.g. 0.5 corresponds to a plot spanning half of the text width
def set_plot_width(fig, text_width_ratio):
	old_width, old_height = fig.get_size_inches()
	new_width = text_width_ratio*text_width_in_inches
	new_height = (new_width/old_width)*old_height
	fig.set_size_inches(new_width, new_height)

# Algorithmic improvements

# Hybrid truncation

def matrix_tensor_product(matrix, tensor, mode):
	transposition = list(range(tensor.ndim))
	transposition[mode] = 0
	transposition[0] = mode
	new_tensor = tensor.transpose(transposition)
	shape = list(new_tensor.shape)
	new_tensor = (matrix @ new_tensor.reshape((new_tensor.shape[0], -1)))
	shape[0] = matrix.shape[0]
	new_tensor = new_tensor.reshape(shape)
	new_tensor = new_tensor.transpose(transposition)
	return new_tensor

def check_tensor_orthogonality(tensor, mode):
	transposition = list(range(tensor.ndim))
	transposition[mode] = 0
	transposition[0] = mode
	matrix = tensor.transpose(transposition)
	matrix = matrix.reshape((matrix.shape[0], -1))
	inner_products = matrix @ matrix.T
	norms = np.sqrt(np.diag(inner_products))
	truncation_rank = norms[norms > 0].size
	norms = norms[:truncation_rank]
	inner_products = inner_products[:truncation_rank, :truncation_rank]
	normalized_inner_products = inner_products/np.outer(norms, norms)
	ref = np.eye(*normalized_inner_products.shape, dtype=normalized_inner_products.dtype)
	return np.linalg.norm(normalized_inner_products - ref)/np.linalg.norm(ref)

def check_orthogonality_approximation(tensor, mode, factor_error):
	transposition = list(range(tensor.ndim))
	transposition[mode] = 0
	transposition[0] = mode
	matrix = tensor.transpose(transposition)
	matrix = matrix.reshape((matrix.shape[0], -1))
	exact = np.linalg.norm(factor_error @ matrix)
	norms = np.array([np.linalg.norm(matrix[i, :]) for i in range(matrix.shape[0])])
	approx = np.linalg.norm(factor_error * norms)
	return abs((approx - exact)/exact)

def sign(val):
	return (-1 if val < 0 else (1 if val > 0 else 1))

def factor_to_reflectors(factor):
	reflectors = np.empty(factor.shape, dtype=factor.dtype)
	new_factor = np.empty(factor.shape, dtype=factor.dtype)
	rows, cols = factor.shape
	R = factor.copy()
	for col in range(cols):
		reflector = R[col:, col].copy()
		reflector[0] += sign(reflector[0])*np.linalg.norm(reflector)
		reflector = reflector/np.linalg.norm(reflector)
		reflectors[col:, col] = reflector
		R[col:, col:] = R[col:, col:] - (2*reflectors[col:, col:col + 1]) @ (reflectors[col:, col:col + 1].T @ R[col:, col:])
		new_factor[:, col] = sign(R[col, col])*factor[:, col]
	return reflectors, new_factor

def reflectors_to_factor(reflectors):
	factor = np.eye(*reflectors.shape, dtype=reflectors.dtype)
	for col in range(reflectors.shape[1] - 1, -1, -1):
		factor[col:, col:] = factor[col:, col:] - (2*reflectors[col:, col:col + 1]) @ (reflectors[col:, col:col + 1].T @ factor[col:, col:])
	return factor

def save_latex_table(data, path, first_row, first_column):
	
	with open(path, "w") as f:
		f.write(first_row + "\n")
		format_string = "{:.2f}"
		def val_to_string(val):
			exponent = int(log10(abs(val)))
			string = "$" + format_string.format(val/10**exponent) + " \\cdot 10^{%s}"%exponent + "$"
			if string == "$1.00 \\cdot 10^{0}$":
				return "1"
			return "$" + format_string.format(val/10**exponent) + " \\cdot 10^{%s}"%exponent + "$"
		for row_index, row in enumerate(data):
			f.write(first_column[row_index] + " & ")
			f.write(" & ".join(map(val_to_string, row)))
			f.write(" \\\\ \\hline\n")

def generate_error_component_table():
	
	# Obtain necessary factor components
	
	# ATC output:
	"""
	[Info]      Target SSE: 48281694200.922112
	[Info]      Truncation SSE: 24038868477.297325
	[Info]      Actual output SSE (without rounding): 48545718684.087784
	Target PSNR: 60 dB
	PSNR (dB): 59.9755
	"""
	
	modes = 3
	actual_sse = 48545718684.087784
	
	sses = np.array(
		[[24038868477.300507, -3.367539788348573e-07, 6847.75626607045, -176818.28100237454, -14322.94871071243],
		[-3.367539788348573e-07, 24242746794.572357, -373988.6306946605, -1711510.6943958218, -818630.353701944],
		[6847.75626607045, -373988.6306946605, 50003539.947933815, 37143.020800519735, 74.30302914605292],
		[-176818.28100237454, -1711510.6943958218, 37143.020800519735, 215590802.31827608, -693.953168387148],
		[-14322.94871071243, -818630.353701944, 74.30302914605292, -693.953168387148, 4614498.010507083]]
	)
	
	if True:
		
		# Load factors
		full_factors = [load_matrix(output_data_path + "moffett-field-full-factor-%s"%i) for i in range(modes)]
		quantized_factors = [load_matrix(output_data_path + "moffett-field-quantized-factor-%s"%i) for i in range(modes)]
		# Orthogonality of full_factors and quantized_factors: OK
		
		# Determine factor errors
		factor_errors = []
		for mode in range(modes):
			# Error in between full_factors columns and quantized_factors columns in the beginning (first 10 columns): OK
			"""full_reflectors, full_factors[mode] = factor_to_reflectors(full_factors[mode])
			quantized_reflectors, _ = factor_to_reflectors(quantized_factors[mode])
			full_quantized_reflectors = np.empty(full_reflectors.shape, dtype=full_reflectors.dtype)
			full_quantized_reflectors[:, :quantized_reflectors.shape[1]] = quantized_reflectors
			full_quantized_reflectors[:, quantized_reflectors.shape[1]:] = full_reflectors[:, quantized_reflectors.shape[1]:]
			full_quantized_factor = reflectors_to_factor(full_quantized_reflectors)
			factor_errors.append(full_quantized_factor - full_factors[mode])"""
			full_quantized_factor = np.empty(full_factors[mode].shape, dtype=full_factors[mode].dtype)
			full_quantized_factor[:, :quantized_factors[mode].shape[1]] = quantized_factors[mode]
			full_quantized_factor[:, quantized_factors[mode].shape[1]:] = full_factors[mode][:, quantized_factors[mode].shape[1]:]
			factor_errors.append(full_quantized_factor - full_factors[mode])
		
		# Obtain necessary core components
		truncated_shape = (502, 962, 36)
		truncated_transposed_shape = (truncated_shape[2], truncated_shape[0], truncated_shape[1])
		truncated_slicing = tuple(slice(0, truncated_shape[mode]) for mode in range(modes))
		
		# Obtain truncation error
		print("Computing truncation error")
		full_core = np.fromfile("../../data/hyperspectral_images/preprocessed/moffett_field_float32", np.float32).reshape((512, 1024, 224)).astype(np.float64)
		for mode in [2, 0, 1]:
			print("Multiplying mode %s"%mode)
			full_core = matrix_tensor_product(full_factors[mode].T, full_core, mode)
		core_shape = full_core.shape
		truncated_core = np.zeros(core_shape, dtype=np.float64)
		truncated_core[truncated_slicing] = np.fromfile(output_data_path + "moffett-field-truncated-core-%sx%sx%s-2,0,1"%(truncated_transposed_shape), np.float64).reshape(truncated_transposed_shape).transpose((1, 2, 0))
		truncation_error_tensor = truncated_core - full_core
		print("Relative truncation error tensor within untruncated part (should be almost zero):", np.linalg.norm(truncation_error_tensor[truncated_slicing])/np.linalg.norm(full_core))
		del full_core
		
		# Obtain quantized core
		quantized_core = np.zeros(core_shape, dtype=np.float64)
		quantized_core[truncated_slicing] = np.fromfile(output_data_path + "moffett-field-quantized-core-%sx%sx%s-2,0,1"%(truncated_transposed_shape), np.float64).reshape(truncated_transposed_shape).transpose((1, 2, 0))
		for mode in range(modes):
			print("Core orthogonality metric along mode %s: %s"%(mode, check_tensor_orthogonality(quantized_core, mode)))
			print("Core orthogonality relative approximation error along mode %s: %s"%(mode, check_orthogonality_approximation(quantized_core, mode, factor_errors[mode])))
		
		# Obtain quantization error
		print("Computing quantization error")
		quantization_error_tensor = quantized_core - truncated_core
		del truncated_core
		
		# Construct all error components
		error_components = [truncation_error_tensor, quantization_error_tensor]
		for mode in range(modes):
			error_components.append(matrix_tensor_product(full_factors[mode].T @ factor_errors[mode], quantized_core, mode))
		
		error_components_amount = len(error_components)
		sses = np.empty((error_components_amount, error_components_amount), dtype=np.float64)
		for table_row in range(error_components_amount):
			line = ""
			for table_col in range(error_components_amount):
				sses[table_row, table_col] = np.tensordot(error_components[table_row], error_components[table_col], axes=(list(range(modes)), list(range(modes))))
				line += "\t%s"%sses[table_row, table_col]
			print(line)
	
	norms = np.sqrt(np.diag(sses))
	normalized_sses = sses/np.outer(norms, norms)
	print(normalized_sses)
	
	print("Total SSE: ", np.sum(sses))
	print("Relative error: ", abs(np.sum(sses) - actual_sse)/actual_sse)
	print("Approximated SSE: ", np.sum(np.diag(sses)))
	print("Relative error: ", abs(np.sum(np.diag(sses)) - actual_sse)/actual_sse)
	
	error_component_strings = ["$\\overline{\\ten{B}} - \\ten{B}$", "$\\delta \\ten{B}$"] + ["$U_<mode>^T \\delta U_<mode> \\times_<mode> \\widetilde{\\ten{B}}$".replace("<mode>", str(mode + 1)) for mode in range(modes)]
	first_row = "\\cline{2-6}\n& %s \\\\ \\hline"%" & ".join(error_component_strings)
	first_column = list(map(lambda string: "\\multicolumn{1}{|l|}{%s}"%string, error_component_strings))
	save_latex_table(sses, table_path + "inner-products.tex", first_row, first_column)
	save_latex_table(normalized_sses, table_path + "normalized-inner-products.tex", first_row, first_column)
	
	"""
	Relative truncation error tensor within untruncated part (should be almost zero): 5.250918748500614e-16
	Computing quantization error
		24038868477.300507	-3.367539788348573e-07	6847.75626607045	-176818.28100237454	-14322.94871071243
		-3.367539788348573e-07	24242746794.572357	-373988.6306946605	-1711510.6943958218	-818630.353701944
		6847.75626607045	-373988.6306946605	50003539.947933815	37143.020800519735	74.30302914605292
		-176818.28100237454	-1711510.6943958218	37143.020800519735	215590802.31827608	-693.953168387148
		-14322.94871071243	-818630.353701944	74.30302914605292	-693.953168387148	4614498.010507083
	[[ 1.00000000e+00 -1.39496982e-17  6.24584077e-06 -7.76703101e-05 -4.30044299e-05]
	 [-1.39496982e-17  1.00000000e+00 -3.39677759e-04 -7.48640969e-04 -2.44756762e-03]
	 [ 6.24584077e-06 -3.39677759e-04  1.00000000e+00  3.57735246e-04  4.89152089e-06]
	 [-7.76703101e-05 -7.48640969e-04  3.57735246e-04  1.00000000e+00 -2.20015267e-05]
	 [-4.30044299e-05 -2.44756762e-03  4.89152089e-06 -2.20015267e-05  1.00000000e+00]]
	Total SSE:  48545720312.58642
	Approximated SSE:  48551824112.14958
	"""

def plot_rtmss_rate_distortion():
	
	data = load_csv(output_data_path + "rtmss_rate_distortion.csv", [str, float])
	
	# Post-processing
	for dataset_name in data.keys():
		if dataset_name == "target_rel_errors":
			continue
		
		for rtmss in data[dataset_name].keys():
			
			# Replace points with zero error with target errors, since we can produce this compression factor with at most the target error
			highest_lossless_compression_factor = None
			for i, rel_error in enumerate(data[dataset_name][rtmss]["actual_rel_errors"]):
				if rel_error == 0:
					if highest_lossless_compression_factor is None:
						highest_lossless_compression_factor = data[dataset_name][rtmss]["compression_factors"][i]
					data[dataset_name][rtmss]["actual_rel_errors"][i] = data["target_rel_errors"][i]
					data[dataset_name][rtmss]["compression_factors"][i] = highest_lossless_compression_factor
		
			# Interpolate data points using linear interpolation on the logarithms of the compression factors (and relative errors)
			data[dataset_name][rtmss]["compression_factors_at_target_rel_errors"] = []
			for target_rel_error in data["target_rel_errors"]:
				actual_rel_errors = data[dataset_name][rtmss]["actual_rel_errors"]
				compression_factors = data[dataset_name][rtmss]["compression_factors"]
				
				# Find line segment for interpolation (or in the extreme cases, extrapolation)
				if target_rel_error >= actual_rel_errors[0]:
					i = 0 # Use first line segment
				elif target_rel_error < actual_rel_errors[-1]:
					i = len(actual_rel_errors) - 2 # Use last line segment
				else:
					for i in range(len(data[dataset_name][rtmss]["actual_rel_errors"]) - 1):
						# Check if this line segment contains the target PSNR
							if target_rel_error < actual_rel_errors[i] and target_rel_error >= actual_rel_errors[i + 1]:
								break
				
				x1, x2 = map(log10, actual_rel_errors[i:i + 2])
				y1, y2 = map(log10, compression_factors[i:i + 2])
				slope = (y2 - y1)/(x2 - x1)
				intercept = y1 - slope*x1
				interpolated_compression_factor = 10**(slope*log10(target_rel_error) + intercept)
				data[dataset_name][rtmss]["compression_factors_at_target_rel_errors"].append(interpolated_compression_factor)
		
		# Generate normalized compression factors
		for rtmss in data[dataset_name].keys():
			data[dataset_name][rtmss]["normalized_compression_factors"] = []
		for i, target_rel_error in enumerate(data["target_rel_errors"]):
			max_compression_factor = 0
			for rtmss in data[dataset_name].keys():
				max_compression_factor = max(max_compression_factor, data[dataset_name][rtmss]["compression_factors_at_target_rel_errors"][i])
			for rtmss in data[dataset_name].keys():
				data[dataset_name][rtmss]["normalized_compression_factors"].append(data[dataset_name][rtmss]["compression_factors_at_target_rel_errors"][i]/max_compression_factor)
	
	"""
	set_output_mode_image()
	for dataset_name in data.keys():
		if dataset_name == "target_psnrs":
			continue
		
		for rtmss in data[dataset_name].keys():
			# Plot one curve
			plt.plot(data[dataset_name][rtmss]["actual_psnrs"], data[dataset_name][rtmss]["compression_factors"], label="RTMSS %s"%rtmss, linestyle="solid" if rtmss < 1 else "dashed")
		
		plt.title(dataset_name)
		plt.legend()
		plt.xlabel("PSNR (dB)")
		plt.ylabel("Compression factor")
		plt.yscale("log")
		plt.savefig(figure_path + "rtmss-rate-distortion-" + dataset_name + ".png", bbox_inches="tight", dpi=dpi)
		plt.close()
	"""
	
	# Plot normalized comparison using interpolated values
	set_output_mode_latex()
	datasets = len(data.keys()) - 1
	fig, axes = plt.subplots(1, datasets, constrained_layout=True, figsize=((datasets*small_plot_size + color_bar_width)/dpi, small_plot_size/dpi), dpi=dpi)
	ax_index = 0
	for dataset_name in data.keys():
		if dataset_name == "target_rel_errors":
			continue
		
		rtmsses = len(data[dataset_name].keys())
		image = np.empty((rtmsses, data["target_rel_errors"].size))
		for j, rtmss in enumerate(data[dataset_name].keys()):
			image[j, :] = data[dataset_name][rtmss]["normalized_compression_factors"]
		
		ax = axes[ax_index]
		im = ax.imshow(image, cmap="gray", aspect="equal", vmin=0.5, vmax=1)
		ax.set_title(dataset_name)
		
		# Set up horizontal axis
		rel_error_ticks = data["target_rel_errors"][::8]
		ax.set_xticks(range(len(data["target_rel_errors"]))[::8])
		ax.set_xticklabels(rel_error_ticks)
		ax.set_xlabel("Relative error")
		ax.invert_xaxis()
		
		# Set up vertical axis
		rtmss_ticks = list(data[dataset_name].keys())[::2]
		ax.invert_yaxis()
		ax.set_yticks(range(len(data[dataset_name].keys()))[::2])
		ax.set_yticklabels(rtmss_ticks)
		ax.set_ylabel("RTMSS")
		
		# Set up line to indicate default RTMSS
		default_rtmss = 0.5
		rtmsses = list(data[dataset_name].keys())
		default_rtmss_height = default_rtmss/(rtmsses[1] - rtmsses[0])
		ax.plot([-0.5, len(data["target_rel_errors"]) - 0.5], [default_rtmss_height, default_rtmss_height], linestyle="--", color="black")
		
		ax_index += 1
	
	fig.colorbar(im, ax=axes)
	set_plot_width(fig, 1)
	fig.savefig(figure_path + "rtmss-normalized-rate-distortion" + ".pgf", bbox_inches="tight", dpi=dpi)

# T-HODCT
# Since these experiments were scrapped from the paper during first submission drafting, I didn't redo them for further submissions

def plot_factor_dcts():
	set_output_mode_latex()
	for apply_dct in (True, False):
		for dataset_name, order in (("isotropic1024fine-pressure-time-128", 4), ("moffett-field", 3)):
			fig, axes = plt.subplots(1, order, constrained_layout=True, figsize=((order*small_plot_size + color_bar_width)/dpi, small_plot_size/dpi), dpi=dpi)
			for factor_index in range(order):
				
				# Get data
				data = load_matrix(output_data_path + dataset_name + "-factor_" + str(factor_index))
				transformed_data = np.abs(scipy.fft.dct(data, type=2, axis=0, norm="ortho") if apply_dct else data)
				
				# Set up subplot
				ax = axes[factor_index]
				im = ax.imshow(transformed_data, cmap="gray", aspect="equal", norm=LogNorm(vmin=10**-4, vmax=10**0))
				ax.set_title("Factor %s"%factor_index)
				ax.set_xlabel("Columns")
				if factor_index == 0:
					ax.set_ylabel("Rows")
			
			fig.colorbar(im, ax=axes)
			set_plot_width(fig, (order*small_plot_size + color_bar_width)/(4*small_plot_size + color_bar_width))
			fig.savefig(figure_path + dataset_name + ("-factor-dcts.pgf" if apply_dct else "-factors.pgf"), bbox_inches="tight", dpi=dpi)

def plot_factor_dcts_poster():
	set_output_mode_image()
	for apply_dct in (True,):
		for dataset_name, order in (("isotropic1024fine-pressure-time-128", 4),):
			for factor_index in range(order):
				data = load_matrix(output_data_path + dataset_name + "-factor_" + str(factor_index))
				transformed_data = np.abs(scipy.fft.dct(data, type=2, axis=0, norm="ortho") if apply_dct else data)
				cmap = plt.cm.coolwarm
				norm = LogNorm(vmin=10**-4, vmax=10**0)
				image = cmap(norm(transformed_data))
				plt.imsave("../../../Presentations/ATC/Poster/images/" + dataset_name + "-factor-" + str(factor_index) + ("-dct" if apply_dct else "") + ".png", image)

def plot_factor_dcts2():
	set_output_mode_image()
	for apply_dct in (True, False):
		for dataset_name, order in (("isotropic1024fine-pressure-time-128", 4), ("moffett-field", 3)):
			fig, axes = plt.subplots(1, order, constrained_layout=True, figsize=((order*small_plot_size + color_bar_width)/dpi, small_plot_size/dpi), dpi=dpi)
			for factor_index in range(order):
				
				# Get data
				data = load_matrix(output_data_path + dataset_name + "-factor_" + str(factor_index))
				transformed_data = np.abs(scipy.fft.dct(data, type=2, axis=0, norm="ortho") if apply_dct else data)
				
				# Set up subplot
				ax = axes[factor_index]
				im = ax.imshow(transformed_data, cmap="gray", aspect="equal", norm=LogNorm(vmin=10**-4, vmax=10**0), interpolation="nearest")
				ax.set_title("Factor %s"%factor_index)
				ax.set_xlabel("Columns")
				if factor_index == 0:
					ax.set_ylabel("Rows")
			
			fig.colorbar(im, ax=axes)
			fig.savefig(figure_path + dataset_name + ("-factor-dcts.png" if apply_dct else "-factors.png"), bbox_inches="tight", dpi=dpi)

# Split bit plane truncation

def parse_list(string):
	return list(map(float, string.strip("[]").split(", ")))

def plot_bit_plane_statistics():
	
	set_output_mode_latex()
	
	# Parse data
	with open(output_data_path + "isotropic1024fine-pressure-time-128-output", "r") as f:
		lines = f.readlines()
		for i, line in enumerate(lines):
			if "Bit costs:" in line:
				leading_bit_costs = parse_list(lines[i + 1].split(": ")[1].strip())
				trailing_bit_costs = parse_list(lines[i + 2].split(": ")[1].strip())
			if "Efficiency:" in line:
				leading_bit_efficiencies = parse_list(lines[i + 1].split(": ")[1].strip())
				trailing_bit_efficiencies = parse_list(lines[i + 2].split(": ")[1].strip())
				break
	
	min_bit_plane = (len(leading_bit_costs) - leading_bit_costs[::-1].index(0)) + 1
	trim_bitplanes = 3
	leading_bit_costs = leading_bit_costs[:-trim_bitplanes]
	trailing_bit_costs = trailing_bit_costs[:-trim_bitplanes]
	leading_bit_efficiencies = leading_bit_efficiencies[:-trim_bitplanes]
	trailing_bit_efficiencies = trailing_bit_efficiencies[:-trim_bitplanes]
	
	for data, ylabel, figure_file_name in (([leading_bit_costs, trailing_bit_costs], "Total bit encoding cost", "bit-plane-bit-costs.pgf"), ([np.array(leading_bit_efficiencies)/np.array(trailing_bit_efficiencies)], "$\\frac{\\mbox{Leading bits encoding efficiency}}{\\mbox{Trailing bits encoding efficiency}}$", "bit-plane-efficiencies.pgf")):
		fig, ax = plt.subplots()
		minx, maxx = min_bit_plane, 63
		if len(data) == 1:
			ax.plot(data[0], marker=".")
			ax.plot([maxx, minx], [1, 1], linestyle="--", color="black")
			miny, maxy = 2e-2, 3
			ax.set_ylim(miny, maxy)
			rel_margin = 0.025
			horizontal_margin = rel_margin*(maxx - minx)
			vertical_margin = 2*rel_margin*(maxy - miny) # Roughly scale according to the plot's aspect ratio
			ax.text(maxx - horizontal_margin, 1 + vertical_margin, "Leading bits\nmore efficient", fontdict={"size": small_font_size}, horizontalalignment="left", verticalalignment="baseline", multialignment="left")
			ax.text(maxx - horizontal_margin, 1 - vertical_margin, "Trailing bits\nmore efficient", fontdict={"size": small_font_size}, horizontalalignment="left", verticalalignment="top", multialignment="left")
		elif len(data) == 2:
			ax.plot(data[0], marker=".", label="Leading bits")
			ax.plot(data[1], marker=".", label="Trailing bits",)
			ax.legend()
		ax.set_xlim(maxx, minx)
		ax.set_yscale("log")
		ax.set_xlabel("Bit plane")
		ax.set_ylabel(ylabel)
		plt.xticks()
		plt.yticks()
		ax.xaxis.set_major_locator(MaxNLocator(integer=True))
		set_plot_width(plt.gcf(), 0.45)
		fig.savefig(figure_path + figure_file_name, bbox_inches="tight")

# Core flattenings

def merge_core_flattenings():
	
	with open(output_data_path + "core_flattenings.csv", "r") as f:
		reader = csv.reader(f, delimiter=",", quotechar="\"")
		rows = [row for row in reader]
	
	new_rows = rows[:3]
	with open(output_data_path + "core_flattenings_only_z_order.csv", "r") as f:
		reader = csv.reader(f, delimiter=",", quotechar="\"")
		rows_block_start_pointer = 3
		extra_rows_block_start_pointer = 0
		extra_rows = [row for row in reader]
		for i in range(len(extra_rows) + 1):
			if i < len(extra_rows):
				row1 = extra_rows[i]
			if i > 2:
				if (i < len(extra_rows) and extra_rows[i][0] != "") or i == len(extra_rows):
					# New dataset, end previous block if there was one
					if extra_rows_block_start_pointer != 0:
						for j in range(len(rows) + 1):
							if (i < len(extra_rows) and j < len(rows) and extra_rows[i][0] == rows[j][0]) or (i == len(extra_rows) and j == len(rows)):
								# Found insertion point for block
								adjusted_extra_block = extra_rows[extra_rows_block_start_pointer:i]
								adjusted_extra_block[0][0] = ""
								new_rows = new_rows + rows[rows_block_start_pointer:j] + adjusted_extra_block
								rows_block_start_pointer = j
					extra_rows_block_start_pointer = i
	
	with open(output_data_path + "core_flattenings_merged.csv", "w") as f:
		writer = csv.writer(f)
		writer.writerows(new_rows)

def plot_core_flattenings():
	
	"""
	Handling dataset: Isotropic-PT
	Rejected relative error: 0.1, largest relative error difference: 0.0037878070491088956
	Accepted relative error: 0.03, largest relative error difference: 0.00047177747938854806
	Accepted relative error: 0.01, largest relative error difference: 0.0002496355321230035
	Accepted relative error: 0.003, largest relative error difference: 0.000143166305976377
	Accepted relative error: 0.001, largest relative error difference: 2.99949008666745e-05
	Mean gain achieved by zigzag vs. unfold: 1.0325282376766183
	Handling dataset: Moffett-Field
	Accepted relative error: 0.1, largest relative error difference: 0.0004186870752273702
	Accepted relative error: 0.03, largest relative error difference: 2.6576748666967374e-05
	Accepted relative error: 0.01, largest relative error difference: 9.981434531711628e-06
	Accepted relative error: 0.003, largest relative error difference: 2.9884248344711214e-05
	Accepted relative error: 0.001, largest relative error difference: 6.82334363333476e-05
	Mean gain achieved by zigzag vs. unfold: 1.0004219847544276
	Handling dataset: Foot
	Rejected relative error: 0.1, largest relative error difference: 0.0041383565537008535
	Accepted relative error: 0.03, largest relative error difference: 0.0007287366324584634
	Rejected relative error: 0.01, largest relative error difference: 0.0016048105314114055
	Rejected relative error: 0.003, largest relative error difference: 0.0014158498302260783
	Rejected relative error: 0.001, largest relative error difference: 0.003990962757816427
	Mean gain achieved by zigzag vs. unfold: 0.9918677618981009
	
	Aggregate mean: 1.008121739
	"""
	
	set_output_mode_latex()
	
	# Parse data
	# data format: {dataset_string: {flattening_string: {"mode_storage_orders": [string], target_rel_error_float: {"actual_rel_errors": [float], "compression_factors": [float]}}}}
	data = {}
	current_dataset_name = None
	current_flattening_name = None
	with open(output_data_path + "core_flattenings.csv", "r") as f:
		reader = csv.reader(f, delimiter=",", quotechar="\"")
		for i, row in enumerate(reader):
			if i == 1:
				# Parse target relative errors
				target_rel_errors = [float(field) for field in row if field != ""]
			elif i >= 3:
				
				# Parse data
				
				# Handle new data structures that need to be created
				if row[0] != "":
					# New dataset
					current_dataset_name = row[0]
					data[current_dataset_name] = {}
				if row[1] != "":
					# New flattening
					current_flattening_name = row[1]
					data[current_dataset_name][current_flattening_name] = {"mode_storage_orders": []}
					for target_rel_error in target_rel_errors:
						data[current_dataset_name][current_flattening_name][target_rel_error] = {"actual_rel_errors": [], "compression_factors": []}
				
				# Handle actual data
				data[current_dataset_name][current_flattening_name]["mode_storage_orders"].append(row[2])
				for i, target_rel_error in enumerate(target_rel_errors):
					data[current_dataset_name][current_flattening_name][target_rel_error]["actual_rel_errors"].append(float(row[3 + 2*i]))
					data[current_dataset_name][current_flattening_name][target_rel_error]["compression_factors"].append(float(row[4 + 2*i]))
	
	# Plot data
	miny, maxy = 0.94, 1.06
	colors = {
		0.1: "tab:blue",
		0.03: "tab:orange",
		0.01: "tab:green",
		0.003: "tab:red",
		0.001: "tab:purple"
	}
	dataset_names = list(data.keys())
	# Make sure we process the biggest plot last, because this is where we'll add the common legend
	last_dataset_name = "Isotropic-PT"
	if last_dataset_name in dataset_names:
		del dataset_names[dataset_names.index(last_dataset_name)]
		dataset_names.append(last_dataset_name)
	legend_rel_errors = set()
	for dataset_name in dataset_names:
		
		# Create one plot per dataset
		print("Handling dataset:", dataset_name)
		
		dataset_data = data[dataset_name]
		mode_storage_orders = dataset_data["Unfold"]["mode_storage_orders"] # Flattening doesn't matter here
		mode_storage_order_indices = range(len(mode_storage_orders))
		
		# Set labels for legend
		unfold_linestyle = "solid"
		zigzag_linestyle = "dashed"
		z_order_linestyle = "dotted"
		for linestyle, label in ((unfold_linestyle, "Lexicographic"), (zigzag_linestyle, "Zigzag"), (z_order_linestyle, "Z-order")):
			plt.plot([], [], color="black", marker=".", linestyle=linestyle, label=label)
		
		# Check which target errors are have sufficiently close error margins to be used
		accepted_target_rel_errors_data = {}
		for target_rel_error in target_rel_errors:
			
			# One pair of lines per relative error
			unfold_data = dataset_data["Unfold"][target_rel_error]
			zigzag_data = dataset_data["Zigzag"][target_rel_error]
			z_order_data = dataset_data["Z-order"][target_rel_error]
			
			# Check if actual relative errors are very close to each other, else the comparison might not be valid
			relative_margin = 0.0015
			merged_rel_errors = unfold_data["actual_rel_errors"] + zigzag_data["actual_rel_errors"] + z_order_data["actual_rel_errors"]
			largest_relative_error_difference = (max(merged_rel_errors) - min(merged_rel_errors))/min(merged_rel_errors)
			if largest_relative_error_difference > relative_margin:
				print("Rejected relative error: " + str(target_rel_error) + ", largest relative error difference: " + str(largest_relative_error_difference))
				continue
			print("Accepted relative error: " + str(target_rel_error) + ", largest relative error difference: " + str(largest_relative_error_difference))
			
			mean = np.mean(unfold_data["compression_factors"] + zigzag_data["compression_factors"] + z_order_data["compression_factors"])
			unfold_factors = np.array(unfold_data["compression_factors"])/mean
			zigzag_factors = np.array(zigzag_data["compression_factors"])/mean
			z_order_factors = np.array(z_order_data["compression_factors"])/mean
			accepted_target_rel_errors_data[target_rel_error] = (unfold_factors, zigzag_factors, z_order_factors)
		
		# If there are too many curves to show on the plot, only pick a few
		if len(accepted_target_rel_errors_data) > 3:
			prioritized_target_errors = [0.1, 0.01, 0.001]
			target_rel_errors_to_check = list(accepted_target_rel_errors_data.keys())
			for target_rel_error in target_rel_errors_to_check:
				if target_rel_error not in prioritized_target_errors:
					del accepted_target_rel_errors_data[target_rel_error]
		
		# Plot the curves after filtering
		for target_rel_error, (unfold_factors, zigzag_factors, z_order_factors) in accepted_target_rel_errors_data.items():
			color = colors[target_rel_error]
			# Plot curves
			plt.plot(mode_storage_order_indices, unfold_factors, marker=".", color=color, linestyle=unfold_linestyle)
			plt.plot(mode_storage_order_indices, zigzag_factors, marker=".", color=color, linestyle=zigzag_linestyle)
			plt.plot(mode_storage_order_indices, z_order_factors, marker=".", color=color, linestyle=z_order_linestyle)
			# Add labels for common legend
			legend_rel_errors.add(target_rel_error)
		
		accepted_target_rel_errors = accepted_target_rel_errors_data.keys()
		
		# Filter core sizes into new TeX file
		filtered_core_sizes_strings = []
		core_sizes = {}
		core_sizes_path = output_data_path + "core_flattenings_core_sizes_" + dataset_name
		with open(core_sizes_path, "r") as f:
			core_sizes_strings = f.read()[1:-1].split("), (")
			for core_sizes_string in core_sizes_strings:
				parts = core_sizes_string.split(": ")
				target_rel_error = float(parts[0])
				if target_rel_error in accepted_target_rel_errors:
					filtered_core_sizes_strings.append("(" + core_sizes_string + ")")
					core_sizes[target_rel_error] = [int(string) for string in parts[1][1:-1].split(", ")]
		with open(core_sizes_path + "_filtered.tex", "w") as f:
			f.write(", ".join(filtered_core_sizes_strings))
		
		# Determine "optimal" mode storage orders according to shortest-mode-first heuristic
		optimal_mode_storage_orders = set()
		for target_rel_error in core_sizes.keys():
			sizes = core_sizes[target_rel_error]
			if len(sizes) != len(set(sizes)):
				print("Warning: Optimal mode storage order is ambiguous for target relative error %s, core sizes: %s"%(target_rel_error, sizes))
			new_optimal_mode_storage_order = str(tuple(list(zip(*sorted(zip(sizes, range(len(sizes))))))[1]))
			if len(optimal_mode_storage_orders) > 0 and new_optimal_mode_storage_order not in optimal_mode_storage_orders:
				print("Warning: new optimal mode storage order %s for core sizes %s (target relative error: %s) is not part of existing set of optimal mode storage orders %s"%(new_optimal_mode_storage_order, sizes, target_rel_error, optimal_mode_storage_orders))
			optimal_mode_storage_orders.add(new_optimal_mode_storage_order)
		for optimal_mode_storage_order in optimal_mode_storage_orders:
			optimal_mode_store_order_index = mode_storage_orders.index(optimal_mode_storage_order)
			plt.plot([optimal_mode_store_order_index, optimal_mode_store_order_index], [miny, maxy], color="black", linestyle="--")
		
		# Determine geometric mean of gain achieved by zigzag vs. unfold
		total_logs = 0
		for target_rel_error in accepted_target_rel_errors:
			total_logs += log10(dataset_data["Zigzag"][target_rel_error]["compression_factors"][optimal_mode_store_order_index]/dataset_data["Unfold"][target_rel_error]["compression_factors"][optimal_mode_store_order_index])
		print("Mean gain achieved by zigzag vs. unfold: %s"%(10**(total_logs/len(accepted_target_rel_errors))))
		
		# In the case of the last dataset, add the common legend
		if dataset_name == last_dataset_name:
			for target_rel_error in sorted(list(legend_rel_errors), reverse=True):
				plt.plot([], [], color=colors[target_rel_error], linestyle="-", label="Target RE: %s"%target_rel_error)
			plt.legend(fontsize=very_small_font_size, loc="best")
		
		if len(mode_storage_orders) > 10:
			plt.xticks(mode_storage_order_indices, map(str, mode_storage_orders), rotation=45)
		plt.xlabel("Mode storage order")
		plt.ylabel("Normalized compression factor")
		plt.ylim(miny, maxy)
		set_plot_width(plt.gcf(), 0.4 if len(mode_storage_orders) <= 10 else 0.6)
		plt.savefig(figure_path + dataset_name.replace(" ", "-") + "-core-flattenings.pgf", bbox_inches="tight")
		plt.close()

# Unified quantization

def plot_unified_quantization_error_control():
	
	set_output_mode_image()
	
	data = load_csv(output_data_path + "unified_quantization_measurements.csv", [str, lambda x: x == "1"])
	
	default_linestyle = "solid"
	unified_quantization_linestyle = "dashed"
	for dataset_name, color in zip(data.keys(), ["tab:blue", "tab:orange", "tab:green"]):
		for unified_quantization in (False, True):
			plt.plot(data["target_psnrs"], np.abs(data[dataset_name][unified_quantization]["actual_psnrs"] - data["target_psnrs"]), color=color, linestyle=(unified_quantization_linestyle if unified_quantization else default_linestyle))
		plt.plot([], [], color=color, linestyle="-", label=dataset_name)
	plt.plot([], [], color="black", linestyle=default_linestyle, label="Default")
	plt.plot([], [], color="black", linestyle=unified_quantization_linestyle, label="Unified quantization")
	plt.legend()
	plt.xlabel("Target PSNR (dB)")
	plt.ylabel("|Actual PSNR (dB) - target PSNR (dB)|")
	plt.savefig(figure_path + "unified-quantization-error-control.png")
	plt.close()

# I don't know why, but you need to run the script separately for each function or else the wrong files get generated for some reason ...
#plot_rtmss_rate_distortion()
#plot_bit_plane_statistics()
plot_core_flattenings()

"""
Results from v1.1.3 (8 November 2022)
Handling dataset: Moffett-Field
Rejected relative error: 0.1, largest relative error difference: 0.09543293585846777
Accepted relative error: 0.03, largest relative error difference: 0.000818012230280406
Accepted relative error: 0.01, largest relative error difference: 0.0008991997122561248
Accepted relative error: 0.003, largest relative error difference: 0.0009173183726240585
Accepted relative error: 0.001, largest relative error difference: 0.0009951899154088707
Mean gain achieved by zigzag vs. unfold: 0.9968676560320842
Handling dataset: Foot
Rejected relative error: 0.1, largest relative error difference: 0.004545483314943215
Accepted relative error: 0.03, largest relative error difference: 0.0013413127877717726
Rejected relative error: 0.01, largest relative error difference: 0.0015382319516031152
Accepted relative error: 0.003, largest relative error difference: 0.0009543207029949285
Rejected relative error: 0.001, largest relative error difference: 0.002444602620800772
Warning: new optimal mode storage order (1, 2, 0) for core sizes [254, 246, 251] (target relative error: 0.003) is not part of existing set of optimal mode storage orders {'(2, 1, 0)'}
Mean gain achieved by zigzag vs. unfold: 0.9906452575160556
Handling dataset: Isotropic-PT
Rejected relative error: 0.1, largest relative error difference: 0.005646661352161981
Accepted relative error: 0.03, largest relative error difference: 0.001465491632834934
Accepted relative error: 0.01, largest relative error difference: 0.0010293309348922747
Accepted relative error: 0.003, largest relative error difference: 0.001016497250458146
Accepted relative error: 0.001, largest relative error difference: 0.0005162302386865557
Mean gain achieved by zigzag vs. unfold: 1.0298697966209618
Geometric mean: 1.005648007
"""

