/*
 * This file contains tests which automatically verify the correctness of most ATC functionality using the C++ interface.
 */

#include <iostream>
#include <fstream>
#include <vector>
#include <functional>
#include <tuple>
#include <experimental/filesystem>
#include <math.h>
#include <cmath>
#include <algorithm>
#include <numeric>
#include <limits>

#include "atc.hpp"

// PRELIMINARIES

/// The path to the directory where the datasets are stored
const std::string input_data_path = "../data/";
/// The path which temporarily stores compressed objects
const std::string compression_path = "/tmp/comp.atc";
/// The path which temporarily stores decompressed data
const std::string decompression_path = "/tmp/decomp.raw";
/// The path which temporarily stores other data
const std::string other_path = "/tmp/other.raw";

/// Stores both the path to a dataset and its metadata (which isn't stored in the file).
template<typename data_type>
struct Dataset {
	const std::string name;
	const std::string data_file;
	const std::vector<size_t> sizes;
};

/// A simple 5D dataset which flattens to [0, 1, 2, ..., 3^5 - 1].
const Dataset<uint8_t> basic_data = {
	"test-data-5D",
	"test_data_5D.raw",
	std::vector<size_t>{3, 3, 3, 3, 3}
};

/// A more complex 4D dataset with some spatial continuity along all modes.
const Dataset<float> complex_data = {
	"test-data-4x8x16x32",
	"test_data_4x8x16x32.raw",
	std::vector<size_t>{4, 8, 16, 32}
};

/// Gets the corresponding ATC data type object based on the dataset's template parameter
template<typename T>
atc::DataType get_data_type(const Dataset<T>&) {
	return atc::get_data_type<T>();
}

/// Loads a data file as a raw buffer of some type.
template<typename T>
std::vector<T> load_data(const std::string& path) {
	
	// Open file
	std::ifstream input_stream(path, std::ios::in | std::ios::binary);
	if (!input_stream.is_open())
		throw std::invalid_argument("Data file at " + path + " could not be opened");
	
	// Check size
	size_t file_size = std::experimental::filesystem::file_size(path);
	if (file_size%sizeof(T) != 0)
		throw std::invalid_argument("Data file at " + path + " contains " + std::to_string(file_size)
			+ " bytes, which is not a multiple of the size of the data type ( " + std::to_string(sizeof(T)) + " bytes)");
	size_t elements = file_size/sizeof(T);
	
	// Load data
	std::vector<T> res(elements);
	input_stream.read(reinterpret_cast<char*>(res.data()), elements*sizeof(T));
	return res;
	
}

/// Stores a raw buffer of some type as a file
template<typename T>
void write_data(const std::vector<T>& data, const std::string& path) {
	
	// Open file
	std::ofstream output_stream(path, std::ios::out | std::ios::binary);
	if (!output_stream.is_open())
		throw std::invalid_argument("Data file at " + path + " could not be opened");
	
	// Write data
	output_stream.write(reinterpret_cast<const char*>(data.data()), data.size()*sizeof(T));
	
}

/// Checks whether or not abs(a - b)/abs(a) < epsilon while also checking some special cases (NaN, Inf)
template<typename T>
bool approx_equal(T expected, T actual, T epsilon = 1e-9) {
	// Handle NaN
	if (isnan(expected))
		return isnan(actual);
	// Handle +-Inf (or normal equality)
	if (expected == actual)
		return true;
	// Handle normal numbers
	return abs(actual - expected)/expected < epsilon;
}

/**
 * Calculates the SSE in between two buffers.
 * @throw std::runtime_error If the buffers don't have the same size
 */
template<typename T>
double sse(const std::vector<T>& original, const std::vector<T>& decompressed) {
	if (original.size() != decompressed.size())
		throw std::runtime_error("Cannot compare two buffers of different size!");
	double sse = 0;
	for (size_t i = 0; i < original.size(); i++)
		sse += pow(original[i] - decompressed[i], 2);
	return sse;
}

/// Calculates the relative error in between two buffers.
template<typename T>
double relative_error(const std::vector<T>& original, const std::vector<T>& decompressed) {
	double norm = 0;
	for (size_t i = 0; i < original.size(); i++)
		norm += pow(original[i], 2);
	return sqrt(sse(original, decompressed)/norm);
}

/// Calculates the RMSE in between two buffers.
template<typename T>
double rmse(const std::vector<T>& original, const std::vector<T>& decompressed) {
	return sqrt(sse(original, decompressed)/original.size());
}

/// Calculates the PSNR in between two buffers.
template<typename T>
double psnr(const std::vector<T>& original, const std::vector<T>& decompressed) {
	double data_min = *std::min_element(original.begin(), original.end());
	double data_max = *std::max_element(original.begin(), original.end());
	return 20*log10((data_max - data_min)/rmse(original, decompressed));
}

/// Calculates the actual error in between the original and decompressed data and verifies if it is close to the target error
template<typename T>
double verify_target_error(double target_error, double error_margin, const std::vector<T>& original_data, const std::vector<T>& decompressed_data, const atc::Options& options) {
	double error = -1;
	double relative_error_on_error;
	bool zero_error;
	if (options.error_metric == atc::ErrorMetric::PSNR) {
		error = psnr(original_data, decompressed_data);
		relative_error_on_error = abs(pow(10, (error - target_error)/20) - 1);
		zero_error = isinf(error);
	} else {
		if (options.error_metric == atc::ErrorMetric::RelativeError)
			error = relative_error(original_data, decompressed_data);
		else if (options.error_metric == atc::ErrorMetric::RMSE)
			error = rmse(original_data, decompressed_data);
		relative_error_on_error = abs((error - target_error)/target_error);
		zero_error = (error == 0);
	}
	if (!zero_error && relative_error_on_error >= error_margin)
		throw std::runtime_error("Compression error " + std::to_string(error) + " deviated significantly from target error " + std::to_string(target_error) + "!");
	return error;
}

/// Takes the results from a compress/decompress/compress_decompress operation and verifies all fields accordingly
template<typename T>
void verify_results(const atc::Results& results, const atc::Options& options, bool file_io, bool compress, bool decompress, const std::vector<char>& compressed_data, const std::vector<T>& original_data = std::vector<T>(), const std::vector<T>& decompressed_data = std::vector<T>(), double target_error = 0, double error_margin = 0.05) {
	
	// Check compression results
	if (compress) {
		// compress/compress_decompress
		if (file_io && results.compressed_data.size() > 0)
			throw std::runtime_error("Compression returned non-empty compressed_data while using file IO!");
		size_t actual_compressed_size = compressed_data.size();
		if (results.compressed_size != actual_compressed_size)
			throw std::runtime_error("Reported compressed size (" + std::to_string(results.compressed_size)
				+ ") did not match actual compressed size (" + std::to_string(actual_compressed_size) + ")");
		double actual_compression_factor = double(original_data.size()*sizeof(T))/actual_compressed_size;
		if (!approx_equal(actual_compression_factor, results.compression_factor))
			throw std::runtime_error("Reported compression factor (" + std::to_string(results.compression_factor)
				+ ") did not match actual compression factor (" + std::to_string(actual_compression_factor) + ")");
	} else {
		// decompress
		if (results.compressed_data.size() > 0)
			throw std::runtime_error("Decompression returned non-empty compressed_data!");
		if (results.compressed_size != 0)
			throw std::runtime_error("compressed_size was set during decompression!");
		if (results.compression_factor != 0)
			throw std::runtime_error("compression_factor was set during decompression!");
	}
	
	// Check decompression results
	if (decompress) {
		// decompress/compress_decompress
		size_t decompressed_size = std::accumulate(results.decompressed_sizes.begin(), results.decompressed_sizes.end(), 1, [] (size_t x, size_t y) {return x*y;});
		if (decompressed_size != decompressed_data.size())
			throw std::runtime_error("decompress returned decompressed_sizes whose product does not match actual decompressed size!");
	} else {
		// compress
		if (results.decompressed_sizes.size() > 0)
			throw std::runtime_error("compress returned non-empty decompressed_sizes!");
	}
	
	// Check compress_decompress results
	if (!options.calculate_error && !isnan(results.error))
		throw std::runtime_error("Compression/decompression resulted in non-NaN error while the error should not have been computed!");
	if (decompress && !options.autocrop && std::all_of(options.cutout.begin(), options.cutout.end(), [] (const atc::Slice& slice) {return slice.is_standard();})) {
		double error = verify_target_error(target_error, error_margin, original_data, decompressed_data, options);
		if (compress && options.calculate_error && !approx_equal(error, results.error))
			throw std::runtime_error("compress_decompress returned error of " + std::to_string(results.error) + ", expected " + std::to_string(error) + "!");
	}
	
	if (!options.timings && results.timings.size() > 0)
		throw std::runtime_error("Compression/decompression returned timings while not requested!");
	else if (options.timings && results.timings.size() == 0)
		throw std::runtime_error("Compression/decompression didn't return timings while requested!");
	
}

/// Verifies if the decompressed sizes stored in results are equal to the expected sizes, and if not throws a descriptive std::runtime_error
void verify_decompressed_sizes(const std::vector<size_t>& expected, const atc::Results& results) {
	if (results.decompressed_sizes != expected) {
		std::stringstream ss;
		ss << "decompressed_sizes did not equal expected sizes! Expected: [";
		for (const auto& size : expected)
			ss << size << ", ";
		ss << "], actual: [";
		for (const auto& size : results.decompressed_sizes)
			ss << size << ", ";
		ss << "]";
		throw std::runtime_error(ss.str());
	}
}

/**
 * Tests compression using file IO and given options.
 * Also verifies the equivalence of compress + decompress versus compress_decompress.
 */
template<typename T>
void test_compress_decompress_files(const Dataset<T>& dataset, const atc::Options& options = atc::Options(), double target_error = 80) {
	
	// Set-up compress and decompress
	atc::Results compress_results = atc::compress(input_data_path + dataset.data_file, get_data_type(dataset), dataset.sizes, target_error, compression_path, options);
	atc::Results decompress_results = atc::decompress(compression_path, decompression_path, options);
	std::vector<T> original_data = load_data<T>(input_data_path + dataset.data_file);
	std::vector<char> compressed_data1 = load_data<char>(compression_path);
	std::vector<T> decompressed_data1 = load_data<T>(decompression_path);
	verify_results(compress_results, options, true, true, false, compressed_data1, original_data);
	verify_results(decompress_results, options, true, false, true, compressed_data1, original_data, decompressed_data1, target_error);
	
	// Set-up compress_decompress
	atc::Options temp_options = options;
	temp_options.calculate_error = true;
	atc::Results compress_decompress_results = atc::compress_decompress(input_data_path + dataset.data_file, get_data_type(dataset), dataset.sizes, target_error, compression_path, decompression_path, temp_options);
	std::vector<char> compressed_data2 = load_data<char>(compression_path);
	std::vector<T> decompressed_data2 = load_data<T>(decompression_path);
	if (options.verbose)
		std::cout << compress_decompress_results.error << std::endl;
	verify_results(compress_decompress_results, temp_options, true, true, true, compressed_data2, original_data, decompressed_data2, target_error);
	
}

/**
 * Tests compression using memory IO and given options.
 * Also verifies the equivalence of compress + decompress versus compress_decompress.
 */
template<typename T>
void test_compress_decompress_memory(const Dataset<T>& dataset, const atc::Options& options = atc::Options(), double target_error = 80) {
	
	// Set-up compress and decompress
	std::vector<T> original_data = load_data<T>(input_data_path + dataset.data_file);
	atc::Results compress_results = atc::compress(original_data.data(), dataset.sizes, target_error, options);
	atc::TypedResults<T> decompress_results = atc::decompress<T>(compress_results.compressed_data, options);
	verify_results(compress_results, options, false, true, false, compress_results.compressed_data, original_data);
	verify_results(decompress_results, options, false, false, true, compress_results.compressed_data, original_data, decompress_results.decompressed_data, target_error);
	
	// Set-up compress_decompress
	atc::Options temp_options = options;
	temp_options.calculate_error = true;
	atc::TypedResults<T> compress_decompress_results = atc::compress_decompress(original_data.data(), dataset.sizes, target_error, temp_options);
	verify_results(compress_decompress_results, temp_options, false, true, true, compress_decompress_results.compressed_data, original_data, compress_decompress_results.decompressed_data, target_error);
	
}

/**
 * Tests decompression with a non-standard cut-out
 * Verifies that the produced cut-out is (almost) equal to cut-out obtained by transforming the full decompressed data
 * Verifies that the produced cut-out has a low error compared to the cut-out obtained by transforming the original data
 */
void test_cutout_complex_data(double target_psnr = 80, bool apply_dct = false, double dct_truncation_max_sse_share = 0, double rank_truncation_max_sse_share = 0) {
	
	const Dataset<float>& dataset = complex_data;
	auto original_data = load_data<float>(input_data_path + dataset.data_file);
	
	atc::Options options;
	options.apply_dct = apply_dct;
	options.dct_truncation_max_sse_share = dct_truncation_max_sse_share;
	options.rank_truncation_max_sse_share = rank_truncation_max_sse_share;
	auto compress_results = atc::compress(original_data.data(), dataset.sizes, target_psnr, options);
	verify_results(compress_results, options, false, true, false, compress_results.compressed_data, original_data);
	auto decompress_results = atc::decompress<float>(compress_results.compressed_data);
	verify_results(decompress_results, atc::Options(), false, false, true, compress_results.compressed_data, original_data, decompress_results.decompressed_data, target_psnr);
	
	options.cutout = {{":"}, {"-1"}, {"-3::-4b"}, {"5:16:5l"}};
	auto decompress_results_cutout = atc::decompress<float>(compress_results.compressed_data, options);
	// Update slices with mode sizes, necessary for next verify_results
	for (size_t i = 0; i < dataset.sizes.size(); i++)
		options.cutout[i].update_mode_size(dataset.sizes[i]);
	verify_results(decompress_results_cutout, options, false, false, true, compress_results.compressed_data, std::vector<float>(), decompress_results_cutout.decompressed_data);
	verify_decompressed_sizes(std::vector<size_t>{4, 1, 4, 3}, decompress_results_cutout);
	
	std::vector<size_t> strides(dataset.sizes.size(), 1);
	strides[2] = dataset.sizes[3]*strides[3];
	strides[1] = dataset.sizes[2]*strides[2];
	strides[0] = dataset.sizes[1]*strides[1];
	auto get_index = [&] (size_t i, size_t j, size_t k, size_t l) {return i*strides[0] + j*strides[1] + k*strides[2] + l*strides[3];};
	size_t cutout_size = std::accumulate(options.cutout.begin(), options.cutout.end(), 1, [] (size_t acc, const atc::Slice& slice) {return acc*slice.get_size();});
	if (decompress_results_cutout.decompressed_data.size() != cutout_size)
		throw std::runtime_error("Expected cut-out of size 4x1x4x3 = " + std::to_string(cutout_size)
			+ ", but received decompressed data array of size " + std::to_string(decompress_results_cutout.decompressed_data.size()));
	std::vector<float> decompressed_data_cutout(cutout_size, 0);
	size_t index = 0;
	for (size_t i = 0; i < dataset.sizes[0]; i++) {
		for (size_t j = dataset.sizes[1] - 1; j < dataset.sizes[1]; j++) {
			for (int64_t k1 = dataset.sizes[2] - 3; k1 >= int64_t(0); k1 -= 4) {
				for (size_t l1 = 5; l1 < 16; l1 += 5) {
					
					// Accumulate weighted average for this cutout cell based on full results
					double total_weight = 0;
					float decompressed_val = 0;
					for (int64_t k2 = std::max(0L, k1 - 2); k2 < std::min<int64_t>(dataset.sizes[2], k1 + 2); k2++) {
						double k_weight = 1;
						for (size_t l2 = std::max(0L, int64_t(l1) - 2*5); l2 < std::min(dataset.sizes[3], l1 + 2*5 + 1); l2++) {
							double x = double(int64_t(l2) - int64_t(l1))/5;
							double l_weight = (x == 0 ? 1 : 10*sin(M_PI*x)*sin(M_PI*x/10)/(M_PI*M_PI*x*x));
							
							// Combine weights to compute weighted average term
							double weight = k_weight*l_weight;
							decompressed_val += weight*decompress_results.decompressed_data[get_index(i, j, k2, l2)];
							total_weight += weight;
						}
					}
					decompressed_val /= total_weight;
					
					// Check if decompressed cut-out matches cut-out of decompressed_data
					float cutout_val = decompress_results_cutout.decompressed_data[index];
					if (!approx_equal(decompressed_val, cutout_val, 1e-6f)) {
						std::stringstream ss;
						ss << "Decompressed cut-out did not match cut-out of full decompressed data at position ("
							<< i << ", " << j << ", " << k1 << ", " << l1
							<< "): expected " << decompressed_val
							<< ", actual " << cutout_val
						<< std::endl;
						throw std::runtime_error(ss.str());
					}
					
					decompressed_data_cutout[index] = decompressed_val;
					index++;
					
				}
			}
		}
	}
	
	if (index != cutout_size)
		throw std::runtime_error("Not all values in the cut-out were verified, buggy test!");
	
}

template<typename T>
void test_LSI_storage_order(const Dataset<T>& dataset, const atc::Options& options = atc::Options()) {
	Dataset<T> mut_dataset = dataset;
	atc::Options mut_options = options;
	mut_options.storage_order = atc::StorageOrder::LSI;
	test_compress_decompress_memory(mut_dataset, mut_options);
}

/// Verifies that the mode storage order is interpreted properly by making sure it improves the compression factor
void test_storage_order_compression_effect() {
	
	// First extract a 3x7x13x11 cut-out from complex_data to make sure LSI/MSI ordering makes a significant difference and causes lots of discontinuities
	const Dataset<float>& dataset = complex_data;
	atc::Options cutout_options;
	cutout_options.error_metric = atc::ErrorMetric::RelativeError;
	cutout_options.cutout = {{":3"}, {"-7:"}, {"1:14"}, {":21:2"}};
	cutout_options.suppress_warnings = true;
	auto cutout_results = atc::compress_decompress<float>(input_data_path + dataset.data_file, dataset.sizes, 0, cutout_options);
	verify_decompressed_sizes(std::vector<size_t>{3, 7, 13, 11}, cutout_results);
	
	// Then try MSI and LSI ordering for compressing this cut-out
	double target_error = 80;
	atc::Options options;
	options.storage_order = atc::StorageOrder::MSI;
	auto msi_results = atc::compress(cutout_results.decompressed_data, cutout_results.decompressed_sizes, target_error, options);
	options.storage_order = atc::StorageOrder::LSI;
	auto lsi_results = atc::compress(cutout_results.decompressed_data, cutout_results.decompressed_sizes, target_error, options);
	
	if (lsi_results.compression_factor >= 0.95*msi_results.compression_factor)
		throw std::runtime_error("Compression factor did not degrade noticeably by erroneously using the LSI storage order: MSI compression factor: "
			+ std::to_string(msi_results.compression_factor) + ", LSI compression factor: " + std::to_string(lsi_results.compression_factor)
		);
	
}

/// Verifies that 32-bit precision floating-point works but only for high errors
void test_ftypes() {
	
	const Dataset<float>& dataset = complex_data;
	atc::Options options;
	options.error_metric = atc::ErrorMetric::RelativeError;
	options.calculate_error = true;
	
	// Test with high error
	
	// Test with 32-bit precision
	double target_relative_error = 1e-1;
	options.float_type = atc::FloatType::float32;
	auto results = atc::compress_decompress<float>(input_data_path + dataset.data_file, dataset.sizes, target_relative_error, options);
	double float32_compression_factor = results.compression_factor;
	if (abs(results.error - target_relative_error)/target_relative_error >= 0.05) {
		throw std::runtime_error("Compression error with 32-bit ftype deviated too much from high target relative error "
			+ std::to_string(target_relative_error) + " (" + std::to_string(results.error) + ")"
		);
	}
	
	// Test with 64-bit precision
	options.float_type = atc::FloatType::float64;
	results = atc::compress_decompress<float>(input_data_path + dataset.data_file, dataset.sizes, target_relative_error, options);
	if (abs(results.error - target_relative_error)/target_relative_error >= 0.05) {
		throw std::runtime_error("Compression error with 64-bit ftype deviated too much from low target relative error "
			+ std::to_string(target_relative_error) + " (" + std::to_string(results.error) + ")"
		);
	}
	if (abs(float32_compression_factor - results.compression_factor)/results.compression_factor >= 0.01) {
		throw std::runtime_error("Compression factor with 32-bit ftype (" + std::to_string(float32_compression_factor)
			+ ") deviated too much from 64-bit compression factor (" + std::to_string(results.compression_factor) + ")"
		);
	}
	
	// Test with low error
	
	// Test with 32-bit precision
	target_relative_error = 1e-6;
	options.float_type = atc::FloatType::float32;
	results = atc::compress_decompress<float>(input_data_path + dataset.data_file, dataset.sizes, target_relative_error, options);
	if (abs(results.error - target_relative_error)/target_relative_error <= 0.05) {
		throw std::runtime_error("Compression error with 32-bit ftype did not deviate much from low target relative error "
			+ std::to_string(target_relative_error) + " (" + std::to_string(results.error) + ")"
		);
	}
	
	// Test with 64-bit precision
	options.float_type = atc::FloatType::float64;
	results = atc::compress_decompress<float>(input_data_path + dataset.data_file, dataset.sizes, target_relative_error, options);
	if (abs(results.error - target_relative_error)/target_relative_error >= 0.05) {
		throw std::runtime_error("Compression error with 64-bit ftype deviated too much from low target relative error "
			+ std::to_string(target_relative_error) + " (" + std::to_string(results.error) + ")"
		);
	}
	
}

/// Verifies that data files can start with arbitrary headers which can be skipped using the skip_bytes parameter
void test_skip_bytes() {
	
	// Reference compression
	const Dataset<float>& dataset = complex_data;
	auto original_data = load_data<float>(input_data_path + dataset.data_file);
	double target_psnr = 80;
	atc::Options options;
	auto ref_results = atc::compress_decompress(original_data.data(), dataset.sizes, target_psnr);
	
	// Create padded data file
	size_t padding = 1358; // arbitrary value
	std::vector<char> padded_data(padding + original_data.size()*sizeof(float), 68 /* arbitrary value */);
	std::copy(original_data.begin(), original_data.end(), reinterpret_cast<float*>(&padded_data[padding]));
	write_data(padded_data, other_path);
	
	// Compress padded file
	options.skip_bytes = padding;
	auto results = atc::compress_decompress<float>(other_path, dataset.sizes, target_psnr, options);
	
	if (double(abs(int64_t(ref_results.compressed_size) - int64_t(results.compressed_size)))/ref_results.compressed_size >= 0.01) {
		throw std::runtime_error("Compressed size for unpadded data file (" + std::to_string(ref_results.compressed_size)
			+ ") deviated too much from compressed size for padded data file (" + std::to_string(results.compressed_size) + ")"
		);
	}
	
	try {
		options.skip_bytes = padding - 1;
		auto results = atc::compress_decompress<float>(other_path, dataset.sizes, target_psnr, options);
		throw std::runtime_error("ATC did not throw exception despite passing invalid skip_bytes parameter!");
	} catch (const std::invalid_argument& e) {
		// ATC detected wrong padding argument, everything is fine
	}
	
}

/// Verifies that all error metrics can be used with reasonable error control
void test_error_metrics() {
	
	const Dataset<float>& dataset = complex_data;
	atc::Options options;
	options.error_metric = atc::ErrorMetric::PSNR;
	test_compress_decompress_files(dataset, options, 100);
	test_compress_decompress_memory(dataset, options, 100);
	options.error_metric = atc::ErrorMetric::RelativeError;
	test_compress_decompress_files(dataset, options, 1e-5);
	test_compress_decompress_memory(dataset, options, 1e-5);
	
	// For RMSE, compute the data norm to find a reasonable RMSE
	options.error_metric = atc::ErrorMetric::RMSE;
	auto original_data = load_data<float>(input_data_path + dataset.data_file);
	std::vector<float> zero_data(original_data.size(), 0);
	double data_norm_sq = sse(original_data, zero_data);
	double rmsv = sqrt(data_norm_sq/original_data.size());
	double rmse = 0.01*rmsv;
	test_compress_decompress_files(dataset, options, rmse);
	test_compress_decompress_memory(dataset, options, rmse);
	
}

/// Verifies that we can pass an arbitrary permutation as the ST-HOSVD compression mode processing order and that invalid permutations will be detected
void test_mode_processing_order() {
	
	const Dataset<float>& dataset = complex_data;
	atc::Options options;
	options.mode_processing_order = {1, 3, 2, 0};
	test_compress_decompress_files(dataset, options, 100);
	test_compress_decompress_memory(dataset, options, 100);
	
	try {
		options.mode_processing_order = {1, 3, 1, 0};
		test_compress_decompress_files(dataset, options, 100);
		throw std::runtime_error("ATC did not throw exception despite passing invalid mode_processing_order parameter (invalid permutation)!");
	} catch (const std::invalid_argument& e) {
		// ATC detected wrong mode processing order, everything is fine
	}
	
	try {
		options.mode_processing_order = {3, 2};
		test_compress_decompress_files(dataset, options, 100);
		throw std::runtime_error("ATC did not throw exception despite passing invalid mode_processing_order parameter (too short)!");
	} catch (const std::invalid_argument& e) {
		// ATC detected wrong mode processing order, everything is fine
	}
	
}

/// Verifies that we can pass an arbitrary permutation as the mode storage order and that invalid permutations will be detected
void test_mode_storage_order() {
	
	const Dataset<uint8_t>& dataset = basic_data;
	atc::Options options;
	options.mode_storage_order = {2, 4, 3, 0, 1};
	test_compress_decompress_files(dataset, options, 100);
	test_compress_decompress_memory(dataset, options, 100);
	
	try {
		options.mode_storage_order = {0, 2, 3, 4, 2};
		test_compress_decompress_files(dataset, options, 100);
		throw std::runtime_error("ATC did not throw exception despite passing invalid mode_storage_order parameter (invalid permutation)!");
	} catch (const std::invalid_argument& e) {
		// ATC detected wrong mode processing order, everything is fine
	}
	
	try {
		options.mode_storage_order = {0};
		test_compress_decompress_files(dataset, options, 100);
		throw std::runtime_error("ATC did not throw exception despite passing invalid mode_storage_order parameter (too short)!");
	} catch (const std::invalid_argument& e) {
		// ATC detected wrong mode processing order, everything is fine
	}
	
}

/// Verifies that 32-bit precision floating-point works but only for medium errors
void test_itypes() {
	
	const Dataset<float>& dataset = complex_data;
	atc::Options options;
	options.error_metric = atc::ErrorMetric::RelativeError;
	options.calculate_error = true;
	
	// Test with high error
	double target_relative_error = 1e-4;
	
	// Test with 32-bit precision
	options.int_type = atc::IntType::uint32;
	auto results = atc::compress_decompress<float>(input_data_path + dataset.data_file, dataset.sizes, target_relative_error, options);
	double uint32_compression_factor = results.compression_factor;
	if (abs(results.error - target_relative_error)/target_relative_error >= 0.05) {
		throw std::runtime_error("Compression error with 32-bit itype deviated too much from high target relative error "
			+ std::to_string(target_relative_error) + " (" + std::to_string(results.error) + ")"
		);
	}
	
	// Test with 64-bit precision
	options.int_type = atc::IntType::uint64;
	results = atc::compress_decompress<float>(input_data_path + dataset.data_file, dataset.sizes, target_relative_error, options);
	if (abs(results.error - target_relative_error)/target_relative_error >= 0.05) {
		throw std::runtime_error("Compression error with 64-bit itype deviated too much from low target relative error "
			+ std::to_string(target_relative_error) + " (" + std::to_string(results.error) + ")"
		);
	}
	if (abs(uint32_compression_factor - results.compression_factor)/results.compression_factor >= 0.01) {
		throw std::runtime_error("Compression factor with 32-bit itype (" + std::to_string(uint32_compression_factor)
			+ ") deviated too much from 64-bit compression factor (" + std::to_string(results.compression_factor) + ")"
		);
	}
	
	// Test with low error
	target_relative_error = 1e-7;
	
	// Test with 32-bit precision: irrelevant, for this dataset we can't go low enough with our target error to the point where error control differs
	
	// Test with 64-bit precision
	options.int_type = atc::IntType::uint64;
	results = atc::compress_decompress<float>(input_data_path + dataset.data_file, dataset.sizes, target_relative_error, options);
	if (abs(results.error - target_relative_error)/target_relative_error >= 0.05) {
		throw std::runtime_error("Compression error with 64-bit itype deviated too much from low target relative error "
			+ std::to_string(target_relative_error) + " (" + std::to_string(results.error) + ")"
		);
	}
	
}

/// Verify that compression works fine even if one modes is of rank 1 (rank truncation will then truncate the mode size to 1)
void test_rank_1_mode() {
	
	std::vector<uint8_t> test_data(27, 0);
	std::iota(test_data.begin(), test_data.begin() + 9, 0);
	std::iota(test_data.begin() + 9, test_data.begin() + 18, 0);
	std::iota(test_data.begin() + 18, test_data.begin() + 27, 0);
	
	double target_psnr = 80;
	auto res1 = atc::compress_decompress(test_data.data(), {3, 3, 3}, target_psnr);
	verify_results(res1, atc::Options(), false, true, true, res1.compressed_data, test_data, res1.decompressed_data, target_psnr);
	
	atc::Options options;
	options.apply_dct = true;
	options.dct_truncation_max_sse_share = 0.1;
	auto res2 = atc::compress_decompress(test_data.data(), {3, 3, 3}, target_psnr);
	verify_results(res2, atc::Options(), false, true, true, res2.compressed_data, test_data, res2.decompressed_data, target_psnr);
	
}

/// Verify that autocrop works properly with different thresholds
void test_autocrop(bool apply_dct = false) {
	
	// We generate a dataset modelling two 3D diamonds, one wrapped inside the other
	std::vector<size_t> sizes = {100, 90, 80};
	std::vector<double> test_data(std::accumulate(sizes.begin(), sizes.end(), 1, [] (size_t a, size_t b) {return a*b;}));
	std::vector<int64_t> center = {34, 56, 49};
	size_t r1 = 15;
	size_t r2 = 25;
	size_t index = 0;
	for (size_t x = 0; x < sizes[0]; x++) {
		for (size_t y = 0; y < sizes[1]; y++) {
			for (size_t z = 0; z < sizes[2]; z++) {
				size_t distance = abs(int64_t(center[0]) - int64_t(x))
					+ abs(int64_t(center[1]) - int64_t(y))
					+ abs(int64_t(center[2]) - int64_t(z));
				test_data[index] = (2*((x + y + z)%2) - 1)*(distance <= r1 ? 1 : (distance <= r2 ? 1e-2 : 0));
				index++;
			}
		}
	}
	
	// Apply compression, this is independent from the autocrop parameter
	double target_psnr = 100;
	atc::Options options;
	options.apply_dct = apply_dct;
	options.dct_truncation_max_sse_share = 0.5; // DCT rank truncation doesn't do much here because of the high discontinuities, DCT base doesn't work well
	auto compress_results = atc::compress(test_data.data(), sizes, target_psnr, options);
	verify_results(compress_results, options, false, true, false, compress_results.compressed_data, test_data);
	
	// Apply decompression without autocrop
	options.autocrop = false;
	auto decompress_results = atc::decompress<double>(compress_results.compressed_data, options);
	// Because this dataset is intentionally very discontinuous, ATC fails to compress it with precise error control, so we relax the error margin in this case
	verify_results(decompress_results, options, false, false, true, compress_results.compressed_data, test_data, decompress_results.decompressed_data, target_psnr, 0.4);
	verify_decompressed_sizes(sizes, decompress_results);
	
	// Apply decompression while cropping very low-energy slices
	options.autocrop = true;
	options.autocrop_threshold = 1e-4;
	decompress_results = atc::decompress<double>(compress_results.compressed_data, options);
	verify_results(decompress_results, options, false, false, true, compress_results.compressed_data, test_data, decompress_results.decompressed_data);
	verify_decompressed_sizes(std::vector<size_t>(3, 2*r2 + 1), decompress_results);
	
	// Apply decompression while cropping low-energy slices
	options.autocrop = true;
	options.autocrop_threshold = 1e-2;
	decompress_results = atc::decompress<double>(compress_results.compressed_data, options);
	verify_results(decompress_results, options, false, false, true, compress_results.compressed_data, test_data, decompress_results.decompressed_data);
	verify_decompressed_sizes(std::vector<size_t>(3, 2*r1 + 1), decompress_results);
	
}

/// Performs all tests and reports success if all tests succeeded.
int main() {
	
	std::vector<std::pair<std::string, std::function<void ()>>> tests = {
		{"test_compress_decompress_files_basic", [] () {test_compress_decompress_files(basic_data);}},
		{"test_compress_decompress_files_complex", [] () {test_compress_decompress_files(complex_data);}},
		{"test_compress_decompress_memory_basic", [] () {test_compress_decompress_memory(basic_data);}},
		{"test_compress_decompress_memory_complex", [] () {test_compress_decompress_memory(complex_data);}},
		{"test_cutout_complex_data", [] () {test_cutout_complex_data(80);}},
		#ifdef ATC_FFTW
			{"test_cutout_complex_data_dct_no_truncation", [] () {test_cutout_complex_data(80, true, 0);}},
			{"test_cutout_complex_data_dct_with_truncation", [] () {test_cutout_complex_data(50, true, 1, 1);}},
		#endif
		{"test_timings", [] () {atc::Options options; options.timings = true; test_compress_decompress_memory(complex_data, options);}},
		{"test_0_threads", [] () {atc::Options options; options.threads = 0; test_compress_decompress_files(basic_data, options);}},
		{"test_1_threads", [] () {atc::Options options; options.threads = 1; test_compress_decompress_memory(complex_data, options);}},
		{"test_4_threads", [] () {atc::Options options; options.threads = 4; test_compress_decompress_files(complex_data, options);}},
		{"test_LSI_storage_order_basic_data", [] () {test_LSI_storage_order(basic_data);}},
		{"test_LSI_storage_order_complex_data", [] () {test_LSI_storage_order(complex_data);}},
		{"test_storage_order_compression_effect", test_storage_order_compression_effect},
		{"test_ftypes", test_ftypes},
		{"test_EPCM_basic", [] () {atc::Options options; options.explicitly_parallellized_core_multiplications = true; test_compress_decompress_files(basic_data, options);}},
		{"test_EPCM_complex", [] () {atc::Options options; options.explicitly_parallellized_core_multiplications = true; test_compress_decompress_files(complex_data, options);}},
		{"test_skip_bytes", test_skip_bytes},
		{"test_error_metrics", test_error_metrics},
		#ifdef ATC_FFTW
			{"test_dct", [] () {atc::Options options; options.apply_dct = true; test_compress_decompress_memory(complex_data, options);}},
			{"test_dct_high_truncation", [] () {atc::Options options; options.apply_dct = true; options.dct_truncation_max_sse_share = 0.8; test_compress_decompress_files(basic_data, options);}},
		#endif
		{"test_no_rank_truncation", [] () {atc::Options options; options.rank_truncation_max_sse_share = 0; test_compress_decompress_files(complex_data, options);}},
		{"test_mode_processing_order", test_mode_processing_order},
		{"test_mode_storage_order", test_mode_storage_order},
		{"test_precise_svd", [] () {atc::Options options; options.precise_svd = true; test_compress_decompress_memory(complex_data, options);}},
		{"test_no_householder", [] () {atc::Options options; options.householder = false; test_compress_decompress_memory(complex_data, options);}},
		{"test_itypes", test_itypes},
		{"test_flattenings_unfold", [] () {atc::Options options; options.core_flattening = atc::CoreFlattening::Unfold; test_compress_decompress_memory(complex_data, options);}},
		{"test_flattenings_zigzag", [] () {atc::Options options; options.core_flattening = atc::CoreFlattening::Zigzag; test_compress_decompress_memory(basic_data, options);}},
		{"test_flattenings_zorder", [] () {atc::Options options; options.core_flattening = atc::CoreFlattening::ZOrder; test_compress_decompress_files(complex_data, options);}},
		{"test_rank_1_mode", test_rank_1_mode},
		{"test_0_qb", [] () {atc::Options options; options.quantization_blocks = 0; test_compress_decompress_files(complex_data, options);}},
		{"test_1_qb", [] () {atc::Options options; options.quantization_blocks = 1; test_compress_decompress_memory(basic_data, options);}},
		{"test_4_qb", [] () {atc::Options options; options.quantization_blocks = 4; test_compress_decompress_files(basic_data, options);}},
		{"test_100000_qb", [] () {atc::Options options; options.quantization_blocks = 100000; test_compress_decompress_memory(complex_data, options);}},
		{"test_1_lpmt", [] () {atc::Options options; options.last_pass_max_threads = 1; test_compress_decompress_memory(complex_data, options);}},
		{"test_4_lpmt", [] () {atc::Options options; options.last_pass_max_threads = 4; test_compress_decompress_memory(complex_data, options);}},
		{"test_20_lpmt", [] () {atc::Options options; options.last_pass_max_threads = 20; test_compress_decompress_memory(complex_data, options);}},
		{"test_0.01_qrem", [] () {atc::Options options; options.quantization_rel_error_margin = 0.01; test_compress_decompress_memory(complex_data, options);}},
		{"test_0.05_qrem", [] () {atc::Options options; options.quantization_rel_error_margin = 0.05; test_compress_decompress_memory(complex_data, options);}},
		{"test_0.0001_qrem", [] () {atc::Options options; options.quantization_rel_error_margin = 0.0001; test_compress_decompress_memory(complex_data, options);}},
		{"test_no_split_bit_plane_truncation", [] () {atc::Options options; options.split_bitplane_truncation = false; test_compress_decompress_memory(complex_data, options);}},
		{"test_entropy_coding_ac", [] () {atc::Options options; options.entropy_coding_method = atc::EntropyCodingMethod::Arithmetic; test_compress_decompress_memory(complex_data, options);}},
		{"test_entropy_coding_ans", [] () {atc::Options options; options.entropy_coding_method = atc::EntropyCodingMethod::ANS; test_compress_decompress_memory(complex_data, options);}},
		{"test_autocrop", [] () {test_autocrop(false);}},
		#ifdef ATC_FFTW
			{"test_autocrop_dct", [] () {test_autocrop(true);}},
		#endif
	};
	
	for (const auto& [test_name, test] : tests) {
		std::cout << "Performing " << test_name << " ... " << std::flush;
		try {
			test(); // Throws an exception in case of failure
			std::cout << "Success!" << std::endl;
		} catch (const std::exception& e) {
			std::cout << "Failed!" << std::endl;
			std::cout << "Reason: " << e.what() << std::endl;
			return 1;
		}
	}
	std::cout << "All test succeeded!" << std::endl;
	return 0;
	
}
