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

#include <cmath>
#include <cstdint>
#include <algorithm>
#include <numeric>
#include <string>

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

namespace atc {
	
	template<typename ftype>
	void Compressor<ftype>::decompress(CompressedIO& compressed_io, RawIOBase<ftype>* raw_io) {
		
		assert(raw_io != nullptr);
		if (options.calculate_error && options.autocrop)
			throw std::invalid_argument("Cannot use autocrop when calculating the error of the output tensor");
		if (options.calculate_error && options.cutout.size() > 0)
			throw std::invalid_argument("Cannot slice the output tensor when calculating the decompression error of the full tensor too");
		if (options.autocrop && options.cutout.size() > 0)
			throw std::invalid_argument("Cannot use autocrop when manually slicing the tensor too");
		
		print_banner("Start of decompression");
		start_timer("decompress");
		
		// Read header and verify validity
		read_header(compressed_io);
		if (data_type_id != raw_io->get_data_type_id())
			throw std::runtime_error("Compressed data represents tensor of a data type different from the expected data type");
		if (options.calculate_error && product(original_sizes) != raw_io->get_stored_data_size())
			throw std::runtime_error("Compressed data represents tensor of with sizes different from the sizes of the tensor which was just compressed");
		#ifndef ATC_FFTW
			if (options.apply_dct)
				throw std::runtime_error("Compressed data was compressed using HODCT preprocessing, however this version of ATC was compiled without FFTW and therefore cannot apply the inverse HODCT. In order to decompress this data, you need to recompile ATC with FFTW available.");
		#endif
		
		// Pre-process given cut-out
		// Cut-outs will already be used to calculate decompressed sizes in core dequantization (for buffer allocation)
		// Non-specified slicings are assumed to be the standard (0, -1, 1)
		options.cutout.resize(modes);
		if (options.storage_order == StorageOrder::MSI)
			std::reverse(options.cutout.begin(), options.cutout.end());
		for (uint8_t mode = 0; mode < modes; mode++)
			options.cutout[mode].update_mode_size(original_sizes[mode]);
		verify_slices();
		
		// Decode and dequantize core
		dequantize(compressed_io, QuantizationTarget::Core);
		
		// Initialize factors array, then decode and dequantize factors
		weighted_factors.resize(modes);
		for (uint8_t mode = 0; mode < modes; mode++)
			dequantize(compressed_io, QuantizationTarget::Factor, mode);
		unweight_factors();
		finalize_decompression_factors();
		
		// Optionally write dequantized core to file for further analysis
		// Filename will contain core sizes in most-significant-index-first order
		if (!options.error_components_path.empty()) {
			// Should also never be used when using cut-outs of any kind!
			if (options.storage_order != StorageOrder::MSI)
				throw std::invalid_argument("Cannot save error components unless storage order is MSI!");
			std::stringstream ss;
			ss << options.error_components_path + "-quantized-core-";
			ss << svd_sizes[current_mode_order[modes - 1]];
			for (uint8_t mode = modes - 1; mode-- > 0; )
				ss << "x" << svd_sizes[current_mode_order[mode]];
			ss << "-" << mode_to_string(current_mode_order[modes - 1]);
			for (uint8_t mode = modes - 1; mode-- > 0; )
				ss << "," << mode_to_string(current_mode_order[mode]);
			std::ofstream stream(ss.str(), std::ios::out | std::ios::binary);
			stream.write(reinterpret_cast<char*>(core.data()), product(svd_sizes)*sizeof(ftype));
			stream.close();
			for (uint8_t mode = modes; mode-- > 0; ) {
				std::ofstream stream(options.error_components_path + "-quantized-factor-" + mode_to_string(mode), std::ios::out | std::ios::binary);
				uint64_t shape[2];
				shape[0] = factors[mode].rows();
				shape[1] = factors[mode].cols();
				stream.write(reinterpret_cast<char*>(shape), 2*sizeof(*shape));
				stream.write(reinterpret_cast<char*>(factors[mode].data()), sizeof(ftype)*factors[mode].size());
				stream.close();
			}
		}
		
		// Only for research/debugging
		double core_quantization_sse = 0;
		std::vector<double> factors_scaled_quantization_sses(modes);
		if (options.test_error_approximation) {
			core_quantization_sse = (Eigen::Map<vtype>(core.data(), core.size()) - Eigen::Map<vtype>(core_copy.data(), core_copy.size())).squaredNorm();
			force_clear(core_copy);
			for (uint8_t mode = 0; mode < modes; mode++) {
				Eigen::MatrixXd factor = factors[mode].template cast<double>();
				Eigen::MatrixXd factor_copy = (options.householder ? householder_to_factor(factors_copy[mode]) : factors_copy[mode]).template cast<double>();
				Eigen::MatrixXd temp = factor_copy.leftCols(factor.cols()); // We need to make this copy to prevent a segfault on the next line
				factor_copy = temp;
				factors_scaled_quantization_sses[mode] = ((factor - factor_copy)*Eigen::Map<Eigen::VectorXd>(core_slice_norms[mode].data(), core_slice_norms[mode].size()).asDiagonal()).squaredNorm();
			}
			force_clear(factors_copy);
		}
		
		// ST-HOSVD
		resize_core(svd_sizes, quantized_sizes, "decompress_resize_core");
		if (options.autocrop && !std::all_of(options.cutout.begin(), options.cutout.end(), [] (const Slice& slice) {return slice.is_standard();})) {
			// Slicing was updated by autocrop, which may change the optimal decompression mode processing order
			size_t size = calculate_required_size_core(options.apply_dct);
			// This in turn may also change the minimum required memory for storing the core
			core.resize(size);
		}
		decompress_st_hosvd();
		force_clear(factors);
		
		// T-HODCT
		#ifdef ATC_FFTW
			if (options.apply_dct && std::any_of(apply_dct_to_mode.begin(), apply_dct_to_mode.end(), [] (const bool& val) {return val;}))
				decompress_t_hodct();
		#endif
		
		print_info("Decompressed tensor sizes: " + sizes_to_string(decompressed_sizes));
		
		// Save currently stored data and optionally calculate SSE
		start_timer("decompress_save_data");
		double sse = raw_io->write(core, options.calculate_error, product(decompressed_sizes));
		stop_timer();
		if (options.calculate_error) {
			// Data normsq, min and max should be up to date with the data stored in DataBuffer
			switch (options.error_metric) {
			case ErrorMetric::RelativeError:
				results.error = sqrt(sse/data_normsq);
				break;
			case ErrorMetric::RMSE:
				results.error = sqrt(sse/core.size());
				break;
			case ErrorMetric::PSNR:
				results.error = 20*log10((data_max - data_min)/sqrt(sse/core.size()));
				break;
			}
		}
		results.decompressed_sizes = std::move(decompressed_sizes);
		if (options.storage_order == StorageOrder::MSI)
			std::reverse(results.decompressed_sizes.begin(), results.decompressed_sizes.end());
		
		if (options.test_error_approximation || !options.error_components_path.empty()) {
			print_info("Target SSE: " + std::to_string(original_target_sse));
			print_info("Truncation SSE: " + std::to_string(rank_truncation_current_sse));
			double total_approximation_sse = 0;
			if (options.test_error_approximation) {
				print_info("Core quantization SSE: " + std::to_string(core_quantization_sse));
				for (uint8_t mode = 0; mode < modes; mode++)
					print_info("Factor " + mode_to_string(mode) + " SSE: " + std::to_string(factors_scaled_quantization_sses[mode]));
				total_approximation_sse = rank_truncation_current_sse + core_quantization_sse + std::accumulate(factors_scaled_quantization_sses.begin(), factors_scaled_quantization_sses.end(), 0.0);
				print_info("Total approximation SSE: " + std::to_string(total_approximation_sse));
			}
			double actual_sse_without_rounding = (Eigen::Map<vtype>(data_copy.data(), data_copy.size()) - Eigen::Map<vtype>(core.data(), core.size())).squaredNorm();
			print_info("Actual output SSE (without rounding): " + std::to_string(actual_sse_without_rounding));
			if (options.test_error_approximation)
				print_info("Relative error of approximation: " + std::to_string(abs(actual_sse_without_rounding - total_approximation_sse)/actual_sse_without_rounding));
		}
		
		stop_timer();
		print_banner("End of decompression");
		
	}
	
	template<typename ftype>
	void Compressor<ftype>::finalize_decompression_factors() {
		
		start_timer("decompress_finalize_factors");
		
		apply_dct_to_mode = std::vector<bool>(modes, options.apply_dct);
		for (uint8_t mode = 0; mode < modes; mode++) {
			
			Slice& slice = options.cutout[mode];
			
			// Perform Householder decompression if necessary
			if (options.householder)
				factors[mode] = householder_to_factor(factors[mode]);
			
			// Apply DCT to factor already if necessary
			// If not applied here, it will be applied to the full core using the fast DCT later
			// The DCT needs to be applied if:
			// - the slice is non-standard, because the slicing, DCT and factor need to be merged
			// - autocrop is enabled, because we need to crop based on the "real" factors (in the time domain, not the frequency domain)
			#ifdef ATC_FFTW
				if (options.apply_dct && (!slice.is_standard() || options.autocrop)) {
					if (dct_sizes[mode] != original_sizes[mode]) {
						// factors[mode] is a dct_sizes[mode] x quantized_sizes[mode] matrix
						// We need to multiply it on the left by an original_sizes[mode] x dct_sizes[mode] DCT matrix
						// However, we can only multiply it with a square DCT matrix (using the fast DCT)
						// So we pad the factor with zero rows so it becomes an original_sizes[mode] x quantized_sizes[mode] matrix
						mtype unpadded_factor = std::move(factors[mode]);
						factors[mode] = mtype::Zero(original_sizes[mode], quantized_sizes[mode]);
						factors[mode].topRows(dct_sizes[mode]) = unpadded_factor;
					}
						orthogonal_dct(factors[mode], true);
					apply_dct_to_mode[mode] = false;
				}
			#endif
			
			// Autocrop based on factors in real domain if requested
			// We do this before applying the cut-out because this step will determine the cut-out!
			if (options.autocrop) {
				size_t start_row = 0, end_row = 0;
				bool start_set = false;
				for (size_t row = 0; row < size_t(factors[mode].rows()); row++) {
					double sq_norm = 0;
					for (size_t col = 0; col < size_t(factors[mode].cols()); col++)
						sq_norm += pow(factors[mode](row, col), 2);
					if (sqrt(sq_norm/factors[mode].cols()) > options.autocrop_threshold) {
						if (!start_set) {
							start_row = row;
							start_set = true;
						}
						end_row = row + 1;
					}
				}
				slice = Slice(start_row, end_row);
				slice.update_mode_size(factors[mode].rows());
			}
			
			if (!slice.is_standard()) {
				
				// Calculate and apply convolution to factor matrix
			
				int8_t sign = (slice.get_step() > 0) - (slice.get_step() < 0);
				int64_t original_size = original_sizes[mode];
				mtype convolution = mtype::Zero(slice.get_size(), original_sizes[mode]); // convolution*U convolves U along the columns
				size_t row = 0;
				for (int64_t i = slice.get_start(); sign*i < sign*slice.get_stop(); i += slice.get_step()) {
					switch (slice.get_reduction()) {
						case Reduction::Downsampling: {
							// Simply sample the value at this point and ignore everything else
							convolution(row, i) = 1;
							break;
						}
						case Reduction::Box: {
							// Average values around the current point
							int64_t positive_step = abs(slice.get_step());
							int64_t start = std::max(i - positive_step/2, 0L);
							int64_t end = std::min(i - positive_step/2 + positive_step, original_size);
							for (int64_t j = start; j < end; j++)
								convolution(row, j) = 1.0/(end - start);
							break;
						}
						case Reduction::Lanczos: {
							int64_t a = 2*slice.get_step(); // Upscaled Lanczos window
							int64_t start = std::max(std::min(i - a, original_size - 1), 0L);
							int64_t end = std::max(std::min(i + a + 1, original_size), -1L);
							double kernel_sum = 0;
							for (int64_t j = start; sign*j < sign*end; j += sign) {
								double x = double(j - i)/slice.get_step(); // Lanczos filter coordinates (peak at 0), sign of x doesn't matter
								if (x == 0)
									convolution(row, j) = 1;
								else
									convolution(row, j) = a*sin(M_PI*x)*sin(M_PI*x/a)/(M_PI*M_PI*x*x);
								kernel_sum += convolution(row, j);
							}
							for (int64_t j = start; sign*j < sign*end; j += sign)
								convolution(row, j) /= kernel_sum; // Normalize the kernel so that it adds up to 1
							break;
						}
					}
					row++;
					
				}
				
				factors[mode] = convolution*factors[mode];
				
			}
			
		}
		
		if (options.autocrop && options.verbose) {
			print_info("Autocrop cut-out: ", false);
			for (uint8_t mode_index = 0; mode_index < modes; mode_index++)
				std::cout << " " << options.cutout[options.storage_order == StorageOrder::MSI ? modes - mode_index - 1 : mode_index];
			std::cout << std::endl;
		}
		
		stop_timer();
		
	}
	
	#include "explicit_instantiations/Compressor.txt"
	
}
