/*
 * Experiments to generate data for figures, tables, ... in the ATC paper
 */

#include <vector>
#include <tuple>
#include <fstream>
#include <streambuf>
#include <string>
#include <iostream>
#include <numeric>
#include <algorithm>
#include <cassert>
#include <cmath>
#include <sstream>
#include <experimental/filesystem>

#include "atc.hpp"

// PRELIMINARIES

const std::string temp_comp_path = "/tmp/comp.atc";
const std::string temp_decomp_path = "/tmp/decomp";
const std::string temp_stdout_path = "/tmp/atc-stdout";
const std::string input_data_path = "../../data/";
const std::string output_data_path = "../../../Manuscripts/atc-paper/data/";

struct Dataset {
	const std::string name;
	const std::string data_file;
	const atc::DataType data_type;
	const std::vector<size_t> sizes;
};

const Dataset turbulence_pressure_time_64 = {
	"Isotropic1024-fine-pressure-time-64",
	"turbulence/preprocessed/isotropic1024fine-pressure-time-64",
	atc::DataType::float32,
	std::vector<size_t>{64, 64, 64, 64}
};

const Dataset turbulence_pressure_time_128 = {
	"Isotropic-PT",
	"turbulence/preprocessed/isotropic1024fine-pressure-time-128",
	atc::DataType::float32,
	std::vector<size_t>{100, 128, 128, 128}
};

const Dataset moffett_field = {
	"Moffett-Field",
	"hyperspectral_images/preprocessed/moffett_field",
	atc::DataType::int16,
	std::vector<size_t>{512, 1024, 224}
};

const Dataset foot = {
	"Foot",
	"volumes/raw/foot",
	atc::DataType::uint8,
	std::vector<size_t>{256, 256, 256}
};

// Redirects cout to file as long as the lifetime of this object
class RedirectCout {
public:
	RedirectCout() : console_cout(std::cout.rdbuf()) {
		m_ss.str("");
		std::cout.rdbuf(m_ss.rdbuf());
	}
	RedirectCout(const std::string& file) : stream(file), console_cout(std::cout.rdbuf()) {
		std::cout.rdbuf(stream.rdbuf());
	}
	// Return the stringstream buffer and reset the local state
	std::stringstream ss() {
		std::stringstream res;
		std::swap(res, m_ss);
		return res;
	}
	~RedirectCout() {
		std::cout.rdbuf(console_cout);
	}
private:
	std::stringstream m_ss;
	std::ofstream stream;
	std::streambuf* console_cout;
};

void compress(const Dataset& dataset, double target_rel_error, const std::string& compressed_file, atc::Options options = atc::Options()) {
	options.error_metric = atc::ErrorMetric::RelativeError;
	atc::compress(input_data_path + dataset.data_file, dataset.data_type, dataset.sizes, target_rel_error, compressed_file, options);
}

void compress(const Dataset& dataset, double target_rel_error, const atc::Options options = atc::Options()) {
	compress(dataset, target_rel_error, temp_comp_path, options);
}

atc::Results compress_decompress(const Dataset& dataset, double target_rel_error, atc::Options options = atc::Options(), bool no_calculate_error = false) {
	options.error_metric = atc::ErrorMetric::RelativeError;
	options.calculate_error = !no_calculate_error;
	return atc::compress_decompress(input_data_path + dataset.data_file, dataset.data_type, dataset.sizes, target_rel_error, temp_comp_path, temp_decomp_path, options);
}

std::vector<double> logspace(double start, double end, size_t amount) {
	assert(amount > 1);
	std::vector<double> res(amount);
	for (size_t i = 0; i < amount - 1; i++)
		res[i] = start*std::pow(end/start, double(i)/double(amount - 1));
	res[amount - 1] = end; // We set this explicitly to avoid round-off errors
	return res;
}

// EXPERIMENTS

// Algorithmic improvements

// Hybrid truncation

void check_hybrid_truncation_error_approximation() {
	
	// Settings
	std::vector<Dataset> datasets = {turbulence_pressure_time_128, moffett_field, foot};
	std::vector<double> target_rel_errors = {0.1, 0.01, 0.001};
	std::vector<double> rtmsses = {0, 0.25, 0.5, 0.75, 1};
	
	// Other initialization
	atc::Options options;
	options.test_error_approximation = true;
	options.rank_truncation = true;
	options.threads = 4;
	options.core_flattening = atc::CoreFlattening::Unfold;
	options.verbose = true; // To make sure we see the output
	std::string output_file = output_data_path + "hybrid_truncation_error_approximation_rel_errors.csv";
	std::ofstream out(output_file + ".temp");
	
	// Print header
	out << "Dataset,RTMSS,Target relative error" << std::endl;
	out << ",";
	for (auto target_rel_error : target_rel_errors)
		out << "," << target_rel_error;
	out << std::endl;
	
	for (const auto& dataset : datasets) {
		std::cout << "Handling dataset: " << dataset.name << std::endl;
		for (const auto& rtmss : rtmsses) {
			options.rank_truncation_max_sse_share = rtmss;
			std::cout << "Handling RTMSS: " << options.rank_truncation_max_sse_share << std::endl;
				
			// Print start of entry
			out << dataset.name << "," << options.rank_truncation_max_sse_share;
			for (auto target_rel_error : target_rel_errors) {
				std::cout << "For target relative error: " << target_rel_error << std::endl;
				
				// Perform compression and capture output
				std::stringstream ss;
				{
					RedirectCout redirect;
					atc::Results res = compress_decompress(dataset, target_rel_error, options, true);
					ss = redirect.ss();
				}
				
				// Filter output to extract relative error on SSE approximation
				std::string line;
				std::string search_string = "Relative error of approximation: ";
				double rel_approx_error = -1; // NaN takes extra includes etcetera ...
				while (std::getline(ss, line)) {
					size_t pos = line.find(search_string);
					if (pos != std::string::npos) {
						rel_approx_error = std::stod(line.substr(pos + search_string.length()));
						break;
					}
				}
				
				out << "," << rel_approx_error << std::flush;
				
			}
			out << std::endl;
			
		}
	}
	
	// Finish and move file
	out.close();
	std::rename((output_file + ".temp").c_str(), output_file.c_str());
	
}

// This experiment can take a few hours
void dump_rtmss_rate_distortion() {
	atc::Options options;
	options.rank_truncation = true;
	//options.apply_dct = false;
	//options.dct_truncation_max_sse_share = 0;
	options.threads = 4;
	options.core_flattening = atc::CoreFlattening::Unfold;
	std::string output_file = output_data_path + "rtmss_rate_distortion.csv";
	std::ofstream out(output_file + ".temp");
	std::vector<double> target_rel_errors = logspace(1e-1, 1e-3, 17);
	
	// Print header
	out << "Dataset,RTMSS,Target relative error" << std::endl;
	out << ",";
	for (auto target_rel_error : target_rel_errors)
		out << "," << target_rel_error << ",";
	out << std::endl;
	out << ",";
	for (auto target_rel_error [[maybe_unused]] : target_rel_errors)
		out << ",Actual relative error,Compression factor";
	out << std::endl;
	
	for (const auto& dataset : std::vector<Dataset>{turbulence_pressure_time_128, moffett_field, foot}) {
		std::cout << "Handling dataset: " << dataset.name << std::endl;
		for (const auto& rtmss : std::vector<double>{0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1}) {
			options.rank_truncation_max_sse_share = rtmss;
			std::cout << "Handling RTMSS: " << options.rank_truncation_max_sse_share << std::endl;
				
			// Print start of entry
			out << dataset.name << "," << options.rank_truncation_max_sse_share;
			for (auto target_rel_error : target_rel_errors) {
				std::cout << "For target relative error: " << target_rel_error << std::endl;
				
				// Perform compression and print data
				atc::Results res = compress_decompress(dataset, target_rel_error, options);
				out << "," << res.error << "," << res.compression_factor << std::flush;
				
			}
			out << std::endl;
			
		}
	}
	
	// Finish and move file
	out.close();
	std::rename((output_file + ".temp").c_str(), output_file.c_str());
}

// T-HODCT

void write_factors_to_file(const Dataset& dataset, double target_psnr, const std::string& factors_file) {
	atc::Options options;
	options.factors_path = output_data_path + factors_file;
	options.rank_truncation = false;
	options.apply_dct = false;
	options.float_type = atc::FloatType::float64;
	options.verbose = true;
	compress(dataset, target_psnr, options);
}

void experiment_write_factors_to_files() {
	write_factors_to_file(turbulence_pressure_time_128, 80, "isotropic1024fine-pressure-time-128-factor");
	write_factors_to_file(moffett_field, 80, "moffett-field-factor");
}

// Split bitplane truncation

void dump_split_bitplane_truncation_output() {
	atc::Options options;
	options.verbose = true;
	options.split_bitplane_truncation = false;
	options.print_encoding_sse_reduction_bit_cost = true;
	options.rank_truncation_max_sse_share = 0.5;
	options.apply_dct = false;
	RedirectCout redirect(output_data_path + "isotropic1024fine-pressure-time-128-output");
	compress(turbulence_pressure_time_128, 0.001, options);
}

// Core reorderings

// This experiment takes quite long (hours)
void test_core_flattenings() {
	atc::Options options;
	options.rank_truncation_max_sse_share = 0.5;
	options.apply_dct = false;
	options.verbose = true;
	options.split_bitplane_truncation = true;
	options.threads = 4;
	std::string output_file = output_data_path + "core_flattenings.csv";
	std::string core_sizes_base_output_file = output_data_path + "core_flattenings_core_sizes";
	std::ofstream out(output_file + ".temp");
	std::vector<double> target_rel_errors{0.1, 0.03, 0.01, 0.003, 0.001};
	out << "Dataset,Core flattening,Mode storage order,Target relative error" << std::endl;
	out << ",,";
	for (auto target_rel_error : target_rel_errors)
		out << "," << target_rel_error << ",";
	out << std::endl;
	out << ",,";
	for (auto target_rel_error [[maybe_unused]]: target_rel_errors)
		out << ",Actual relative error,Compression factor";
	out << std::endl;
	for (const auto& dataset : std::vector<Dataset>{turbulence_pressure_time_128, moffett_field, foot}) {
		std::cout << "Handling dataset: " << dataset.name << std::endl;
		bool first_flattening = true;
		bool stored_core_sizes = false;
		for (auto flattening : std::vector<atc::CoreFlattening>{atc::CoreFlattening::Unfold, atc::CoreFlattening::Zigzag, atc::CoreFlattening::ZOrder}) {
			
			std::string flattening_string = (flattening == atc::CoreFlattening::Unfold ? "Unfold" : (flattening == atc::CoreFlattening::Zigzag ? "Zigzag" : "Z-order"));
			std::cout << "Handling core flattening: " << flattening_string << std::endl;
			
			if (first_flattening)
				out << dataset.name;
			first_flattening = false;
			bool first_it2 = true;
			
			options.core_flattening = flattening;
			uint8_t modes = dataset.sizes.size();
			options.mode_storage_order.resize(modes);
			std::iota(options.mode_storage_order.begin(), options.mode_storage_order.end(), 0);
			do {
				
				std::stringstream ss;
				ss << "(" << uint64_t(options.mode_storage_order[0]);
				for (uint8_t i = 1; i < modes; i++)
					ss << ", " << uint64_t(options.mode_storage_order[i]);
				ss << ")";
				std::cout << "Handling mode order: " << ss.str() << std::endl;
				
				out << ",";
				if (first_it2)
					out << flattening_string;
				first_it2 = false;
				out << ",\"" << ss.str() << "\"";
				
				std::string core_sizes_output_file = core_sizes_base_output_file + "_" + dataset.name;
				std::ofstream core_sizes_out;
				if (!stored_core_sizes)
					core_sizes_out = std::ofstream(core_sizes_output_file + ".temp");
				for (size_t i = 0; i < target_rel_errors.size(); i++) {
					double target_rel_error = target_rel_errors[i];
					std::cout << "Handling relative error: " << target_rel_error << std::endl;
					{
						
						RedirectCout redirect;
						atc::Results results = compress_decompress(dataset, target_rel_error, options);
						out << "," << results.error << "," << results.compression_factor << std::flush;
						
						if (!stored_core_sizes) {
							// Parse output from stdout
							std::string str(redirect.ss().str());
							std::string keyword = "ST-HOSVD-truncated core sizes: ";
							size_t start = str.find(keyword) + keyword.size();
							size_t end = str.find("\n", start);
							
							// Output extracted information to file
							core_sizes_out << "(" << target_rel_error << ": " << str.substr(start, end - start) << ")";
							if (i < target_rel_errors.size() - 1)
								core_sizes_out << ", ";
							core_sizes_out << std::flush;
						}
						
					}
				}
				if (!stored_core_sizes) {
					core_sizes_out.close();
					std::rename((core_sizes_output_file + ".temp").c_str(), core_sizes_output_file.c_str());
					stored_core_sizes = true;
				}
				
				out << std::endl;
				
			} while (std::next_permutation(options.mode_storage_order.begin(), options.mode_storage_order.end()));
			
		}
	}
	out.close();
	std::rename((output_file + ".temp").c_str(), output_file.c_str());
}

/*void print_truncated_sizes_for_core_flattenings() {
	atc::Options options;
	options.verbose = true;
	options.rank_truncation_max_sse_share = 0.5;
	options.apply_dct = false;
	options.split_bitplane_truncation = true;
	options.threads = 4;
	options.core_flattening = atc::CoreFlattening::Unfold;
	std::vector<double> target_rel_errors{0.1, 0.03, 0.01, 0.003, 0.001};
	for (const auto& dataset : std::vector<Dataset>{turbulence_pressure_time_128, moffett_field, foot}) {
		std::cout << "Handling dataset: " << dataset.name << std::endl;
		
		for (auto target_rel_error : target_rel_errors) {
			
			// Perform compression
			std::stringstream ss;
			{
				RedirectCout redirect;
				compress(dataset, target_rel_error, options);
				ss = redirect.ss();
			}
			
			// Parse output from stdout
			std::string str(ss.str());
			std::string keyword = "ST-HOSVD-truncated core sizes: ";
			size_t start = str.find(keyword) + keyword.size();
			size_t end = str.find("\n", start);
			
			// Print extracted information to stdout
			std::cout << "(" << target_rel_error << ": " << str.substr(start, end - start) << "), ";
			
		}
		
		std::cout << std::endl;
		
	}
}*/

// Unified quantization

/*void dump_unified_quantization_measurements() {
	atc::Options options;
	options.rank_truncation_max_sse_share = 0.5;
	options.apply_dct = false;
	options.split_bitplane_truncation = true;
	options.threads = 4;
	options.core_flattening = atc::CoreFlattening::Unfold;
	std::string output_file = output_data_path + "unified_quantization_measurements.csv";
	std::ofstream out(output_file + ".temp");
	std::vector<double> target_psnrs{
		40, 42.5, 45, 47.5,
		50, 52.5, 55, 57.5,
		60, 62.5, 65, 67.5,
		70, 72.5, 75, 77.5,
		80
	};
	
	// Print header
	out << "Dataset,Unified quantization,Target PSNR" << std::endl;
	out << ",";
	for (auto target_psnr : target_psnrs)
		out << "," << target_psnr << ",";
	out << std::endl;
	out << ",";
	for (auto target_psnr [[maybe_unused]]: target_psnrs)
		out << ",Actual PSNR,Compression factor";
	out << std::endl;
	for (const auto& dataset : std::vector<Dataset>{turbulence_pressure_time_128, moffett_field, foot}) {
		std::cout << "Handling dataset: " << dataset.name << std::endl;
		
		for (auto unified_quantization : std::vector<bool>{false, true}) {
			options.unified_quantization = unified_quantization;
			std::cout << "With unified quantization: " << options.unified_quantization << std::endl;
			
			// Print start of entry
			out << dataset.name << "," << options.unified_quantization;
			for (auto target_psnr : target_psnrs) {
				std::cout << "For target PSNR: " << target_psnr << std::endl;
				
				// Perform compression and print data
				atc::Results res = compress_decompress(dataset, target_psnr, options);
				out << "," << res.error << "," << res.compression_factor << std::flush;
				
			}
			out << std::endl;
		
		}
		
	}
	out.close();
	std::rename((output_file + ".temp").c_str(), output_file.c_str());
}*/

// Implementation

// Type support

void dump_type_precision_measurements() {
	atc::Options options;
	options.rank_truncation_max_sse_share = 0.5;
	options.apply_dct = false;
	options.threads = 4;
	options.core_flattening = atc::CoreFlattening::Unfold;
	std::string output_file = output_data_path + "type_precision_measurements.csv";
	std::ofstream out(output_file + ".temp");
	std::vector<double> target_psnrs{
		40, 42.5, 45, 47.5,
		50, 52.5, 55, 57.5,
		60, 62.5, 65, 67.5,
		70, 72.5, 75, 77.5,
		80
	};
	
	// Print header
	out << "Dataset,Float type,Int type,Target relative error" << std::endl;
	out << ",,";
	for (auto target_psnr : target_psnrs)
		out << "," << target_psnr << ",";
	out << std::endl;
	out << ",,";
	for (auto target_psnr [[maybe_unused]]: target_psnrs)
		out << ",Actual relative error,Compression factor";
	out << std::endl;
	
	for (const auto& dataset : std::vector<Dataset>{turbulence_pressure_time_128, moffett_field, foot}) {
		std::cout << "Handling dataset: " << dataset.name << std::endl;
		for (const auto& float_type : std::vector<atc::FloatType>{atc::FloatType::float32, atc::FloatType::float64}) {
			options.float_type = float_type;
			std::string float_type_str = (options.float_type == atc::FloatType::float32 ? "float32" : "float64");
			std::cout << "Handling float_type: " << float_type_str << std::endl;
			
			for (const auto& int_type : std::vector<atc::IntType>{atc::IntType::uint32, atc::IntType::uint64}) {
				options.int_type = int_type;
				std::string int_type_str = (options.int_type == atc::IntType::uint32 ? "uint32" : "uint64");
				std::cout << "Handling int_type: " << int_type_str << std::endl;
				
				// Print start of entry
				out << dataset.name << "," << float_type_str << "," << int_type_str;
				for (auto target_psnr : target_psnrs) {
					std::cout << "For target PSNR: " << target_psnr << std::endl;
					
					// Perform compression and print data
					atc::Results res = compress_decompress(dataset, target_psnr, options);
					out << "," << res.error << "," << res.compression_factor << std::flush;
					
				}
				out << std::endl;
				
			}
		}
	}
	
	// Finish and move file
	out.close();
	std::rename((output_file + ".temp").c_str(), output_file.c_str());
}

// Asymmetric numeral systems

// General

int main() {
	check_hybrid_truncation_error_approximation();
	dump_rtmss_rate_distortion();
	dump_split_bitplane_truncation_output();
	test_core_flattenings();
	return 0;
}
