import json
import os
import sys
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
import mpl_toolkits.axisartist as axisartist
import numpy as np
from math import log10, ceil, isnan, sqrt
from PIL import Image

from datasets import *
from compressors import *
from util import *

# Constants
results_path = "../results/"
figure_path = "../../../Manuscripts/atc-paper/figures/"
table_path = "../../../Manuscripts/atc-paper/tables/"
poster_figure_path = "../../../Presentations/ATC/Poster/images/"
numa_seminar_figure_path = "../../../Presentations/ATC/NUMA seminar/figures/"
midterm_report_figure_path = "../../../Administration/Midterm report/figures/"
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
very_small_marker_size = 2

# Pre-define line colors and styles for compressors
color_wheel = ["tab:blue", "tab:orange", "tab:green", "tab:red", "tab:purple", "tab:brown", "tab:pink"]
colors = {
	ATC.name: "tab:blue",
	generate_atc_compressor(blas=False).name: "tab:blue",
	generate_atc_compressor(optimization_level=0).name: "tab:green",
	generate_atc_compressor(optimization_level=1).name: "tab:red",
	generate_atc_compressor(optimization_level=2).name: "tab:purple",
	generate_atc_compressor(optimization_level=3).name: "tab:brown",
	generate_atc_compressor(optimization_level=4).name: "tab:pink",
	generate_atc_compressor(optimization_level=5).name: "tab:gray",
	generate_atc_compressor(optimization_level=5, blas=True, run=1).name: "tab:pink",
	generate_atc_compressor(optimization_level=5, blas=False, run=1).name: "tab:pink",
	generate_atc_compressor(optimization_level=5, blas=True, run=2).name: "tab:pink",
	generate_atc_compressor(optimization_level=5, blas=False, run=2).name: "tab:pink",
	generate_atc_compressor(blas=True, run=1).name: "tab:blue",
	generate_atc_compressor(blas=False, run=1).name: "tab:blue",
	generate_atc_compressor(blas=True, run=2).name: "tab:blue",
	generate_atc_compressor(blas=False, run=2).name: "tab:blue",
	generate_atc_compressor(blas=True, run=3).name: "tab:orange",
	generate_atc_compressor(blas=False, run=3).name: "tab:orange",
	generate_atc_compressor(blas=True, run=4).name: "tab:green",
	generate_atc_compressor(blas=False, run=4).name: "tab:green",
	generate_atc_compressor(blas=True, run=5).name: "tab:red",
	generate_atc_compressor(blas=False, run=5).name: "tab:red",
	generate_atc_compressor(blas=True, run=6).name: "tab:purple",
	generate_atc_compressor(blas=False, run=6).name: "tab:purple",
	generate_atc_compressor(blas=True, epcm=False, run=6).name: "tab:purple",
	generate_atc_compressor(blas=False, epcm=False, run=6).name: "tab:purple",
	generate_atc_compressor(blas=True, epcm=False, run=7).name: "tab:brown",
	generate_atc_compressor(blas=False, epcm=False, run=7).name: "tab:brown",
	generate_atc_compressor(blas=True, epcm=False, run=8).name: "tab:pink",
	generate_atc_compressor(blas=False, epcm=False, run=8).name: "tab:pink",
	generate_atc_compressor(blas=True, epcm=False, run=9).name: "tab:blue",
	generate_atc_compressor(blas=False, epcm=False, run=9).name: "tab:blue",
	generate_atc_compressor(hwm="csn").name: "tab:blue",
	TTHRESH.name: "tab:orange",
	x265.name: "tab:green",
	SZ.name: "tab:red",
	ZFP.name: "tab:purple",
	FPZIP.name: "tab:brown", # Not shown anymore in newer versions of the paper, not competitive
	TuckerMPI.name: "tab:brown",
	generate_tuckermpi_compressor(single_threaded=True).name: "tab:brown",
	generate_tuckermpi_compressor(shared_memory_flags=False).name: "tab:brown"
}
for rtmss in [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]:
	colors[generate_atc_compressor(rtmss=rtmss).name] = "tab:blue"
for dtmss in [0.1, 0.3, 0.5, 0.7, 0.9]:
	colors[generate_atc_compressor(dtmss=dtmss).name] = "tab:blue"
for epcm in (True, False):
	colors[generate_atc_compressor(blas=True, epcm=epcm).name] = "tab:blue"
linestyles = {
	ATC.name: "solid",
	generate_atc_compressor(rtmss=0.0).name: "dotted",
	generate_atc_compressor(rtmss=0.1).name: "dotted",
	generate_atc_compressor(rtmss=0.2).name: "dashed",
	generate_atc_compressor(rtmss=0.3).name: "dashed",
	generate_atc_compressor(rtmss=0.4).name: "solid",
	generate_atc_compressor(rtmss=0.6).name: "dashdot",
	generate_atc_compressor(rtmss=0.7).name: "dashdot",
	generate_atc_compressor(rtmss=0.8).name: (0, (3, 1, 1, 1, 1, 1)),
	generate_atc_compressor(rtmss=0.9).name: (0, (3, 1, 1, 1, 1, 1)),
	generate_atc_compressor(dtmss=0.1).name: "dotted",
	generate_atc_compressor(dtmss=0.3).name: "dashed",
	generate_atc_compressor(dtmss=0.5).name: "solid",
	generate_atc_compressor(dtmss=0.7).name: "dashdot",
	generate_atc_compressor(dtmss=0.9).name: (0, (3, 1, 1, 1, 1, 1)),
	generate_atc_compressor(blas=False, epcm=True).name: "dashed",
	generate_atc_compressor(blas=False, epcm=False).name: "dotted",
	generate_atc_compressor(optimization_level=5, blas=True, run=1).name: "dashed",
	generate_atc_compressor(optimization_level=5, blas=False, run=1).name: "dotted",
	generate_atc_compressor(optimization_level=5, blas=True, run=2).name: "solid",
	generate_atc_compressor(optimization_level=5, blas=False, run=2).name: "dashed",
	generate_atc_compressor(hwm="csn").name: "dashed",
	TTHRESH.name: "solid",
	x265.name: "solid",
	SZ.name: "solid",
	ZFP.name: "solid",
	FPZIP.name: "solid",
	TuckerMPI.name: "solid",
	generate_tuckermpi_compressor(single_threaded=True).name: "solid",
	generate_tuckermpi_compressor(shared_memory_flags=False).name: "dashed"
}
monochrome_linestyles = {
	ATC.name: "solid",
	TTHRESH.name: "dotted",
	x265.name: (0, (3, 1, 1, 1)),
	SZ.name: "dashed",
	ZFP.name: (0, (3, 1, 1, 1, 1, 1)),
	FPZIP.name: "solid",
	TuckerMPI.name: (0, (7, 3))
} # Only used when we can't use markers, e.g. when using errorbars
markers = {
	ATC.name: "o",
	TTHRESH.name: "v",
	x265.name: "^",
	SZ.name: "s",
	ZFP.name: "p",
	FPZIP.name: "*",
	TuckerMPI.name: "D",
	generate_tuckermpi_compressor(single_threaded=True).name: "D",
}
for level in range(ATC_MAX_OPTIMIZATION_LEVEL):
	linestyles[generate_atc_compressor(optimization_level=level).name] = "solid"
for run in range(1, 10):
	linestyles[generate_atc_compressor(blas=True, run=run).name] = "solid"
	linestyles[generate_atc_compressor(blas=False, run=run).name] = "dashed"
	linestyles[generate_atc_compressor(blas=True, epcm=False, run=run).name] = "dotted"
	linestyles[generate_atc_compressor(blas=False, epcm=False, run=run).name] = "dashdot"
for i, lpmt in enumerate([1, 2, 4, 8, 16, 36]):
	for qrem in [None, 0.001]:
		compressor = generate_atc_compressor(blas=True, epcm=False, qrem=qrem, lpmt=lpmt, run=9)
		colors[compressor.name] = color_wheel[i + 1]
		linestyles[compressor.name] = "solid"

# Preliminaries

IMAGE = 0
LATEX = 1
POSTER = 2
PRESENTATION = 3
def set_output_mode(output_mode):
	if output_mode == IMAGE:
		matplotlib.use("TkAgg") # Default value I got from checking in a Python interpreter
		matplotlib.rcParams.update(matplotlib.rcParamsDefault)
	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': very_small_marker_size,
		})
	elif output_mode == POSTER:
		# Set up Latex poster figure integration
		matplotlib.use("pgf")
		matplotlib.rcParams.update({
			"pgf.texsystem": "pdflatex",
			'font.family': 'sans-serif',
			'text.usetex': True,
			'pgf.rcfonts': False,
			'axes.labelsize': small_font_size,
			'axes.titlesize': small_font_size,
			'xtick.labelsize': small_font_size,
			'ytick.labelsize': small_font_size,
			'legend.fontsize': small_font_size,
			'lines.linewidth': thin_line_width,
			'lines.markersize': very_small_marker_size,
		})
	elif output_mode == PRESENTATION:
		# Set up Beamer presentation figure integration
		matplotlib.use("pgf")
		matplotlib.rcParams.update({
			"pgf.texsystem": "pdflatex",
			'font.family': 'sans-serif',
			'text.usetex': True,
			'pgf.rcfonts': False,
			'axes.labelsize': small_font_size,
			'axes.titlesize': small_font_size,
			'xtick.labelsize': small_font_size,
			'ytick.labelsize': small_font_size,
			'legend.fontsize': small_font_size,
			'lines.linewidth': thin_line_width,
			'lines.markersize': very_small_marker_size,
			'legend.fontsize': small_font_size
		})

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

# 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)

def safe_json_file_update(path, data):
	temp_path = path + ".tmp"
	with open(temp_path, "w") as f:
		json.dump(data, f)
	os.rename(temp_path, path)

def add_relative_errors(dataset):
	original = load_dataset(dataset)
	data_range = compute_data_range(original)
	data_norm = np.linalg.norm(original)
	for compressor in compressors:
		path = results_path + "%s_%s.json"%(dataset.name, compressor.name)
		if os.path.isfile(path):
			with open(path, "r") as f:
				data = json.load(f)
			for key in data[dataset.name][compressor.name]:
				data[dataset.name][compressor.name][key]["relative_error"] = psnr_to_relative_error(data[dataset.name][compressor.name][key]["psnr"], data_range, original.size, data_norm)
			safe_json_file_update(path, data)

def select_keys(dictionary, keys):
	return {key: dictionary[key] for key in keys}

# End of preliminaries

def generate_rate_distortion_plot(dataset, output_mode=IMAGE, base_path=figure_path):
	
	set_output_mode(output_mode)
	
	# Collect data
	data = {dataset.name: {}}
	for compressor in compressors:
		path = results_path + "%s_%s.json"%(dataset.name, compressor.name)
		if os.path.isfile(path):
			with open(path, "r") as f:
				compressor_data = json.load(f)
			data[dataset.name].update(compressor_data[dataset.name])
	
	# Plot a line for each compressor
	rel_error_lower_bound = 1e-6
	rel_error_upper_bound = 1e-1
	max_compression_factor = 1
	for compressor in compressors:
		if compressor.name in data[dataset.name] and len(data[dataset.name][compressor.name]) > 0:
			data_pairs = set()
			for entry in data[dataset.name][compressor.name].values():
				data_pairs.add((entry["relative_error"], entry["compression"]["compression_factor"]))
				if entry["relative_error"] >= rel_error_lower_bound and entry["relative_error"] < rel_error_upper_bound:
					max_compression_factor = max(max_compression_factor, entry["compression"]["compression_factor"])
			data_pairs = sorted(list(data_pairs))
			rel_errors, compression_factors = zip(*data_pairs)
		else:
			rel_errors = []
			compression_factors = []
		plt.plot(rel_errors, compression_factors, label=compressor.name, marker=".")
	
	# Finalize plot
	plt.xlabel("Relative error")
	plt.xscale("log")
	plt.xlim(rel_error_lower_bound, rel_error_upper_bound)
	plt.ylabel("Compression factor")
	plt.yscale("log")
	plt.ylim(1, max_compression_factor)
	plt.title(dataset.name + " (" + (dataset.original_dtype if dataset.original_dtype is not None else dataset.dtype) + ", " + " x ".join(list(map(str, dataset.sizes))) + ")")
	if output_mode == IMAGE:
		plt.legend()
		plt.show()
		plt.close()
	elif output_mode == LATEX:
		plt.legend(fontsize=very_small_font_size)
		set_plot_width(plt.gcf(), 0.3)
		plt.savefig(base_path + "rate-distortion_" + dataset.name.replace(" ", "-") + ".pgf", bbox_inches="tight")
		plt.close()

def generate_rate_distortions_plots(datasets, selected_compressors=compressors, output_mode=IMAGE, output_path=None, close_comparison=False):
	
	set_output_mode(output_mode)
	
	rel_error_lower_bound = 1e-5
	rel_error_upper_bound = 1e-1
	target_rel_errors = np.logspace(log10(rel_error_lower_bound), log10(rel_error_upper_bound), 1001)
	rows = 1 if close_comparison else 2
	columns = len(datasets)
	wide_view = (output_mode != POSTER)
	if not wide_view:
		columns, rows = rows, columns
	if output_mode == IMAGE:
		fig, axes = plt.subplots(rows, columns, squeeze=False)
	elif output_mode == POSTER or output_mode == PRESENTATION:
		fig, axes = plt.subplots(rows, columns, squeeze=False, figsize=(1.1*columns*small_plot_size/dpi, rows*small_plot_size/dpi))
	for dataset_index, dataset in enumerate(datasets):
		
		dataset_string = " x ".join(list(map(str, dataset.sizes)))
		
		# Collect data
		data = {dataset.name: {}}
		for compressor in selected_compressors:
			path = results_path + ("%s_%s.json"%(dataset.name, compressor.name)).replace(" ", "_")
			if os.path.isfile(path):
				with open(path, "r") as f:
					compressor_data = json.load(f)
				data[dataset.name].update(compressor_data[dataset.name])
		
		compression_factors_data = {}
		max_compression_factor = 1
		row, column = 0, dataset_index
		if not wide_view:
			column, row = row, column
		if not close_comparison:
			absolute_ax = axes[row, column]
		for compressor_name in data[dataset.name]:
				
				# Filter and sort unique data points from collected data
				data_pairs = set()
				for entry in data[dataset.name][compressor_name].values():
					data_pairs.add((entry["relative_error"], entry["compression"]["compression_factor"]))
				data_pairs = sorted(list(data_pairs))
				
				if len(data_pairs) > 0:
					
					rel_errors, compression_factors = zip(*data_pairs)
					
					# Plot absolute compression factors
					for i in range(len(rel_errors)):
						if rel_errors[i] >= rel_error_lower_bound and rel_errors[i] < rel_error_upper_bound:
							max_compression_factor = max(max_compression_factor, compression_factors[i])
					if rel_errors[0] == 0:
						# Don't plot first data point (at infinity on logarithmic scale)
						# Instead, draw line to indicate minimal achievable compression (at lossless compression)
						if not close_comparison:
							absolute_ax.loglog([rel_error_lower_bound, rel_errors[1]], [compression_factors[0],]*2, color=colors[compressor_name], linestyle=linestyles[compressor_name])
						rel_errors_to_plot = rel_errors[1:]
						compression_factors_to_plot = compression_factors[1:]
					else:
						rel_errors_to_plot = rel_errors
						compression_factors_to_plot = compression_factors
					if not close_comparison:
						absolute_ax.loglog(rel_errors_to_plot, compression_factors_to_plot, label=compressor_name, marker=".", color=colors[compressor_name], linestyle=linestyles[compressor_name])
				
					# Interpolate data points using linear interpolation on the logarithms of the compression factors and logarithms of relative errors
					compression_factors_data[compressor_name] = {"rel_errors": [], "interpolated_compression_factors": [], "normalized_compression_factors": []}
					for target_rel_error in target_rel_errors:
						
						# Find line segment for interpolation (or in the extreme cases, extrapolation)
						if target_rel_error < rel_errors[0]:
							# No data available for this point since we don't extrapolate, skip
							continue
						elif target_rel_error >= rel_errors[-1]:
							# No data available for this point since we don't extrapolate, skip
							continue
						else:
							for i in range(len(rel_errors) - 1):
								# Check if this line segment contains the target PSNR
								if target_rel_error >= rel_errors[i] and target_rel_error < rel_errors[i + 1]:
									break
						
						if rel_errors[i] == 0:
							# In this case, the slope of the line is 0 on a logarithmic scale but it should be drawn at the compression factor achieved at 0 error
							interpolated_compression_factor = compression_factors[i]
						else:
							# No weird behavior with logarithms here
							x1, x2 = map(log10, 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)
						compression_factors_data[compressor_name]["rel_errors"].append(target_rel_error)
						compression_factors_data[compressor_name]["interpolated_compression_factors"].append(interpolated_compression_factor)
		
		# Finish up absolute compression factors subplot
		if not close_comparison:
			absolute_ax.set_xlim(rel_error_lower_bound, rel_error_upper_bound)
			absolute_ax.set_ylim(1, 1.5*max_compression_factor) # Add a bit of a margin on top
			quantity = "Compression factor"
			if wide_view:
				if column == 0:
					absolute_ax.set_ylabel(quantity)
			else:
				if row == 0:
					absolute_ax.set_title(quantity)
			if wide_view:
				absolute_ax.set_title(dataset_string)
			else:
				absolute_ax.set_ylabel(dataset_string)
		
		# Generate normalized compression factors
		indices = {compressor_name: 0 for compressor_name in compression_factors_data}
		for target_rel_error in target_rel_errors:
			
			# Find corresponding compression factor in each compressor dataset and calculate maximum
			max_compression_factor = 0
			for compressor_name in compression_factors_data:
				while indices[compressor_name] < len(compression_factors_data[compressor_name]["rel_errors"]) and compression_factors_data[compressor_name]["rel_errors"][indices[compressor_name]] < target_rel_error:
					indices[compressor_name] += 1
				if indices[compressor_name] < len(compression_factors_data[compressor_name]["rel_errors"]) and compression_factors_data[compressor_name]["rel_errors"][indices[compressor_name]] == target_rel_error:
					max_compression_factor = max(max_compression_factor, compression_factors_data[compressor_name]["interpolated_compression_factors"][indices[compressor_name]])
			
			# Index already matches for each compressor, so just compute normalized compression factors at the current indices
			for compressor_name in compression_factors_data:
				if indices[compressor_name] < len(compression_factors_data[compressor_name]["rel_errors"]) and compression_factors_data[compressor_name]["rel_errors"][indices[compressor_name]] == target_rel_error:
					compression_factors_data[compressor_name]["normalized_compression_factors"].append(compression_factors_data[compressor_name]["interpolated_compression_factors"][indices[compressor_name]]/max_compression_factor)
		
		# Plot normalized compression factors
		row, column = (0 if close_comparison else 1), dataset_index
		if not wide_view:
			column, row = row, column
		ax = axes[row, column]
		for compressor in selected_compressors:
			if compressor.name in compression_factors_data:
				ax.semilogx(compression_factors_data[compressor.name]["rel_errors"], compression_factors_data[compressor.name]["normalized_compression_factors"], label=compressor.name, color=colors[compressor.name], linestyle=linestyles[compressor.name])
		
		# Finalize normalized compression factor plot
		ax.set_xlim(rel_error_lower_bound, rel_error_upper_bound)
		zoom = 0.4 if close_comparison else 1
		ax.set_ylim(1 - zoom, 1 + 0.1*zoom)
		if row == rows - 1:
			ax.set_xlabel("Relative error")
		quantity = "Relative compression factor"
		if wide_view:
			if row == 0:
				ax.set_title(dataset_string)
			if column == 0:
				ax.set_ylabel(quantity)
		else:
			if row == 0:
				ax.set_title(quantity)
	
	# Finalize full plot
	if output_mode == IMAGE:
		plt.legend()
		plt.show()
		plt.close()
	elif output_mode == POSTER:
		plt.subplots_adjust(wspace=0.3, hspace=0.3)
		set_plot_width(fig, 0.45)
		fig.savefig(output_path, bbox_inches="tight", transparent=True)
		plt.close()
	elif output_mode == PRESENTATION:
		plt.subplots_adjust(wspace=0.3, hspace=0.3)
		set_plot_width(fig, 0.7)
		
		# Add legend correctly
		handles, labels = axes[0, 0].get_legend_handles_labels()
		if "ATC (0.1)" in labels:
			labels = ["ATC (0.5)" if label == "ATC" else label for label in labels]
		loc = "lower center"
		fig.legend(handles, labels, loc="lower center", bbox_to_anchor=(0.5, -0.3 if rows == 1 else -0.2), bbox_transform=plt.gcf().transFigure, ncol=len(selected_compressors)//2)
		
		fig.savefig(output_path, bbox_inches="tight", transparent=True)
		plt.close()

def generate_times_plot(datasets, selected_compressors=compressors, output_mode=IMAGE, output_path=None):
	
	set_output_mode(output_mode)
	
	rel_error_lower_bound = 1e-5
	rel_error_upper_bound = 1e-1
	rows = 2
	columns = len(datasets)
	if output_mode == IMAGE:
		fig, axes = plt.subplots(rows, columns, squeeze=False)
	elif output_mode == PRESENTATION:
		fig, axes = plt.subplots(rows, columns, squeeze=False, figsize=(1.2*columns*small_plot_size/dpi, rows*small_plot_size/dpi))
	for column, dataset in enumerate(datasets):
		
		dataset_string = " x ".join(list(map(str, dataset.sizes)))
		
		# Collect data
		data = {dataset.name: {}}
		for compressor in selected_compressors:
			if dataset.name == "Isotropic1024-fine-pressure-time-128" and compressor.name == "x265":
				# x265 compression broke here for some reason, ignore data for now
				continue
			path = results_path + ("ATC-wrong-mode-storage-order/" if "ATC" in compressor.name else "") + ("%s_%s.json"%(dataset.name, compressor.name)).replace(" ", "_")
			if os.path.isfile(path):
				with open(path, "r") as f:
					compressor_data = json.load(f)
				data[dataset.name].update(compressor_data[dataset.name])
		
		for row, ylabel, key in ((0, "Compression time (s)", "compression"), (1, "Decompression time (s)", "decompression")):
			
			ax = axes[row, column]
			
			for compressor_name in data[dataset.name]:
				
				data_pairs = set()
				for entry in data[dataset.name][compressor_name].values():
					data_pairs.add((entry["relative_error"], entry[key]["real_time"]))
				data_pairs = sorted(list(data_pairs))
				
				if len(data_pairs) > 0:
					
					rel_errors, times = zip(*data_pairs)
					
					# Plot times
					if rel_errors[0] == 0:
						ax.loglog([rel_error_lower_bound, rel_errors[1]], [max(times[0], times[1]),]*2, color=colors[compressor_name], linestyle=linestyles[compressor_name])
						xs = rel_errors[1:]
						ys = times[1:]
					else:
						xs = rel_errors
						ys = times
					ax.loglog(xs, ys, label=compressor_name, marker=".", color=colors[compressor_name], linestyle=linestyles[compressor_name])
			
			# Finalize subplot
			ax.set_xlim(rel_error_lower_bound, rel_error_upper_bound)
			ax.tick_params(axis='y', which='major', pad=-1)
			if row == 0:
				ax.set_title(dataset_string)
			if row == rows - 1:
				ax.set_xlabel("Relative error")
			if column == 0:
				ax.set_ylabel(ylabel)
	
	# Finalize full plot
	if output_mode == IMAGE:
		plt.legend()
		plt.show()
		plt.close()
	elif output_mode == PRESENTATION:
		plt.subplots_adjust(wspace=0.3, hspace=0.3)
		set_plot_width(fig, 0.75)
		
		# Add legend correctly
		handles, labels = axes[0, 0].get_legend_handles_labels()
		if "ATC (0.1)" in labels:
			labels = ["ATC (0.5)" if label == "ATC" else label for label in labels]
		fig.legend(handles, labels, loc="lower center", bbox_to_anchor=(0.5, -0.3 if rows == 1 else -0.2), bbox_transform=plt.gcf().transFigure, ncol=(len(selected_compressors) + 1)//2)
		
		fig.savefig(output_path, bbox_inches="tight", transparent=True)
		plt.close()

def generate_memory_usage_error_control_plot(datasets, selected_compressors=compressors, output_mode=IMAGE, output_path=None):
	
	set_output_mode(output_mode)
	
	rel_error_lower_bound = 1e-5
	rel_error_upper_bound = 1e-1
	rows = 2
	columns = len(datasets)
	if output_mode == IMAGE:
		fig, axes = plt.subplots(rows, columns, squeeze=False)
	elif output_mode == PRESENTATION:
		fig, axes = plt.subplots(rows, columns, squeeze=False, figsize=(1.2*columns*small_plot_size/dpi, rows*small_plot_size/dpi))
	for column, dataset in enumerate(datasets):
		
		dataset_string = " x ".join(list(map(str, dataset.sizes)))
		dataset_size = reduce(lambda x, y: x*y, dataset.sizes)
		
		# Collect data
		data = {dataset.name: {}}
		for compressor in selected_compressors:
			if dataset.name == "Isotropic1024-fine-pressure-time-128" and compressor.name == "x265":
				# x265 compression broke here for some reason, ignore data for now
				continue
			path = results_path + ("ATC-wrong-mode-storage-order/" if "ATC" in compressor.name else "") + ("%s_%s.json"%(dataset.name, compressor.name)).replace(" ", "_")
			if os.path.isfile(path):
				with open(path, "r") as f:
					compressor_data = json.load(f)
				data[dataset.name].update(compressor_data[dataset.name])
		
		for row, ylabel, key in ((0, "$\\frac{\\mbox{Peak memory (bytes)}}{\\mbox{\# tensor elements}}$", "peak_memory"), (1, "$\\frac{\\mbox{Relative error}}{\\mbox{Target relative error}}$", "error_control")):
			
			ax = axes[row, column]
			
			for compressor_name in data[dataset.name]:
				
				if key == "error_control" and compressor_name not in ("ATC", "TTHRESH", "SZ"):
					continue
				
				data_pairs = set()
				for target_psnr, entry in data[dataset.name][compressor_name].items():
					if key == "peak_memory":
						quantity = max(entry["compression"]["peak_memory"], entry["decompression"]["peak_memory"])/dataset_size
					elif key == "error_control":
						target_psnr = float(target_psnr)
						if compressor_name == "TTHRESH":
							# TTHRESH uses different definition of PSNR, so adjust target PSNR so it equals what TTHRESH tries to approximate internally
							# Specifically, it tries to approximate an SSE which is 4 times lower, so its target relative error is 2 times lower, so its target PSNR is 20*log10(2) dB higher
							target_psnr += 20*log10(2)
						quantity = 10**((target_psnr - entry["psnr"])/20)
					data_pairs.add((entry["relative_error"], quantity))
				data_pairs = sorted(list(data_pairs))
				
				if len(data_pairs) > 0:
					
					rel_errors, quantities = zip(*data_pairs)
					
					# Plot times
					if rel_errors[0] == 0:
						ax.plot([rel_error_lower_bound, rel_errors[1]], [max(abs(quantities[0]), abs(quantities[1])),]*2, color=colors[compressor_name], linestyle=linestyles[compressor_name])
						xs = rel_errors[1:]
						ys = quantities[1:]
					else:
						xs = rel_errors
						ys = quantities
					ax.plot(xs, ys, label=compressor_name, marker=".", color=colors[compressor_name], linestyle=linestyles[compressor_name])
			
			# Finalize subplot
			ax.set_xlim(rel_error_lower_bound, rel_error_upper_bound)
			ax.set_xscale("log")
			if key == "error_control":
				ax.set_yscale("log")
				ax.set_ylim(5e-1, 2e0)
				ax.tick_params(axis='y', which='major', pad=-1)
			if row == 0:
				ax.set_title(dataset_string)
			if row == rows - 1:
				ax.set_xlabel("Relative error")
			if column == 0:
				ax.set_ylabel(ylabel.replace("\\", "") if output_mode == IMAGE else ylabel)
	
	# Finalize full plot
	if output_mode == IMAGE:
		plt.legend()
		plt.show()
		plt.close()
	elif output_mode == PRESENTATION:
		plt.subplots_adjust(wspace=0.3, hspace=0.3)
		set_plot_width(fig, 0.75)
		
		# Add legend correctly
		handles, labels = axes[0, columns - 1].get_legend_handles_labels()
		if "ATC (0.1)" in labels:
			labels = ["ATC (0.5)" if label == "ATC" else label for label in labels]
		fig.legend(handles, labels, loc="lower center", bbox_to_anchor=(0.5, -0.3 if rows == 1 else -0.2), bbox_transform=plt.gcf().transFigure, ncol=(len(selected_compressors) + 1)//2)
		
		fig.savefig(output_path, bbox_inches="tight", transparent=True)
		plt.close()

# Below here we use the new format with multiple samples per target error
# Targets are mid-term report and paper
	
"""
PSNR bounds determined using the following data:
Relative error bounds: (1e-05, 0.1)
PSNR bounds: [37.92750183642599, 115.0181797698151]
Relative error bounds: (1e-05, 0.1)
PSNR bounds: [46.53818767826563, 125.07186885392485]
Relative error bounds: (1e-05, 0.1)
PSNR bounds: [44.436349142191965, 120.78772498825037]
Relative error bounds: (1e-05, 0.1)
PSNR bounds: [39.73435114698919, 117.15346543547756]
Relative error bounds: (1e-05, 0.1)
PSNR bounds: [38.07925370554744, 115.65456886605364]
Relative error bounds: (1e-05, 0.1)
PSNR bounds: [63.609286649753955, 140.26054473706944]
Relative error bounds: (1e-05, 0.1)
PSNR bounds: [60.01588376999797, 137.58333667956873]
Relative error bounds: (1e-05, 0.1)
PSNR bounds: [41.776562634732635, 120.81026690706065]
"""

use_psnr = False
rel_error_filter_bound = 0.5 # Any data points with rel. error above this are discarded
rel_error_lower_bound = 1e-4
rel_error_upper_bound = 1e-1
psnr_lower_bound = 40
psnr_upper_bound = 120
target_interpolation_errors_amount = 1001
target_interpolation_errors = np.linspace(psnr_lower_bound, psnr_upper_bound, target_interpolation_errors_amount) if use_psnr else np.logspace(log10(rel_error_upper_bound), log10(rel_error_lower_bound), target_interpolation_errors_amount)

def has_error_control(compressor):
	return compressor.name in ("ATC", "TTHRESH", "SZ", "TuckerMPI") or compressor.name.startswith("ATC") or compressor.name.startswith("TuckerMPI")

def generate_plot_grid(datasets, selected_compressors=compressors, metrics=["compression factor", "relative compression factor", "compression time", "decompression time"], output_mode=IMAGE, output_path=None, target="report", venue=None):
	
	monochrome = False
	if target == "paper" and output_mode == LATEX:
		if venue is None:
			generate_plot_grid(datasets, selected_compressors=selected_compressors, metrics=metrics, output_mode=output_mode, output_path=output_path, target=target, venue="acm-toms")
			generate_plot_grid(datasets, selected_compressors=selected_compressors, metrics=metrics, output_mode=output_mode, output_path=output_path, target=target, venue="arxiv")
			return
		elif venue != "shared":
			output_path_parts = output_path.split("/")
			output_path = "/".join(output_path_parts[:-1] + [venue, output_path_parts[-1]])
		if venue == "acm-toms":
			monochrome = True
	
	set_output_mode(output_mode)
	
	max_columns = 4
	wide_view = (len(datasets) == 1 and len(metrics) <= max_columns)
	if wide_view:
		wrap_factor = 1
		rows = 1
		columns = len(metrics)
	else:
		wrap_factor = ceil(len(datasets)/max_columns)
		rows = len(metrics)*wrap_factor
		columns = min(max_columns, len(datasets))
	if output_mode == IMAGE:
		fig = plt.figure()
		#fig, axes = plt.subplots(rows, columns, squeeze=False)
	else:
		fig = plt.figure(figsize=(((1.2 if wide_view else 0.99) if target == "paper" else 1.1)*columns*small_plot_size/dpi, (0.7 if len(metrics) > 1 else 0.8)*rows*small_plot_size/dpi))
		#fig, axes = plt.subplots(rows, columns, squeeze=False, figsize=(((1.2 if wide_view else 1) if target == "paper" else 1.1)*columns*small_plot_size/dpi, rows*small_plot_size/dpi))
	gridspecs = []
	gap_height = 0.1 if len(metrics) > 1 else 0.2 # Total height of all gridspec gaps combined
	gap_height_per_gridspec = gap_height/(wrap_factor - 1) if wrap_factor > 1 else 0
	gridspec_height = (1 - gap_height)/wrap_factor
	for i in range(wrap_factor):
		gs = GridSpec(len(metrics), columns)
		bottom = i*(gridspec_height + gap_height_per_gridspec)
		gs.update(bottom=bottom, top=bottom + gridspec_height, wspace=0.3, hspace=0.1)
		gridspecs.append(gs)
	gridspecs = gridspecs[::-1]
	
	if not use_psnr:
		# This function was written in the old format, so instead of changing the code here we simply reverse the errors such that the error values are always increasing ...
		local_target_interpolation_errors = target_interpolation_errors[::-1]
	else:
		local_target_interpolation_errors = target_interpolation_errors
	
	for dataset_index, dataset in enumerate(datasets):
		
		if wide_view:
			row = dataset_index
		else:
			column = dataset_index%max_columns
		dataset_string = " x ".join(list(map(str, dataset.sizes))) if target == "report" else dataset.name
		
		# Collect data
		data = {dataset.name: {}}
		for compressor in selected_compressors:
			path = results_path + ("%s_%s.json"%(dataset.name, compressor.name)).replace(" ", "_")
			if os.path.isfile(path):
				with open(path, "r") as f:
					compressor_data = json.load(f)
				data[dataset.name].update(compressor_data[dataset.name])
		
		for metric_index, metric in enumerate(metrics):
			
			if wide_view:
				column = metric_index
			else:
				row = (dataset_index//max_columns)*len(metrics) + metric_index
			#ax = axes[row, column]
			ax = plt.subplot(gridspecs[dataset_index//max_columns][metric_index, dataset_index%max_columns], axes_class=axisartist.Axes if metric.endswith("time") else None)
			label_y_axis = (wide_view or column == 0)
			
			if metric == "compression factor" or (metric == "relative compression factor" and "compression factor" not in metrics[:metric_index]):
				
				compression_factors_data = {}
				max_compression_factor = 1
				for compressor_name in data[dataset.name]:
					
					# Filter and sort unique data points from collected data
					data_pairs = set()
					for target_error, samples in data[dataset.name][compressor_name].items():
						for sample in samples:
							if sample["relative_error"] < rel_error_filter_bound:
								if use_psnr:
									data_pairs.add((sample["psnr"], sample["compression"]["compression_factor"]))
									if sample["psnr"] >= psnr_lower_bound and sample["psnr"] < psnr_upper_bound:
										max_compression_factor = max(max_compression_factor, sample["compression"]["compression_factor"])
								else:
									data_pairs.add((sample["relative_error"], sample["compression"]["compression_factor"]))
									if sample["relative_error"] >= rel_error_lower_bound and sample["relative_error"] < rel_error_upper_bound:
										max_compression_factor = max(max_compression_factor, sample["compression"]["compression_factor"])
					data_pairs = sorted(list(data_pairs))
					
					if len(data_pairs) > 0:
						
						errors, compression_factors = zip(*data_pairs)
						
						if metric == "compression factor":
							# Plot absolute compression factors
							if use_psnr and errors[-1] == float("inf"):
								ax.plot([errors[-2], psnr_upper_bound], [compression_factors[-1],]*2, color=colors[compressor_name], linestyle=linestyles[compressor_name])
								errors_to_plot = errors[:-1]
								compression_factors_to_plot = compression_factors[:-1]
							elif not use_psnr and errors[0] == 0:
								# Don't plot first data point (at infinity on logarithmic scale)
								# Instead, draw line to indicate minimal achievable compression (at lossless compression)
								ax.plot([rel_error_lower_bound, errors[1]], [compression_factors[0],]*2, color=colors[compressor_name], linestyle=linestyles[compressor_name])
								errors_to_plot = errors[1:]
								compression_factors_to_plot = compression_factors[1:]
							else:
								errors_to_plot = errors
								compression_factors_to_plot = compression_factors
							ax.plot(errors_to_plot, compression_factors_to_plot, label=compressor_name, marker=markers[compressor_name] if monochrome else "o", color=colors[compressor_name], linestyle=linestyles[compressor_name])
						
						if "relative compression factor" in metrics:
							
							# Interpolate data points using linear interpolation on the logarithms of the compression factors and logarithms of relative errors
							if use_psnr and errors[-1] == float("inf"):
								sample_errors = errors[:-1]
							elif not use_psnr and errors[0] == 0:
								sample_errors = errors[1:]
							else:
								sample_errors = errors
							compression_factors_data[compressor_name] = {"errors": [], "valid_errors": [], "sample_errors": sample_errors, "valid_sample_errors": [], "interpolated_compression_factors": [], "relative_compression_factors": [], "sample_relative_compression_factors": []}
							for target_error in local_target_interpolation_errors:
								
								# Find line segment for interpolation (or in the extreme cases, extrapolation)
								if target_error < errors[0]:
									# No data available for this point since we don't extrapolate, skip
									continue
								elif target_error >= errors[-1]:
									# No data available for this point since we don't extrapolate, skip
									continue
								
								for i in range(len(errors) - 1):
									# Check if this line segment contains the target PSNR
									if target_error >= errors[i] and target_error < errors[i + 1]:
										break
								
								if (errors[i + 1] == float("inf")) if use_psnr else (errors[i] == 0):
									# In this case, the slope of the line is 0 on a logarithmic scale but it should be drawn at the minimum compression factor (normally the one achieved at 0 error)
									interpolated_compression_factor = min(compression_factors[i:i + 2])
								else:
									# No weird behavior with logarithms here
									error_mapping_function = (lambda psnr: psnr) if use_psnr else log10
									x1, x2 = map(error_mapping_function, 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*error_mapping_function(target_error) + intercept)
								compression_factors_data[compressor_name]["errors"].append(target_error)
								compression_factors_data[compressor_name]["interpolated_compression_factors"].append(interpolated_compression_factor)
				
				if metric == "compression factor":
					# Finish up absolute compression factors subplot
					ax.set_ylim(1, 1.5*max_compression_factor) # Add a bit of a margin on top
					ax.set_yscale("log")
					if label_y_axis:
						ax.set_ylabel("Compressiefactor" if target == "report" else "Compression factor")
			
			if metric == "relative compression factor":
				
				# Generate relative compression factors
				
				indices = {compressor_name: 0 for compressor_name in compression_factors_data.keys()}
				marker_error_indices = {compressor_name: 0 for compressor_name in compression_factors_data.keys()}
				compressor_names = compression_factors_data.keys()
				reference_compressor_name = "ATC" if "ATC" in compressor_names else ("TTHRESH" if "TTHRESH" in compressor_names else "ATC (RTMSS = 0.4)")
				interpolation_error_resolution = (local_target_interpolation_errors[1] - local_target_interpolation_errors[0])/local_target_interpolation_errors[0]
				for target_error in local_target_interpolation_errors:
					
					# Find corresponding compression factor in each compressor dataset
					place_marker_here = False
					max_compression_factor = 0
					for compressor_name in compression_factors_data.keys():
						while indices[compressor_name] < len(compression_factors_data[compressor_name]["errors"]) and compression_factors_data[compressor_name]["errors"][indices[compressor_name]] < target_error:
							indices[compressor_name] += 1
					
					# Only compute relative compression factor if a proper reference is present, else skip
					if indices[reference_compressor_name] < len(compression_factors_data[reference_compressor_name]["errors"]) and compression_factors_data[reference_compressor_name]["errors"][indices[reference_compressor_name]] == target_error:
						# Index already matches for each compressor, so just compute relative compression factors at the current indices
						for compressor_name in compression_factors_data.keys():
							if indices[compressor_name] < len(compression_factors_data[compressor_name]["errors"]) and compression_factors_data[compressor_name]["errors"][indices[compressor_name]] == target_error:
								compression_factors_data[compressor_name]["valid_errors"].append(target_error)
								interpolated_compression_factor = compression_factors_data[compressor_name]["interpolated_compression_factors"][indices[compressor_name]]
								compression_factors_data[compressor_name]["relative_compression_factors"].append(interpolated_compression_factor/compression_factors_data[reference_compressor_name]["interpolated_compression_factors"][indices[reference_compressor_name]])
								if marker_error_indices[compressor_name] < len(compression_factors_data[compressor_name]["sample_errors"]):
									marker_error = compression_factors_data[compressor_name]["sample_errors"][marker_error_indices[compressor_name]]
									if target_error > marker_error:
										if (target_error - marker_error)/marker_error < 2*interpolation_error_resolution:
											# This error is close enough to the marker error we were looking for, so use it
											compression_factors_data[compressor_name]["valid_sample_errors"].append(target_error)
											compression_factors_data[compressor_name]["sample_relative_compression_factors"].append(compression_factors_data[compressor_name]["relative_compression_factors"][-1])
										marker_error_indices[compressor_name] += 1
								
				
				# Plot relative compression factors
				for compressor in selected_compressors:
					if compressor.name in compressor_names:
						ax.plot(compression_factors_data[compressor.name]["valid_errors"], compression_factors_data[compressor.name]["relative_compression_factors"], color=colors[compressor.name], linestyle=linestyles[compressor.name], label=None if monochrome else compressor.name)
						if monochrome:
							ax.scatter(compression_factors_data[compressor.name]["valid_sample_errors"], compression_factors_data[compressor.name]["sample_relative_compression_factors"], marker=markers[compressor.name], color=colors[compressor.name])
							ax.plot([], [], marker=markers[compressor.name], linestyle=linestyles[compressor.name], color=colors[compressor.name], label=compressor.name)
				
				# Finalize relative compression factor plot
				if "SZ" in data[dataset.name].keys():
					# Full compressor comparison, wide range of results so use log plot
					ax.set_ylim(0.01, 2)
					ax.set_yscale("log")
					ticks = [0.01, 0.02, 0.03, 0.04, 0.05, 0.06, 0.07, 0.08, 0.09, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1, 2]
					ax.set_yticks(ticks)
					if output_mode == LATEX:
						ax.set_yticklabels([str(ticks[i]) if i%9 == 0 else "" for i in range(len(ticks))])
				else:
					# TTHRESH comparison, similar values so linear plot shows more detail
					ax.set_ylim(0.5, 1.5)
				if label_y_axis:
					ax.set_ylabel("Relatieve compressiefactor" if target == "report" else "Relative compression factor")
			
			if metric.endswith("time") or metric.endswith("memory") or metric == "error control":
				
				use_error_bars = metric.endswith("time") and "ATC (RTMSS = 0.0)" not in data[dataset.name].keys() and "ATC (DTMSS = 0.1)" not in data[dataset.name].keys()
				
				# Process data
				max_error_margin = 0 # Maximum relative deviation from average relative error within same set of samples
				for compressor_name in data[dataset.name]:
					
					if metric == "error control" and not has_error_control(compressor):
						continue
					
					# Helpers
					def aggregate_fun(agg_fun, metric, samples, target_error=None):
						if metric == "error":
							sample_fun = (lambda sample: sample["psnr"]) if use_psnr else (lambda sample: sample["relative_error"])
						elif metric == "error control":
							if compressor_name.startswith("TuckerMPI"):
								sample_to_psnr_deviation = lambda sample: 20*log10(sample["relative_error"]/float(target_error))
							else:
								sample_to_psnr_deviation = lambda sample: (float(target_error) + (20*log10(2) if compressor_name == "TTHRESH" else 0) - sample["psnr"]) # Correction due to TTHRESH's different definition of PSNR
							if use_psnr:
								sample_fun = sample_to_psnr_deviation
							else:
								sample_fun = lambda sample: 10**(sample_to_psnr_deviation(sample)/20)
						else:
							sample_fun = lambda sample: sample[metric.split(" ")[0]]["real_time" if metric.endswith("time") else "peak_memory"]
						return agg_fun(list(map(sample_fun, samples)))/(np.prod(dataset.sizes) if metric.endswith("memory") else 1)
					arithmetic_avg = lambda vals: sum(vals)/len(vals)
					geometric_avg = lambda vals: 10**(sum(map(log10, vals))/len(vals))
					
					data_tuples = set() # Contains tuples (avg_rel_error, avg_real_time, min_real_time, max_real_time)
					for target_error, samples in data[dataset.name][compressor_name].items():
						if metric.endswith("time") and len(samples) < 5:
							# We only use data points with at least 5 samples
							continue
						if len(samples) == 0:
							print("Warning: 0 samples for dataset %s for compressor %s and target error %s"%(dataset.name, compressor_name, target_error))
							continue
						min_error = aggregate_fun(min, "error", samples)
						avg_error = aggregate_fun(arithmetic_avg if use_psnr else geometric_avg, "error", samples)
						max_error = aggregate_fun(max, "error", samples)
						if not use_psnr or [min_error, avg_error, max_error] != [float("inf"),]*3:
							error_margin = max(abs(min_error - avg_error), abs(max_error - avg_error))/(avg_error if (not use_psnr and avg_error != 0) else 1)
							if isnan(error_margin):
								print(target_psnr, compressor_name, samples)
								print(aggregate_fun(min, "error", samples), avg_error, aggregate_fun(max, "error", samples), (avg_error if (not use_psnr and avg_error != 0) else 1))
							if error_margin > (20*log10(1.1) if use_psnr else 0.1):
								# If there's a large deviation on the actual error, skip this set of samples since we can't represent it properly on the graph without horizontal error bars
								# Only happens in the case of x265 anyway and there we have plenty of data points already
								continue
							max_error_margin = max(error_margin, max_error_margin)
						scale = 1/np.prod(dataset.sizes) if metric.endswith("memory") else 1
						if samples[0]["relative_error"] < rel_error_filter_bound:
							data_tuples.add((avg_error, aggregate_fun(np.median, metric, samples, target_error=target_error), aggregate_fun(min, metric, samples, target_error=target_error), aggregate_fun(max, metric, samples, target_error=target_error))) # Before: aggregate_fun(arithmetic_avg if metric == "error control" else geometric_avg, ...)
					
					# Process data tuples
					data_tuples = sorted(list(data_tuples))
					if len(data_tuples) == 0:
						print("Warning: obtained no data points for compressor %s, dataset %s and metric \"%s\""%(compressor_name, dataset.name, metric))
						continue
					if use_psnr and data_tuples[-1][0] == float("inf"):
						# Infinite PSNR does not fit on scale, remove
						del data_tuples[-1]
					elif not use_psnr and (data_tuples[0][0] == 0):
						# Relative error 0 does not fit on log scale, remove
						del data_tuples[0]
					if metric.endswith("time") and compressor_name == "x265":
						# There's too much data to visualize, so space out the data points such that there are around 4 left in our range
						data_points_in_range = sum(map(lambda data_tuple: int((data_tuple[0] >= psnr_lower_bound) if use_psnr else (data_tuple[0] < rel_error_upper_bound)), data_tuples))
						data_tuples = data_tuples[::max(1, data_points_in_range//4)] # If there are no data points in range, just keep all data
					
					# Plot data tuples
					avg_errors, avg_vals, min_vals, max_vals = zip(*data_tuples)
					ax.errorbar(avg_errors, avg_vals, yerr=[np.array(avg_vals) - np.array(min_vals), np.array(max_vals) - np.array(avg_vals)] if use_error_bars else None, capsize=1.5, label=compressor_name, color=colors[compressor_name], linestyle=(monochrome_linestyles if monochrome and use_error_bars else linestyles)[compressor_name], marker=markers[compressor_name] if monochrome and not use_error_bars else None)
				
				print("Maximum relative error margin for samples from the same target error for dataset %s and metric %s: %s"%(dataset.name, metric, max_error_margin))
				
				if metric.endswith("time"):
					ax.set_yscale("log")
					ax.axis["left"].major_ticklabels.set_ha("left")
					ax.axis["left"].minor_ticklabels.set_ha("left")
					ax.axis["left"].major_ticklabels.set_pad(0)
					for pos in ["left", "bottom"]:
						ax.axis[pos].major_ticks.set_tick_out(True)
						ax.axis[pos].minor_ticks.set_tick_out(True)
					for pos in ["top", "right"]:
						ax.axis[pos].major_ticks.set(visible=False)
						ax.axis[pos].minor_ticks.set(visible=False)
					#print(matplotlib.rcParams)
					# Optional: fix scale across plots, makes it easier to put error bars into perspective
					"""
					ylims = ax.get_ylim()
					midy = (ylims[0]*ylims[1])**0.5
					half_fixed_scale = 10**0.5
					ax.set_ylim(midy/half_fixed_scale, midy*half_fixed_scale)
					"""
				elif metric.endswith("memory"):
					ax.set_ylim(0, 45)
				elif metric == "error control":
					if use_psnr:
						ax.set_yscale("linear")
						ax.set_ylim(-15, 5)
						ticks = [-15, -10, -5, 0, 5]
						ax.set_yticks(ticks)
						if output_mode == LATEX:
							ax.set_yticklabels(["$%s$"%tick for tick in ticks]) # We set these explicitly because else there's a rendering problem with the minus sign when rendering to PGF
					else:
						ax.set_yscale("log")
						ax.set_ylim(0.5, 1.5)
						ticks = [0.5, 0.6, 0.7, 0.8, 0.9, 1, 1.2, 1.4]
						ax.set_yticks(ticks)
						if output_mode == LATEX:
							ax.set_yticklabels(map(str, ticks))
				if label_y_axis:
					ax.set_ylabel({
						"report": {
							"compression time": "Compressietijd (s)",
							"decompression time": "Decompressietijd (s)",
							"compression memory": "Compressie (bytes/element)$",
							"decompression memory": "Decompressie (bytes/element)",
							"error control": "$\\frac{\\mbox{Relatieve fout}}{\\mbox{Relatieve doelfout}}$" if output_mode == LATEX else "(Relatieve fout)/(Relatieve doelfout)"
						}, "paper": {
							"compression time": "Compression time (s)",
							"decompression time": "Decompression time (s)",
							"compression memory": "$\\frac{\\mbox{Peak compr. mem. (bytes)}}{\\mbox{\# tensor elements}}$" if output_mode == LATEX else "Peak comp. mem. (b/el)",
							"decompression memory": "$\\frac{\\mbox{Peak decompr. mem. (bytes)}}{\\mbox{\# tensor elements}}$" if output_mode == LATEX else "Peak comp. mem. (b/el)",
							"error control": "Target PSNR - PSNR (dB)" if use_psnr else ("$\\frac{\\mbox{Relative error}}{\\mbox{Target relative error}}$" if output_mode == LATEX else "(Relative error)/(target relative error)")
						}
					}[target][metric])
			
			if use_psnr:
				ax.set_xscale("linear")
				ax.set_xlim(psnr_lower_bound, psnr_upper_bound)
				ax.set_xticks(np.linspace(psnr_lower_bound, psnr_upper_bound, 5))
			else:
				ax.set_xscale("log")
				ax.set_xlim(rel_error_lower_bound, rel_error_upper_bound)
			if dataset_string == "Hurricane" and metric == "decompression time" and target == "paper":
				ax.set_yticks([40, 50, 60, 70, 80, 90, 100, 200, 300])
				ax.set_yticklabels(["", "", "", "", "", "", "$10^2$", "", ""])
			if metric_index != len(metrics) - 1:
				ax.set_xticklabels([])
			if len(datasets) > 1 and (row == 0 or (wrap_factor > 1 and row%len(metrics) == 0)):
				ax.set_title(dataset_string)
			if row == rows - 1:
				ax.set_xlabel("Relatieve fout" if target == "report" else ("PSNR (dB)" if use_psnr else "Relative error"))
	
	if output_mode == IMAGE:
		plt.subplot(gridspecs[0][0, 0]).legend()
		plt.show()
		plt.close()
	else:
		
		plt.subplots_adjust(wspace=0.5 if wide_view else 0.3, hspace=0.4) # If wide_view, reserve space for y-labels
		set_plot_width(fig, 1.13)
		
		# Add legend correctly
		handles, labels = plt.subplot(gridspecs[0][0, 0]).get_legend_handles_labels()
		legend_rows = 2 if "ATC (RTMSS = 0.0)" in labels else 1
		legend_columns = ceil(len(labels)/legend_rows)
		if "ATC (RTMSS = 0.1)" in labels:
			labels = ["ATC (RTMSS = 0.5)" if label == "ATC" else label for label in labels]
		loc = "lower center"
		position = {
			(4, 1): -0.14,
			(2, 2): -0.31,
			(2, 1): -0.25
		}[(rows, legend_rows)]
		fig.legend(handles, labels, loc="lower center", bbox_to_anchor=(0.5, position), bbox_transform=plt.gcf().transFigure, ncol=legend_columns) # (-0.3 if rows == 1 else -0.2) earlier
		
		fig.savefig(output_path, bbox_inches="tight", transparent=True)
		plt.close()

# Verified to almost correspond with the gain computations done earlier in generate_plot_grid
# For some metrics, the computed values can differ by a relative margin of e.g. 0.004% due to a slightly different way of computing these values
# In other words, these values aren't significantly different in any way and can be used to represent the gains achieved
def compute_performance(datasets, selected_compressors, optimize_atc=False, directory=""):
	
	# Stores performance metric values in absolute numbers (as in, not relative to other compressors)
	# All performance metrics apart from error control are computed as the geometric means of their respective values,
	# allowing us to easily compute the geometric mean of the achieved gains by dividing two of these means.
	# When accumulating logarithms, we use log10, not log or ln!
	# Also, memory metrics are represented in bytes/element.
	# Average error control is computed as the arithmetic average of the absolute values of the logarithmic relative differences, then raised to the power of 10 again if (not use_psnr).
	performance = {}
	metrics=["compression factor", "compression time", "decompression time", "compression memory", "decompression memory", "error control"]
	optimized_atc_name = "ATC (optimal RTMSS)"
	use_medians = True
	for metric in metrics:
		performance[metric] = {}
		for compressor in selected_compressors:
			performance[metric][compressor.name] = 0
		if optimize_atc:
			performance[metric][optimized_atc_name] = 0
	performance_samples_amount = 0
	
	INTERPOLATION = 0
	ACCUMULATION = 1
	REDUCTION = 2
	identity = lambda x: x
	def domain_map(metric, target_domain, reduction_amount=None):
		if metric == "error control":
			if target_domain == INTERPOLATION:
				return identity if use_psnr else log10
			elif target_domain == ACCUMULATION:
				return abs
			elif target_domain == REDUCTION:
				return (lambda value: value/reduction_amount) if use_psnr else (lambda value: 10**(value/reduction_amount) - 1)
		else:
			if target_domain == INTERPOLATION:
				return log10
			elif target_domain == ACCUMULATION:
				return identity
			elif target_domain == REDUCTION:
				return (lambda value: 10**(value/reduction_amount))
	def metric_mean(metric, vals):
		# This function is used to compute the mean value of a set of values of a certain metric belonging to a single average-error-point
		if use_medians:
			return np.median(vals)
		else:
			if metric == "error control" and use_psnr:
				return np.mean(vals)
			else:
				return 10**(np.mean(np.log10(vals)))
	
	for dataset in datasets:
		
		# Collect and filter data
		# Format:
		# data[compressor.name] = [(error, {"compression factor": ..., "error control": ..., "compression time": ..., "decompression time": ..., "compression memory": ..., "decompression memory": ...}), ...]
		# where error is expressed in PSNR if use_psnr, else relative error,
		# error control is expressed in (target PSNR - actual PSNR) if use_psnr or (relative_error/target_relative_error),
		# and memory is expressed in bytes/element.
		# All tuples are also ordered in order of decreasing error (so decreasing relative error, or increasing PSNR)
		# "samples" contains data points per error, while "average samples" contains average data points (arithmetic means error control in terms of PSNR deviation, geometric means for everything else) over a set of samples for the same target error
		# When aggregating data, we ensure that, per error, there is no deviation for target errors, and not too much deviation for the compression factor (< 1%) and memory metrics (< 35% for compression, < 20% for decompression)
		# We also ensure that we only aggregate data over errors that are very close to each other (< 2.6% margin)
		data = {}
		metric_mappings = {comp_string + " " + base_metric: (comp_string, {"time": "real_time", "memory": "peak_memory"}[base_metric]) for comp_string in ("compression", "decompression") for base_metric in ("time", "memory")}
		metric_mappings["compression factor"] = ("compression", "compression_factor")
		# We don't include error control in metric mappings, since the mapping is more complex and should only happen for certain compressors
		all_compressors_available = True
		for compressor in selected_compressors:
			path = results_path + directory + ("/" if directory != "" else "") + ("%s_%s.json"%(dataset.name, compressor.name)).replace(" ", "_")
			if os.path.isfile(path):
				
				with open(path, "r") as f:
					compressor_data = json.load(f)[dataset.name][compressor.name]
				
				samples_per_error = {}
				for target_error, samples in compressor_data.items():
					
					log_errors = []
					sample_data = {metric: [] for metric in metrics}
					if len(samples) == 0:
						continue
					for sample in samples:
						if sample["relative_error"] < rel_error_filter_bound:
							
							# Handle error and general metrics
							error = sample["psnr" if use_psnr else "relative_error"]
							log_errors.append(error if use_psnr else -20*log10(error))
							for key, (comp_string, sample_key) in metric_mappings.items():
								sample_data[key].append(sample[comp_string][sample_key])
							
							# Handle error control
							if has_error_control(compressor):
								compressor_uses_psnr = not compressor.name.startswith("TuckerMPI")
								target_error = float(target_error)
								psnr_deviation = target_error + (20*log10(2) if compressor.name == "TTHRESH" else 0) - sample["psnr"] if compressor_uses_psnr else 20*log10(sample["relative_error"]/target_error)
								sample_data["error control"].append(psnr_deviation if use_psnr else 10**(psnr_deviation/20))
					
					deviation_of_log_errors = max(log_errors) - min(log_errors)
					relative_deviation_of_error = 10**(deviation_of_log_errors/20) - 1
					if relative_deviation_of_error > 0.026:
						print("Warning for dataset %s, compressor %s and target error %s: samples contain log errors = %s, which deviate by %s, skipping this set of samples"%(dataset.name, compressor.name, target_error, log_errors, relative_deviation_of_error))
						continue
					avg_log_error = np.mean(log_errors)
					avg_error = avg_log_error if use_psnr else 10**(-avg_log_error/20)
					
					if avg_error not in samples_per_error.keys():
						samples_per_error[avg_error] = {metric: [] for metric in metrics}
					for metric in metrics:
						samples_per_error[avg_error][metric].extend(sample_data[metric])
				
				avg_errors_to_check = list(samples_per_error.keys())
				for avg_error in avg_errors_to_check:
					metric_deviation_margins = {"compression factor": 0.1, "compression memory": 0.35, "decompression memory": 0.2}
					for metric in metrics:
						vals = samples_per_error[avg_error][metric]
						avg_val = metric_mean(metric, vals)
						if metric in metric_deviation_margins:
							# Check if deviation isn't too large, should be quite deterministic
							min_val, max_val = min(vals), max(vals)
							rel_deviation = (max_val - min_val)/avg_val
							if rel_deviation > metric_deviation_margins[metric]:
								print("Warning for dataset %s, compressor %s, target error %s and metric \"%s\": samples contain values = %s, which deviate by %s, more than the allowed margin of %s. Discarding this set of samples."%(dataset.name, compressor.name, target_error, metric, vals, rel_deviation, metric_deviation_margins[metric]))
								del samples_per_error[avg_error]
								break
						samples_per_error[avg_error][metric] = avg_val
						if metric.endswith("memory"):
							samples_per_error[avg_error][metric] /= dataset.size()
				
				# Convert to lists of tuples, sorted in order of decreasing error
				data[compressor.name] = sorted(samples_per_error.items(), reverse=not use_psnr)
				
			else:
				# No data available for this compressor
				all_compressors_available = False
				print("Compressor %s not available for datasets %s"%(compressor.name, dataset.name))
				break
		
		if not all_compressors_available:
			continue
		
		# For this dataset, iterate through the target interpolation errors
		# Then: if, in this point, data is available for all selected compressors for both standard metrics and averaged metrics,
		# compute the interpolated metric value for each compressor and add it to the corresponding accumulator
		def higher_error(x, y):
			return x < y if use_psnr else x > y
		indices = {}
		for compressor in selected_compressors:
			indices[compressor.name] = -1
		for target_error in target_interpolation_errors:
			
			# Compute correct indices for all compressors for both sample strings
			all_data_available = True
			for compressor in selected_compressors:
				data_points = data[compressor.name]
				for index in range(indices[compressor.name], len(data_points)):
					if index == -1:
						if higher_error(target_error, data_points[0][0]):
							# Target error is higher that largest error in data, so stop
							all_data_available = False
							break
					elif index == len(data_points) - 1:
						# If we reached this point the target error is too low for this data, so stop
						all_data_available = False
						break
					elif not higher_error(target_error, data_points[index][0]) and higher_error(target_error, data_points[index + 1][0]):
						# We just found the correct interval
						break
				indices[compressor.name] = index
				if not all_data_available:
					break
			
			# If all indices were correctly updated, we can compare all metrics and compressors here, so we can apply the interpolation
			if all_data_available:
				performance_samples_amount += 1
				optimal_atc_compressor_name = None # ATC variant with the highest compression factor at this point
				if optimize_atc and metrics[0] != "compression factor":
					raise Exception("Compression factor should be the first metric when trying to optimize ATC over it!")
				for metric in metrics:
					atc_log_compression_factors = {}
					for compressor in selected_compressors:
						if metric == "error control" and not has_error_control(compressor):
							# This metric is irrelevant for this compressor, skip
							continue
						index = indices[compressor.name]
						data_points = data[compressor.name][index:index + 2]
						error_to_interpolation_domain = identity if use_psnr else log10
						x1, x2 = map(error_to_interpolation_domain, (data_points[0][0], data_points[1][0]))
						y1, y2 = map(domain_map(metric, INTERPOLATION), (data_points[0][1][metric], data_points[1][1][metric]))
						slope = (y2 - y1)/(x2 - x1)
						intercept = y1 - slope*x1
						interpolated_value = slope*error_to_interpolation_domain(target_error) + intercept
						val_to_be_added = domain_map(metric, ACCUMULATION)(interpolated_value)
						performance[metric][compressor.name] += val_to_be_added
						if optimize_atc and compressor.name.startswith("ATC"):
							if metric == "compression factor":
								atc_log_compression_factors[compressor.name] = val_to_be_added
							elif optimal_atc_compressor_name is not None and compressor.name == optimal_atc_compressor_name:
								performance[metric][optimized_atc_name] += val_to_be_added
					if optimize_atc and metric == "compression factor":
						optimal_atc_compressor_name = max(atc_log_compression_factors, key=atc_log_compression_factors.get)
						performance[metric][optimized_atc_name] += atc_log_compression_factors[optimal_atc_compressor_name]
	
	# After accumulating all valid data points within the accumulation domain,
	# apply the corresponding reduction to them in order to obtain the corresponding absolute performance metric
	for metric in metrics:
		for compressor_name in performance[metric]:
			if performance_samples_amount > 0:
				performance[metric][compressor_name] = domain_map(metric, REDUCTION, reduction_amount=performance_samples_amount)(performance[metric][compressor_name])
			else:
				performance[metric][compressor_name] = None
	
	if optimize_atc:
		for metric in metrics:
			for compressor_name in list(performance[metric].keys()):
				if compressor_name.startswith("ATC") and compressor_name != optimized_atc_name:
					del performance[metric][compressor_name]
	
	return performance

# Computes the geometric mean of the gains achieved in relative metrics (compression factor, times) compared to the first compressor
# Performance values for the first compressor are kept in absolute numbers
def compute_relative_performance(datasets, selected_compressors, optimize_atc=False, directory=""):
	
	performance = compute_performance(datasets, selected_compressors, optimize_atc=optimize_atc, directory=directory)
	reference = selected_compressors[0].name
	relative_metrics = ["compression factor", "compression time", "decompression time"]
	for metric in relative_metrics:
		for compressor_name in performance[metric]:
			if compressor_name != reference:
				performance[metric][compressor_name] /= performance[metric][reference]
	return performance

def generate_incremental_performance_table(datasets, compressor_sets, path, directory=""):
	
	all_compressors = [compressor for compressor_set in compressor_sets for compressor in compressor_set]
	if TTHRESH in all_compressors:
		datasets = list(filter((lambda dataset: dataset != hurricane), datasets))
	performance = compute_performance(datasets, all_compressors, directory=directory)
	if TuckerMPI in all_compressors:
		# For a proper memory comparison with TuckerMPI, TuckerMPI-1-process should be used
		# However, this compressor only produces data for the 5 smallest datasets, so we only want to compare memory usage in these cases
		# For the other metrics we still want to do a comparison that's as general as possible
		TuckerMPI_1_process = generate_tuckermpi_compressor(single_threaded=True)
		adjusted_compressors = [compressor if compressor != TuckerMPI else TuckerMPI_1_process for compressor in all_compressors]
		memory_performance = compute_performance([foot, brain, moffett_field, isotropic_fine_pressure_time, isotropic_fine_velocity], adjusted_compressors, directory=directory)
		for metric in ["compression memory", "decompression memory"]:
			performance[metric][TuckerMPI.name] = memory_performance[metric][TuckerMPI_1_process.name]
	
	table_header = """\\begin{tabular}{l|rrrrrrrrr}
\\multicolumn{1}{c|}{\\multirow{2}{*}{Compressor}} & \\multicolumn{1}{c}{\\multirow{2}{*}{ \\begin{tabular}{@{}c@{}} Rel. \\\\ comp. \\end{tabular} }} & \\multicolumn{2}{c}{Speed-up}                            & \\multicolumn{4}{c}{Peak memory (bytes/element)}                   & \\multicolumn{2}{c}{\\multirow{2}{*}{ \\begin{tabular}{@{}c@{}} Error \\\\ control \\end{tabular} }} \\\\
\\multicolumn{1}{c|}{}                            & \\multicolumn{1}{c}{}                                   & \\multicolumn{1}{c}{Comp.} & \\multicolumn{1}{c}{Decomp.} & \\multicolumn{2}{c}{Comp.} & \\multicolumn{2}{c}{Decomp.} & \\multicolumn{2}{c}{}                               \\\\ \\hline
"""
	table_footer = """\\end{tabular}"""
	
	compressor_descriptions = {
		"TTHRESH": "TTHRESH",
		"ATC (OPT 0)": "Baseline ATC",
		"ATC (OPT 1)": "+ Rank truncation",
		"ATC (OPT 2)": "+ Parallel quant./encoding ",
		"ATC (OPT 3)": "+ Predict dequant. correction",
		"ATC (OPT 4)": "+ Householder compression",
		"ATC (OPT 5)": "+ Mode storage order heuristic",
		"ATC": "+ Split bit plane truncation",
		"ATC (OPT 0, no BLAS)": "Baseline ATC (no BLAS)",
		"ATC (OPT 1, no BLAS)": "+ Rank truncation (no BLAS)",
		"ATC (OPT 2, no BLAS)": "+ Parallel quant./encoding (no BLAS)",
		"ATC (OPT 3, no BLAS)": "+ Predict dequant. correction (no BLAS)",
		"ATC (OPT 4, no BLAS)": "+ Householder compression (no BLAS)",
		"ATC (OPT 5, no BLAS)": "+ Mode storage order heuristic (no BLAS)",
		"ATC (no BLAS)": "+ Split bit plane truncation (no BLAS)",
		"TuckerMPI": "TuckerMPI",
		"TuckerMPI (SMF)": "+ SMF",
	}
	reference_compressor_name = compressor_sets[0][0].name
	metrics = ["compression factor", "compression time", "decompression time", "compression memory", "decompression memory", "error control"]
	absolute_metrics = ["compression memory", "decompression memory", "error control"]
	higher_is_better_metrics = ["compression factor", "compression time", "decompression time"]
	invert_metrics = ["compression time", "decompression time"] # We want to show the average speed-ups instead of average times
	convert_to_percentage_metrics = ["error control"]
	
	# First apply inversion and *100 to all relevant metrics
	for compressor in all_compressors:
		for metric in invert_metrics:
			performance[metric][compressor.name] = 1/performance[metric][compressor.name]
		for metric in convert_to_percentage_metrics:
			performance[metric][compressor.name] *= 100
	
	def highlight_gain_string(gain_string, gain_percentage, higher_is_better, speedup=False):
		if abs(gain_percentage) >= (10 if speedup else 1):
			color_string = "dark-green" if ((gain_percentage > 0) == higher_is_better) else "red"
			gain_string = "\\textcolor{" + color_string + "}{" + gain_string + "}"
		if abs(gain_percentage) >= (50 if speedup else 10):
			gain_string = "\\textbf{" + gain_string + "}"
		return gain_string
	
	# Then generate the full string
	string = table_header
	for i, compressor_set in enumerate(compressor_sets):
		for j, compressor in enumerate(compressor_set):
			
			last_compressor_in_set = (j == len(compressor_set) - 1)
			max_runs = 1 + int(len(compressor_set) > 1 and last_compressor_in_set)
			for run in range(max_runs):
				
				compare_to_reference_compressor = (j == 0 or run == 1)
				if compare_to_reference_compressor:
					if run == 1:
						compressor_description = "Full " + compressor.name
					else:
						compressor_description = compressor_descriptions[compressor.name]
					string += "\\textbf{%s}"%compressor_description
				else:
					string += compressor_descriptions[compressor.name]
				
				for metric in metrics:
					performance_value = performance[metric][compressor.name]
					# Print absolute metric first if necessary
					if metric in absolute_metrics:
						string += (" & %.2f"%performance_value if performance_value < 0.15 else " & %.1f"%performance_value) + int(metric in convert_to_percentage_metrics)*"\\%"
					# Print gain afterwards
					if compressor.name == reference_compressor_name:
						string += " & -"
					else:
						# Print incremental gain
						current_reference_compressor_name = reference_compressor_name if compare_to_reference_compressor else compressor_sets[i][j - 1].name
						reference_value = performance[metric][current_reference_compressor_name]
						gain_to_show = performance_value/reference_value
						gain_percentage = (gain_to_show - 1)*100
						string += " & "
						gain_string = ("+" if gain_percentage >= 0 else "") + ("%.1f\\%%"%(gain_percentage) if abs(gain_percentage) < 10 else "%s\\%%"%(round(gain_percentage)))
						string += highlight_gain_string(gain_string, gain_percentage, metric in higher_is_better_metrics, speedup=metric.endswith(" time"))
				
				# Finish line
				if i < len(compressor_sets) - 1 or j < len(compressor_set) - 1 + max_runs - 1:
					string += " \\\\"
					if last_compressor_in_set and run == max_runs - 1:
						string += " \\hline"
				string += "\n"
	
	string += table_footer
	
	with open(path, "w") as f:
		f.write(string)

def generate_ATC_compression_examples(dataset, target_rel_errors, path, selected_modes=[0, 1, 2]):
	
	# Import tensor plotter
	sys.path.append("../../../Other/Code/")
	from tensor_plotter import plot_tensor
	
	# Load dataset and get basic metrics
	original = load_data(dataset.path, dataset.dtype)
	data_norm = np.linalg.norm(original)
	data_range = compute_data_range(original)
	
	threads = 4
	print_commands = True
	decompressed_paths = [dataset.path]
	image_paths = [path + "compression-example_%s_original.png"%dataset.name]
	for target_rel_error in target_rel_errors:
		
		print("Target relative error:", target_rel_error)
		target_rel_error_string = str(target_rel_error).replace(".", "-").ljust(5, "0")
		make_temp_paths_unique("compression_example_%s"%target_rel_error_string)
		if not os.path.isfile(get_decompressed_path()):
			
			# Perform compression
			target_psnr = relative_error_to_psnr(target_rel_error, data_range, original.size, data_norm)
			compression_results = compress(ATC, dataset, target_psnr, threads=threads, print_commands=print_commands)
			decompress(ATC, dataset, threads=threads, print_commands=print_commands)
			
			# Check error
			decompressed = load_data(get_decompressed_path(), dataset.dtype)
			psnr = compute_psnr(original, decompressed, data_range)
			rel_error =  psnr_to_relative_error(psnr, data_range, original.size, data_norm)
			print("Compressor ATC resulted in relative error %s (PSNR: %s dB) and compression factor %s"%(rel_error, psnr, compression_results["compression_factor"]))
			
			# Clean-up
			os.remove(get_compressed_path())
		
		decompressed_paths.append(get_decompressed_path())
		image_paths.append(path + "compression-example_%s_%s.png"%(dataset.name, target_rel_error_string))
	
	plot_tensor(decompressed_paths, dataset.dtype, dataset.sizes, image_paths, selected_modes=selected_modes)
	
	for decompressed_path in decompressed_paths[1:]:
		os.remove(decompressed_path)

def generate_slice_images(dataset, amount=4):
	
	selected_mode = dataset.most_constant_mode
	data = load_dataset(dataset).reshape(dataset.sizes)
	
	slice_indices = np.linspace(0, dataset.sizes[selected_mode], amount + 1)[:-1]
	slice_indices = slice_indices + (slice_indices[1] - slice_indices[0])/2 # Shift indices to middle of each interval
	slice_indices = list(map(lambda index: max(0, min(dataset.sizes[selected_mode], int(index))), list(slice_indices)))
	slices = None
	for i, slice_index in enumerate(slice_indices):
		new_slice = data[tuple((slice_index if mode == selected_mode else slice(None)) for mode in range(dataset.modes))]
		if new_slice.shape[0] < new_slice.shape[1]:
			# Make sure image has rows >= columns (portrait mode)
			new_slice = new_slice.transpose()
		if slices is None:
			slices = np.empty((len(slice_indices), new_slice.shape[0], new_slice.shape[1]), dtype=data.dtype)
		slices[i, :, :] = new_slice
	
	for slice_index_in_list, slice_index_in_data in enumerate(slice_indices):
		data_slice = slices[slice_index_in_list, :, :]
		rmsv = np.linalg.norm(data_slice)/np.sqrt(data_slice.size)
		vmax = 2*rmsv
		path = figure_path + dataset.name + "_slice-%s.png"%slice_index_in_data
		im = Image.fromarray(np.clip(np.rint(data_slice/vmax*256), 0, 255).astype("uint8"))
		im.save(path)

# Figure generator suites

def generate_figures_paper():
	
	introduction = False
	gains = False
	rate_distortion = True
	timings = False
	memory_usage = False
	error_control = False
	opt_levels_comparison = False
	
	datasets = [foot, brain, moffett_field, isotropic_fine_pressure_time, isotropic_fine_velocity, deforest_globe_8, deforest_globe_33, hurricane]
	
	# Introduction
	if introduction:
		generate_slice_images(moffett_field)
		generate_ATC_compression_examples(isotropic_fine_pressure_time, [0.001, 0.01, 0.1], figure_path, selected_modes=[1, 2, 3])
		"""
		Results with ATC v1.1.2 (31 October 2022)
		Target relative error: 0.001
		../compressors/atc-blas --verbose --print-results --of ../data/turbulence/isotropic1024fine-pressure-time-128 --cf ../tmp/compressed-compression_example_0-001 --type float32 --storage-order msi --sizes 100 128 128 128 --psnr 79.3251316268004 --threads 4 > ../logs/ATC_compress_2022-10-31_16:14:26.txt
		../compressors/atc-blas --verbose --print-results --cf ../tmp/compressed-compression_example_0-001 --df ../tmp/decompressed-compression_example_0-001 --threads 4 > ../logs/ATC_decompress_2022-10-31_16:14:32.txt
		Compressor ATC resulted in relative error 0.0009997951544609414 (PSNR: 79.32691107480773 dB) and compression factor 389.2764613863233
		Target relative error: 0.01
		../compressors/atc-blas --verbose --print-results --of ../data/turbulence/isotropic1024fine-pressure-time-128 --cf ../tmp/compressed-compression_example_0-010 --type float32 --storage-order msi --sizes 100 128 128 128 --psnr 59.3251316268004 --threads 4 > ../logs/ATC_compress_2022-10-31_16:14:40.txt
		../compressors/atc-blas --verbose --print-results --cf ../tmp/compressed-compression_example_0-010 --df ../tmp/decompressed-compression_example_0-010 --threads 4 > ../logs/ATC_decompress_2022-10-31_16:14:45.txt
		Compressor ATC resulted in relative error 0.010004808975211896 (PSNR: 59.320955608042205 dB) and compression factor 3568.6485382704286
		Target relative error: 0.1
		../compressors/atc-blas --verbose --print-results --of ../data/turbulence/isotropic1024fine-pressure-time-128 --cf ../tmp/compressed-compression_example_0-100 --type float32 --storage-order msi --sizes 100 128 128 128 --psnr 39.3251316268004 --threads 4 > ../logs/ATC_compress_2022-10-31_16:14:53.txt
		../compressors/atc-blas --verbose --print-results --cf ../tmp/compressed-compression_example_0-100 --df ../tmp/decompressed-compression_example_0-100 --threads 4 > ../logs/ATC_decompress_2022-10-31_16:15:00.txt
		Compressor ATC resulted in relative error 0.10538248897334544 (PSNR: 38.86976259220275 dB) and compression factor 95033.51081907783
		Plotting 4 versions of tensor ...
		Extracted slices from ../data/turbulence/isotropic1024fine-pressure-time-128
		Extracted slices from ../tmp/decompressed-compression_example_0-001
		Extracted slices from ../tmp/decompressed-compression_example_0-010
		Extracted slices from ../tmp/decompressed-compression_example_0-100
		Saved image as ../../../Manuscripts/atc-paper/figures/compression-example_Isotropic-PT_original.png
		Saved image as ../../../Manuscripts/atc-paper/figures/compression-example_Isotropic-PT_0-001.png
		Saved image as ../../../Manuscripts/atc-paper/figures/compression-example_Isotropic-PT_0-010.png
		Saved image as ../../../Manuscripts/atc-paper/figures/compression-example_Isotropic-PT_0-100.png
		"""
	
	# Other figures from Algorithmic improvements section: see atc/research/experiments.py
	
	# Results
	
	# Performance gains compared to TTHRESH
	if gains:
		generate_incremental_performance_table(datasets, [[TTHRESH]] + [[generate_atc_compressor(optimization_level=level) for level in range(ATC_MAX_OPTIMIZATION_LEVEL + 1)]] + [[TuckerMPI]], table_path + "incremental-performance.tex")
		#print(compute_relative_performance(datasets, [TTHRESH, ATC, SZ]))
		#print(compute_relative_performance(datasets, [TTHRESH] + [generate_atc_compressor(rtmss=rtmss) for rtmss in (0.0, 0.2, 0.4, 0.6, 0.8)], optimize_atc=True))
		#print(compute_relative_performance(datasets, [SZ, ATC]))
		#generate_plot_grid([moffett_field, isotropic_fine_pressure_time], selected_compressors=[ATC, TTHRESH, SZ, TuckerMPI], metrics=["error control"], target="paper", output_mode=IMAGE)
		"""
		Results with ATC v1.1.2 (31 October 2022):
		{'compression factor': {'TTHRESH': 24.023859450577596, 'ATC': 1.0106867443691758, 'SZ': 0.5008679282649509}, 'compression time': {'TTHRESH': 39.627381909046086, 'ATC': 0.2874198305290824, 'SZ': 0.0921476800151344}, 'decompression time': {'TTHRESH': 19.069331757523575, 'ATC': 0.4519189887861299, 'SZ': 0.05344642653027258}, 'compression memory': {'TTHRESH': 40.26837423809019, 'ATC': 18.577948816409098, 'SZ': 24.434102757546487}, 'decompression memory': {'TTHRESH': 38.66679369451881, 'ATC': 16.911430625231837, 'SZ': 8.003448982276446}, 'error control': {'TTHRESH': 0.33828808988780357, 'ATC': 0.01446038491929258, 'SZ': 0.3169098144096467}}
		{'compression factor': {'TTHRESH': 24.023859450577596, 'ATC (optimal RTMSS)': 1.080225644029098}, 'compression time': {'TTHRESH': 39.627381909046086, 'ATC (optimal RTMSS)': 0.35781233099166}, 'decompression time': {'TTHRESH': 19.069331757523575, 'ATC (optimal RTMSS)': 0.5451940072474609}, 'compression memory': {'TTHRESH': 40.26837423809019, 'ATC (optimal RTMSS)': 18.982493374756746}, 'decompression memory': {'TTHRESH': 38.66679369451881, 'ATC (optimal RTMSS)': 16.904070096837877}, 'error control': {'TTHRESH': 0.33828808988780357, 'ATC (optimal RTMSS)': 0.013804716506198744}}
		{'compression factor': {'SZ': 15.893831298537373, 'ATC': 1.9783691634403695}, 'compression time': {'SZ': 5.786135789871105, 'ATC': 3.0721896534545774}, 'decompression time': {'SZ': 1.570566115721239, 'ATC': 8.398727943354219}, 'compression memory': {'SZ': 23.836751781419604, 'ATC': 18.235849660695468}, 'decompression memory': {'SZ': 7.912353450347964, 'ATC': 16.80514412214238}, 'error control': {'SZ': 0.38248570751426825, 'ATC': 0.01868801151177557}}
		"""
		"""
		# Temporary test to compare old results
		print("Comparing current results with old results:")
		keys = ["compression time", "decompression time"]
		compressor = SZ#generate_atc_compressor(optimization_level=0)
		for dataset in datasets:
			old_results = select_keys(compute_performance([dataset], [compressor], directory="may-2022"), keys)
			new_results = select_keys(compute_performance([dataset], [compressor]), keys)
			if old_results[keys[0]][compressor.name] is not None and new_results[keys[0]][compressor.name] is not None:
				print("Dataset:", dataset.name)
				for key in keys:
					print("%s: %s"%(key, new_results[key][compressor.name]/old_results[key][compressor.name]))
		#print(select_keys(compute_performance(datasets, [ATC, generate_atc_compressor(optimization_level=0), TTHRESH, SZ, ZFP]), ["compression time", "decompression time"]))
		#print(select_keys(compute_performance(datasets, [ATC, generate_atc_compressor(optimization_level=0), TTHRESH, SZ, ZFP], directory="may-2022"), ["compression time", "decompression time"]))
		"""
		#generate_plot_grid(datasets, selected_compressors=[generate_atc_compressor(optimization_level=level) for level in range(ATC_MAX_OPTIMIZATION_LEVEL)] + [ATC], metrics=["compression time", "decompression time"], target="paper", output_mode=IMAGE)
		#generate_plot_grid(datasets, selected_compressors=[generate_atc_compressor(blas=True, epcm=False, run=run) for run in [9]] + [generate_atc_compressor(blas=True, epcm=False, lpmt=lpmt, qrem=0.001, run=9) for lpmt in [4, 8, 16]], metrics=["compression time", "decompression time"], target="paper", output_mode=IMAGE)
		#generate_plot_grid(datasets, selected_compressors=[generate_atc_compressor(blas=True, epcm=False, run=run) for run in [6, 7, 8, 9]] + [generate_atc_compressor(blas=True, epcm=False, qrem=0.001, lpmt=lpmt, run=9) for lpmt in [4, 8, 16]], metrics=["error control"], target="paper", output_mode=IMAGE)
		#generate_plot_grid(datasets, selected_compressors=[generate_atc_compressor(blas=blas, epcm=False, run=run) for run in [9] for blas in [True, False]], metrics=["compression time", "decompression time"], target="paper", output_mode=IMAGE)
		#generate_plot_grid(datasets, selected_compressors=[generate_atc_compressor(blas=blas, epcm=epcm, run=6) for blas in [True, False] for epcm in [True, False]], metrics=["compression memory", "decompression memory"], target="paper", output_mode=IMAGE)
	
	# Rate-distortion (general comparison + RTMSS effect)
	if rate_distortion:
		generate_plot_grid(datasets, selected_compressors=[ATC, TTHRESH, x265, SZ, ZFP, TuckerMPI], metrics=["compression factor", "relative compression factor"], target="paper", output_mode=LATEX, output_path=figure_path + "rate-distortion.pgf")
		generate_plot_grid(datasets, selected_compressors=[TTHRESH] + [generate_atc_compressor(rtmss=rtmss) for rtmss in (0.0, 0.2, 0.4, 0.6, 0.8)], metrics=["relative compression factor"], target="paper", output_mode=LATEX, output_path=figure_path + "rate-distortion-rtmss.pgf", venue="shared")
	
	# Compression/decompression time
	if timings:
		generate_plot_grid(datasets, selected_compressors=[ATC, TTHRESH, x265, SZ, ZFP, TuckerMPI], metrics=["compression time", "decompression time",], target="paper", output_mode=LATEX, output_path=figure_path + "times.pgf")
	
	# Peak memory usage
	if memory_usage:
		generate_plot_grid(datasets, selected_compressors=[ATC, TTHRESH, x265, SZ, ZFP, generate_tuckermpi_compressor(single_threaded=True)], metrics=["compression memory", "decompression memory"], target="paper", output_mode=LATEX, output_path=figure_path + "memory.pgf")

	# Error control
	if error_control:
		generate_plot_grid(datasets, selected_compressors=[ATC, TTHRESH, SZ, TuckerMPI], metrics=["error control",], target="paper", output_mode=LATEX, output_path=figure_path + "error-control.pgf")
	
	# Just for viewing, not actually in the paper for now
	if opt_levels_comparison:
		generate_plot_grid(datasets, selected_compressors=[TTHRESH] + ATC_ALL_OPT_LEVELS + [ATC], metrics=["compression time", "decompression time",], target="paper", output_mode=IMAGE, output_path=None)
		

def generate_figures_paper_extras():
	
	# T-HODCT
	"""
	generate_plot_grid([foot, brain, moffett_field, isotropic_fine_pressure_time, isotropic_fine_velocity, deforest_globe_8, deforest_globe_33, hurricane], selected_compressors=[ATC, generate_atc_compressor(dtmss=0.1), generate_atc_compressor(dtmss=0.3), generate_atc_compressor(dtmss=0.7), generate_atc_compressor(dtmss=0.9)], metrics=["compression time", "decompression time"], target="paper", output_mode=IMAGE, output_path=figure_path + "rate-distortion-dtmss.pgf")
	# Result: for all datasets significantly slower
	generate_plot_grid([hurricane,], selected_compressors=[ATC, generate_atc_compressor(dtmss=0.1), generate_atc_compressor(dtmss=0.3), generate_atc_compressor(dtmss=0.7), generate_atc_compressor(dtmss=0.9)], metrics=["relative compression factor", "compression time", "decompression time"], target="paper", output_mode=LATEX, output_path=figure_path + "rate-distortion-dtmss-hurricane.pgf")
	# Result: for most datasets, time losses of factor 2 to up to factor 10, similar compression rate usually
	"""
	
	# Effect of EPCM on compression/decompression time
	"""
	generate_plot_grid([foot, brain, moffett_field, isotropic_fine_pressure_time, isotropic_fine_velocity, deforest_globe_8, deforest_globe_33, hurricane], selected_compressors=[generate_atc_compressor(epcm=True), TTHRESH, x265, SZ, ZFP, FPZIP], metrics=["compression time",], target="paper", output_mode=IMAGE, output_path=figure_path + "times.pgf")
	# Result: Geometric average of ATC gains according to metric ['compression time']: 0.3876074978784951, EPCM speed-up: 1.102864511
	generate_plot_grid([foot, brain, moffett_field, isotropic_fine_pressure_time, isotropic_fine_velocity, deforest_globe_8, deforest_globe_33, hurricane], selected_compressors=[generate_atc_compressor(epcm=True), TTHRESH, x265, SZ, ZFP, FPZIP], metrics=["decompression time",], target="paper", output_mode=IMAGE, output_path=figure_path + "times.pgf")
	# Result: Geometric average of ATC gains according to metric ['decompression time']: 0.2761630387501534, EPCM speed-up: 1.343743998
	#generate_plot_grid([foot, brain, moffett_field, isotropic_fine_pressure_time, isotropic_fine_velocity, deforest_globe_8, deforest_globe_33, hurricane], selected_compressors=[ATC, generate_atc_compressor(epcm=True), TTHRESH, x265, SZ, ZFP, FPZIP], metrics=["compression time", "decompression time",], target="paper", output_mode=IMAGE, output_path=figure_path + "times.pgf")
	"""
	
	# Effect of RTMSS on timings
	"""
	generate_plot_grid([foot, brain, moffett_field, isotropic_fine_pressure_time, isotropic_fine_velocity, deforest_globe_8, deforest_globe_33, hurricane], selected_compressors=[generate_atc_compressor(rtmss=0.0), generate_atc_compressor(rtmss=0.2), generate_atc_compressor(rtmss=0.4), generate_atc_compressor(rtmss=0.6), generate_atc_compressor(rtmss=0.8), TTHRESH], metrics=["compression time", "decompression time", ], target="paper", output_mode=IMAGE, output_path=figure_path + "times-rtmss.pgf")
	"""
	
	# Effect of EPCM on peak memory usage
	"""
	generate_plot_grid([foot, brain, moffett_field, isotropic_fine_pressure_time, isotropic_fine_velocity, deforest_globe_8, deforest_globe_33, hurricane], selected_compressors=[generate_atc_compressor(epcm=True), TTHRESH, x265, SZ, ZFP, FPZIP], metrics=["compression memory",], target="paper", output_mode=IMAGE, output_path=figure_path + "memory.pgf")
	# Result: Geometric average of ATC gains according to metric ['compression memory']: 0.46659101302756073
	generate_plot_grid([foot, brain, moffett_field, isotropic_fine_pressure_time, isotropic_fine_velocity, deforest_globe_8, deforest_globe_33, hurricane], selected_compressors=[generate_atc_compressor(epcm=True), TTHRESH, x265, SZ, ZFP, FPZIP], metrics=["decompression memory",], target="paper", output_mode=IMAGE, output_path=figure_path + "memory.pgf")
	# Result: Geometric average of ATC gains according to metric ['decompression memory']: 0.4860484438156369
	"""
	
def generate_figures_readme():
	"""
	Results with ATC v1.0.1 (17 March 2021):
	Target relative error: 0.003
	../compressors/atc --verbose --print-results --of ../data/turbulence/isotropic1024fine-pressure-512 --cf ../tmp/compressed-compression_example_0-003 --type float32 --storage-order msi --sizes 512 512 512 --psnr 74.29210085574512 --threads 4
	../compressors/atc --verbose --print-results --cf ../tmp/compressed-compression_example_0-003 --df ../tmp/decompressed-compression_example_0-003 --threads 4
	Compressor ATC resulted in relative error 0.0030038062441371653 (PSNR: 74.28108763526603 dB) and compression factor 48.87858689620545
	Target relative error: 0.01
	../compressors/atc --verbose --print-results --of ../data/turbulence/isotropic1024fine-pressure-512 --cf ../tmp/compressed-compression_example_0-010 --type float32 --storage-order msi --sizes 512 512 512 --psnr 63.834525950138364 --threads 4
	../compressors/atc --verbose --print-results --cf ../tmp/compressed-compression_example_0-010 --df ../tmp/decompressed-compression_example_0-010 --threads 4
	Compressor ATC resulted in relative error 0.010018581557030674 (PSNR: 63.81840119129224 dB) and compression factor 157.39339369272412
	Target relative error: 0.03
	../compressors/atc --verbose --print-results --of ../data/turbulence/isotropic1024fine-pressure-512 --cf ../tmp/compressed-compression_example_0-030 --type float32 --storage-order msi --sizes 512 512 512 --psnr 54.29210085574512 --threads 4
	../compressors/atc --verbose --print-results --cf ../tmp/compressed-compression_example_0-030 --df ../tmp/decompressed-compression_example_0-030 --threads 4
	Compressor ATC resulted in relative error 0.030082988016335518 (PSNR: 54.26810653610949 dB) and compression factor 640.6107813154413
	Target relative error: 0.1
	../compressors/atc --verbose --print-results --of ../data/turbulence/isotropic1024fine-pressure-512 --cf ../tmp/compressed-compression_example_0-100 --type float32 --storage-order msi --sizes 512 512 512 --psnr 43.834525950138364 --threads 4
	../compressors/atc --verbose --print-results --cf ../tmp/compressed-compression_example_0-100 --df ../tmp/decompressed-compression_example_0-100 --threads 4
	Compressor ATC resulted in relative error 0.1014959606501937 (PSNR: 43.705550780471796 dB) and compression factor 4956.066982995772
	Saved image as ../../atc/images/compression-example_Isotropic1024-fine-pressure-512_original.png
	Saved image as ../../atc/images/compression-example_Isotropic1024-fine-pressure-512_0-003.png
	Saved image as ../../atc/images/compression-example_Isotropic1024-fine-pressure-512_0-010.png
	Saved image as ../../atc/images/compression-example_Isotropic1024-fine-pressure-512_0-030.png
	Saved image as ../../atc/images/compression-example_Isotropic1024-fine-pressure-512_0-100.png
	"""
	generate_ATC_compression_examples(isotropic_fine_pressure, [0.01, 0.03, 0.1], "../../atc/doc/images/")

def generate_figures_poster_and_numa_seminar():
	generate_rate_distortion_plot(isotropic_fine_pressure, output_mode=IMAGE, base_path=poster_figure_path)
	#generate_rate_distortion_plot(foot, output_mode=IMAGE, base_path=poster_figure_path)
	#generate_ATC_compression_examples(isotropic_fine_pressure_time, [0.001, 0.01, 0.1], poster_figure_path, selected_modes=[1, 2, 3])
	#generate_ATC_compression_examples(foot, [0.003, 0.01, 0.03, 0.1], poster_figure_path)
	#generate_rate_distortions_plots([isotropic_fine_pressure_time, moffett_field, foot], selected_compressors=[ATC, TTHRESH, x265, SZ, ZFP, FPZIP], output_mode=PRESENTATION, output_path=numa_seminar_figure_path + "normalized-rate-distortion-plots.pgf")
	#generate_rate_distortions_plots([isotropic_fine_pressure_time, moffett_field], selected_compressors=[generate_atc_compressor(0.1), generate_atc_compressor(0.3), ATC, generate_atc_compressor(0.7), generate_atc_compressor(0.9), TTHRESH], output_mode=PRESENTATION, output_path=numa_seminar_figure_path + "normalized-rate-distortion-plots-rtmss.pgf", close_comparison=True)
	#generate_times_plot([isotropic_fine_pressure_time, moffett_field, foot], selected_compressors=[ATC, TTHRESH, SZ, ZFP, FPZIP], output_mode=PRESENTATION, output_path=numa_seminar_figure_path + "times.pgf")
	#generate_memory_usage_error_control_plot([isotropic_fine_pressure_time, moffett_field, foot], selected_compressors=[ATC, TTHRESH, x265, SZ, ZFP, FPZIP], output_mode=IMAGE, output_path=numa_seminar_figure_path + "memory-usage-and-error-control.pgf")

def generate_figures_midterm_report():
	generate_plot_grid([foot, isotropic_fine_pressure_time, deforest_globe_8, hurricane], selected_compressors=[ATC, TTHRESH, x265, SZ, ZFP, FPZIP], metrics=["compression factor", "relative compression factor"], output_mode=LATEX, output_path=midterm_report_figure_path + "rate-distortion.pgf")
	#generate_plot_grid([foot, brain, moffett_field, isotropic_fine_pressure_time, isotropic_fine_velocity, deforest_globe_8, deforest_globe_33, hurricane], selected_compressors=[generate_atc_compressor(0.1), generate_atc_compressor(0.3), ATC, generate_atc_compressor(0.7), generate_atc_compressor(0.9), TTHRESH], metrics=["compression factor", "relative compression factor"], output_mode=IMAGE, output_path=midterm_report_figure_path + "rate-distortion-rtmss.pgf")
	#generate_plot_grid([foot, isotropic_fine_pressure_time, deforest_globe_8, hurricane], selected_compressors=[ATC, TTHRESH, x265, SZ, ZFP, FPZIP], metrics=["compression time", "decompression time"], output_mode=LATEX, output_path=midterm_report_figure_path + "times.pgf")
	#generate_plot_grid([foot, isotropic_fine_pressure_time, deforest_globe_8, hurricane], selected_compressors=[ATC, TTHRESH, x265, SZ, ZFP, FPZIP], metrics=["compression memory", "decompression memory"], output_mode=LATEX, output_path=midterm_report_figure_path + "memory.pgf")
	#generate_plot_grid([foot, isotropic_fine_pressure_time, deforest_globe_8, hurricane], selected_compressors=[ATC, TTHRESH, x265, SZ, ZFP, FPZIP], metrics=["error control"], output_mode=LATEX, output_path=midterm_report_figure_path + "error-control.pgf")

generate_figures_paper()
