/*
 * This file implements the atc::Compressor general compression methods.
 */

#include <string>
#include <sstream>
#include <Eigen/Core>

#include "Compressor.hpp"
#include "CompressedIO.hpp"
#include "RawIO.hpp"
#include "util.hpp"

namespace atc {
	
	template<typename ftype>
	void Compressor<ftype>::compress(RawIOBase<ftype>* raw_io, const std::vector<size_t>& sizes, double target_error, CompressedIO& compressed_io) {
		
		print_banner("Start of compression");
		start_timer("compress");
		
		// Input verification
		if (sizes.size() < 3)
			throw std::invalid_argument("Cannot compress data with less than 3 modes");
		if (sizes.size() > 255)
			throw std::invalid_argument("Cannot compress data with more than 255 modes");
		if (std::find(sizes.begin(), sizes.end(), 0) != sizes.end())
			throw std::invalid_argument("Mode cannot have a length of 0");
		if (std::find(sizes.begin(), sizes.end(), 1) != sizes.end())
			throw std::invalid_argument("Mode cannot have a length of 1");
		for (size_t size : sizes)
			if (size > size_t(std::numeric_limits<int>::max()))
				throw std::invalid_argument("Mode cannot have a length >= 2^31");
		if (target_error < 0 && (options.error_metric == ErrorMetric::RelativeError || options.error_metric == ErrorMetric::RMSE))
			throw std::invalid_argument("Error cannot be negative");
		#ifndef ATC_FFTW
			if (options.apply_dct)
				throw std::invalid_argument("Options.apply_dct was set to true, however this version of ATC was compiled without FFTW and therefore does not support DCT preprocessing. Set this option to false or recompile with FFTW available.");
		#endif
		if (options.rank_truncation) {
			if (options.rank_truncation_max_sse_share < 0 || options.rank_truncation_max_sse_share > 1)
				throw std::invalid_argument("Given ST-HOSVD truncation maximal SSE share is not in the interval [0, 1]");
			if (options.dct_truncation_max_sse_share < 0 || options.dct_truncation_max_sse_share > 1)
				throw std::invalid_argument("Given DCT truncation maximal SSE share is not in the interval [0, 1]");
		}
		if (options.quantization_rel_error_margin <= 0)
			throw std::invalid_argument("Given quantization relative error margin must be strictly positive");
		if (options.print_squared_singular_values && !options.verbose) {
			warn("The option print_squared_singular_values is set to true, while the option verbose is set to false, which means nothing will be printed.");
			options.print_squared_singular_values = false;
		}
		modes = sizes.size();
		verify_mode_permutation(options.mode_processing_order, "mode processing order");
		verify_mode_permutation(options.mode_storage_order, "mode storage order");
		
		// Interpret modes in reverse if necessary because ATC uses LSI mode ordering internally
		original_sizes = sizes;
		if (options.storage_order == StorageOrder::MSI) {
			std::reverse(original_sizes.begin(), original_sizes.end());
			auto transform = [this] (uint8_t mode) { return modes - 1 - mode; };
			std::transform(options.mode_processing_order.begin(), options.mode_processing_order.end(), options.mode_processing_order.begin(), transform);
			// Mode storage order needs to be interpreted differently in two ways:
			// - Each mode index has a different interpretation (mode 0 in row-major is mode (modes - 1) in column-major)
			// - The storage order must be interpreted in reverse
			std::transform(options.mode_storage_order.begin(), options.mode_storage_order.end(), options.mode_storage_order.begin(), transform);
			std::reverse(options.mode_storage_order.begin(), options.mode_storage_order.end());
		}
		
		print_info("Original tensor sizes: " + sizes_to_string(original_sizes));
		
		start_timer("compress_load_data");
		core = raw_io->read(original_sizes, options.calculate_error);
		stop_timer();
		
		// Calculate input statistics
		start_timer("compress_input_statistics");
		data_normsq = squared_norm(core);
		data_min = *std::min_element(core.begin(), core.end());
		data_max = *std::max_element(core.begin(), core.end());
		stop_timer();
		print_info("Input statistics:");
		print_info("- Minimum: " + std::to_string(data_min));
		print_info("- Maximum: " + std::to_string(data_max));
		print_info("- Norm: " + std::to_string(sqrt(data_normsq)));
		
		// Calculate target SSE
		switch (options.error_metric) {
			case ErrorMetric::RelativeError:
				original_target_sse = pow(target_error, 2)*data_normsq;
				break;
			case ErrorMetric::RMSE:
				original_target_sse = pow(target_error, 2)*core.size();
				break;
			case ErrorMetric::PSNR:
				original_target_sse = pow(data_max - data_min, 2)/pow(10, target_error/10)*core.size();
				break;
		}
		target_sse_left = original_target_sse;
		if (options.verbose) {
			double rel_error = sqrt(target_sse_left/data_normsq);
			double rmse = sqrt(target_sse_left/core.size());
			double psnr = 20*log10((data_max - data_min)/rmse);
			print_info("Target error:");
			print_info("- Relative error: " + std::to_string(rel_error));
			print_info("- RMSE: " + std::to_string(rmse));
			print_info("- PSNR (dB): " + std::to_string(psnr));
		}
		if (options.rank_truncation) {
			rank_truncation_current_sse = 0;
			rank_truncation_target_sse = options.rank_truncation_max_sse_share*target_sse_left;
		}
		if (options.test_error_approximation || !options.error_components_path.empty()) {
			data_copy = core;
			original_target_sse = target_sse_left;
		}
		
		// Initialize mode order
		current_mode_order.resize(modes);
		std::iota(current_mode_order.begin(), current_mode_order.end(), 0);
		
		// Optional DCT pre-processing
		dct_sizes = original_sizes;
		#ifdef ATC_FFTW
			if (options.apply_dct)
				compress_t_hodct();
		#endif
		
		// ST-HOSVD
		compress_st_hosvd();
		
		if (options.test_error_approximation) {
			core_copy = core;
		}
		if (options.test_error_approximation || options.householder_weighting_method == HouseholderWeightingMethod::Exact) {
			factors_copy = factors;
		}
		
		// Open IO and write header
		data_type_id = raw_io->get_data_type_id();
		write_header(compressed_io);
		
		// Quantization and encoding of core
		// Also calculates quantized core slice norms and ranks for factor weighting
		quantize(compressed_io, QuantizationTarget::Core);
		force_clear(core);
		
		// Weight, quantize and encode factors
		weight_factors();
		force_clear(factors);
		for (uint8_t mode = 0; mode < modes; mode++)
			quantize(compressed_io, QuantizationTarget::Factor, mode);
		force_clear(weighted_factors);
		
		stop_timer();
		
		// Report sizes
		compressed_io.close(); // Make sure the compressed stream is finalized before checking the number of bytes written
		results.compressed_size = compressed_io.get_bytes_written();
		size_t original_size = product(original_sizes);
		size_t uncompressed_byte_size = original_size*raw_io->get_data_type_size();
		size_t compressed_byte_size = results.compressed_size;
		results.compression_factor = uncompressed_byte_size/double(compressed_byte_size);
		if (options.verbose) {
			print_info("Compression statistics:");
			print_info("- Uncompressed size (bytes): " + std::to_string(uncompressed_byte_size));
			print_info("- Compressed size (bytes): " + std::to_string(compressed_byte_size));
			print_info("- Compression factor: " + std::to_string(results.compression_factor));
			print_info("- Bitrate (bits per value): " + std::to_string(8*double(compressed_byte_size)/original_size));
		}
		
		print_banner("End of compression");
		
	}
	
	#include "explicit_instantiations/Compressor.txt"
	
}
