/*
 * This file handles the ST-HOSVD phase of the pipeline, both during compression and decompression.
 */

#include <stdexcept>
#include <cmath>
#include <chrono>
#include <algorithm>
#include <numeric>
#include <fstream>
#include <sstream>
#include <Eigen/Eigenvalues>
#include <Eigen/Householder>
#include <Eigen/SVD>

#include "hptt/hptt.h"
#include "Compressor.hpp"
#include "util.hpp"

namespace atc {
	
	template<typename ftype>
	void Compressor<ftype>::compress_st_hosvd() {
		
		start_timer("compress_st_hosvd");
		
		// Preliminaries
		
		factors = std::vector<mtype>(modes);
		sq_singular_values = std::vector<std::vector<ftype>>(modes);
		svd_sizes = dct_sizes; // Default value, even if there is no truncation
		
		// Mode processing order
		std::vector<uint8_t> default_mode_order(modes);
		std::iota(default_mode_order.begin(), default_mode_order.end(), 0);
		if (options.mode_processing_order.size() == 0) {
			// If no custom mode processing order is given, use heuristic from ST-HOSVD paper: increasing mode sizes
			std::vector<std::pair<size_t, uint8_t>> sorted_sizes(modes);
			for (uint8_t mode = 0; mode < modes; mode++)
				sorted_sizes[mode] = {svd_sizes[mode], mode};
			std::sort(sorted_sizes.begin(), sorted_sizes.end());
			options.mode_processing_order.resize(modes);
			for (uint8_t mode = 0; mode < modes; mode++)
				options.mode_processing_order[mode] = sorted_sizes[mode].second;
		}
		
		// Timings
		if (options.timings) {
			results.timings["compress_st_hosvd_transpose"] = 0;
			results.timings["compress_st_hosvd_select_base"] = 0;
			if (options.householder)
				results.timings["compress_st_hosvd_convert_householder"] = 0;
			results.timings["compress_st_hosvd_project"] = 0;
		}
		
		// Iterate over modes, alternatingly transposing and projecting
		// Because of the circular shift, the mode order is rotated modes times, so after the full loop we're back to the original mode order
		// So we don't need to record these changes
		temp_core = std::vector<ftype>(core.size()); // Temporary buffer already needs to be allocated here for initial transposition
		print_status("Transposing to mode order " + mode_order_to_string(options.mode_processing_order) + " for processing" + (options.storage_order == StorageOrder::MSI ? " (reverse mode processing order)" : ""));
		transpose(options.mode_processing_order, svd_sizes, "compress_st_hosvd_transpose"); // Transpose to processing order
		for (uint8_t mode_index = 0; mode_index < modes; mode_index++)
			st_hosvd_compress_mode(mode_index);
		
		if (options.rank_truncation)
			print_info("ST-HOSVD-truncated core sizes: " + sizes_to_string(svd_sizes));
		
		// Determine mode storage order
		if (options.mode_storage_order.size() == 0) {
			// If no custom mode storage order is given, use heuristic: increasing mode sizes in MSI order, but we use LSI internally so decreasing mode sizes
			std::vector<std::pair<size_t, uint8_t>> sorted_sizes(modes);
			for (uint8_t mode = 0; mode < modes; mode++)
				sorted_sizes[mode] = {svd_sizes[mode], mode};
			std::sort(sorted_sizes.begin(), sorted_sizes.end());
			std::reverse(sorted_sizes.begin(), sorted_sizes.end());
			options.mode_storage_order.resize(modes);
			for (uint8_t mode = 0; mode < modes; mode++)
				options.mode_storage_order[mode] = sorted_sizes[mode].second;
		}

		// Transpose to mode storage order for quantization
		print_status("Transposing to storage order " + mode_order_to_string(options.mode_storage_order));
		transpose(options.mode_storage_order, svd_sizes, "compress_st_hosvd_transpose");
		force_clear(temp_core); // Clear temporary buffer
		core.resize(product(svd_sizes));
		
		// Optionally write core to file for further analysis
		// Filename will contain core sizes in most-significant-index-first order
		if (!options.core_path.empty() || !options.error_components_path.empty()) {
			if (options.storage_order != StorageOrder::MSI)
				throw std::invalid_argument("Cannot save core unless storage order is MSI!");
			std::stringstream ss;
			ss << (!options.core_path.empty() ? options.core_path + "_" : options.error_components_path + "-truncated-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]];
			if (!options.error_components_path.empty()) {
				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()), core.size()*sizeof(ftype));
			stream.close();
		}
		
		// Adjust SSE for optional ST-HOSVD rank truncation
		target_sse_left -= rank_truncation_current_sse;
		
		stop_timer();
		
		if (options.timings) {
			print_info("ST-HOSVD timings:");
			print_info("- Tensor transposition time: " + std::to_string(results.timings["compress_st_hosvd_transpose"]) + " s");
			print_info("- Base selection: " + std::to_string(results.timings["compress_st_hosvd_select_base"]) + " s");
			if (options.householder)
				print_info("- Convert to Householder reflectors: " + std::to_string(results.timings["compress_st_hosvd_convert_householder"]) + " s");
			print_info("- Core projection: " + std::to_string(results.timings["compress_st_hosvd_project"]) + " s");
		}
		
	}
	
	template<typename ftype>
	void Compressor<ftype>::st_hosvd_compress_mode(uint8_t mode_index) {
		
		// Preliminaries
		auto time = time_now();
		uint8_t mode = options.mode_processing_order[mode_index];
		print_status("Compressing mode " + mode_to_string(mode));
		size_t rows = dct_sizes[mode];
		size_t cols = product(svd_sizes)/rows; // svd_sizes equals dct_sizes in all yet uncompressed modes
		Eigen::Map<mtype> core_matrix(core.data(), rows, cols);
		
		// Calculate and sort SVD
		std::vector<std::pair<ftype, size_t>> sorted_sq_singular_values(rows);
		mtype singular_vectors;
		if (options.precise_svd) {
			// Calculate SVD directly with BDCSVD
			Eigen::BDCSVD<mtype> svd(core_matrix, Eigen::ComputeThinU);
			singular_vectors = svd.matrixU();
			for (size_t i = 0; i < rows; i++)
				sorted_sq_singular_values[i] = {pow(svd.singularValues()[i], 2), i};
		} else {
			// Calculate SVD indirectly using eigenvalue decomposition of Gramian matrix
			// Faster and almost always accurate enough for our purposes
			Eigen::SelfAdjointEigenSolver<mtype> solver(core_matrix * core_matrix.transpose());
			vtype eigenvalues = solver.eigenvalues().real();
			singular_vectors = solver.eigenvectors().real();
			for (size_t i = 0; i < rows; i++)
				sorted_sq_singular_values[i] = {std::max(ftype(0), ftype(eigenvalues(i))), i};
			std::sort(sorted_sq_singular_values.begin(), sorted_sq_singular_values.end());
			std::reverse(sorted_sq_singular_values.begin(), sorted_sq_singular_values.end());
		}
		
		sq_singular_values[mode].resize(sorted_sq_singular_values.size());
		for (size_t i = 0; i < sorted_sq_singular_values.size(); i++)
			sq_singular_values[mode][i] = sorted_sq_singular_values[i].first;
		if (options.print_squared_singular_values) {
			print_status("Squared singular values for mode " + mode_to_string(mode) + ": ", false); print_vector(sq_singular_values[mode]);
		}
		
		// Determine compression rank
		double target_mode_truncation_sse = (options.static_rank_truncation_thresholding ?
			rank_truncation_target_sse/modes
			: (rank_truncation_target_sse - rank_truncation_current_sse)/(modes - mode_index));
		size_t truncation_rank = rows; // truncation_rank = r means that the first r singular vectors are kept
		if (options.rank_truncation) {
			// Perform rank truncation using the ST-HOSVD rule of thumb for error distribution per mode
			double current_sse = 0;
			while (truncation_rank > 1) {
				double new_sse = current_sse + sorted_sq_singular_values[truncation_rank - 1].first;
				if (new_sse > target_mode_truncation_sse) {
					// Exceeded target SSE, stop
					rank_truncation_current_sse += current_sse;
					break;
				}
				current_sse = new_sse;
				truncation_rank--;
			}
		}
		svd_sizes[mode] = truncation_rank;
		
		// Sort and truncate factor
		mtype factor(rows, truncation_rank);
		for (size_t i = 0; i < truncation_rank; i++)
			factor.col(i) = singular_vectors.col(sorted_sq_singular_values[i].second);
		
		if (options.timings)
			results.timings["compress_st_hosvd_select_base"] += elapsed_time(time);
		
		if (options.householder) {
			// Perform Householder compression if necessary
			time = time_now();
			mtype reflectors = factor_to_householder(factor);
			// Adjust factor to match householder_to_factor(reflectors) later on
			for (size_t i = 0; i < truncation_rank; i++)
				if (reflectors(i, i) < 0)
					factor.col(i) *= -1;
			if (options.timings)
				results.timings["compress_st_hosvd_convert_householder"] += elapsed_time(time);
			factors[mode] = reflectors;
		} else {
			factors[mode] = factor;
		}
		
		// Optionally write factors to file for further analysis
		if (!options.factors_path.empty() || !options.error_components_path.empty()) {
			std::ofstream stream((!options.factors_path.empty() ? options.factors_path + "_" : options.error_components_path + "-full-factor-") + mode_to_string(mode), std::ios::out | std::ios::binary);
			uint64_t shape[2];
			shape[0] = factor.rows();
			shape[1] = (!options.factors_path.empty() ? factor.cols() : factor.rows());
			stream.write(reinterpret_cast<char*>(shape), 2*sizeof(*shape));
			stream.write(reinterpret_cast<char*>(factor.data()), sizeof(ftype)*factor.size());
			// Add truncated columns if necessary
			if (!options.error_components_path.empty())
				for (size_t i = truncation_rank; i < rows; i++)
					stream.write(reinterpret_cast<char*>(&(singular_vectors(0, sorted_sq_singular_values[i].second))), sizeof(ftype)*factor.rows());
			stream.close();
		}
		
		// Apply projection
		multiply_core(factor.transpose(), core_matrix, "compress_st_hosvd_project");
		
	}
	
	template<typename ftype>
	Eigen::Matrix<ftype, Eigen::Dynamic, Eigen::Dynamic> Compressor<ftype>::factor_to_householder(const mtype& factor) {
		size_t rows = factor.rows();
		size_t cols = factor.cols();
		Eigen::HouseholderQR<mtype> qr(factor);
		mtype res = qr.matrixQR();
		vtype h = qr.hCoeffs();
		// Scale the reflectors to follow our convention
		for (size_t i = 0; i < cols; i++) {
			res.block(i + 1, i, rows - i - 1, 1) *= sqrt(h(i)/2);
			if (h(i) == 0) {
				// LAPACK will set h(i) to 0 if no reflection was necessary for this column (the column was already equal to +-e_i)
				// However, our decompressor doesn't know the h-coefficients so it will assume a reflection was performed with the given reflector (e_i)
				// Therefore, by reflecting over e_i, the decompressor will cause a sign switch in this factor column and this column alone
				// To compensate for this, we can switch the sign of the corresponding diagonal element in the QR-matrix,
				// thus causing a sign switch of the corresponding column in the factor adjustment step later
				res(i, i) *= -1;
			}
		}
		return res;
	}
	
	template<typename ftype>
	void Compressor<ftype>::decompress_st_hosvd() {
		
		start_timer("decompress_st_hosvd");
		
		// Preliminaries
		
		decompressed_sizes = quantized_sizes;
		
		// Timings
		if (options.timings) {
			results.timings["decompress_st_hosvd_transpose"] = 0;
			results.timings["decompress_st_hosvd_unproject"] = 0;
		}
		
		// Iterate over modes, compressing (and implicitly transposing) each mode
		temp_core = std::vector<ftype>(core.size()); // Allocate temporary buffer already because of the next transposition
		transpose(decompression_mode_processing_order, quantized_sizes, "decompress_st_hosvd_transpose"); // Transpose to processing order
		for (uint8_t mode_index = 0; mode_index < modes; mode_index++)
			st_hosvd_decompress_mode(decompression_mode_processing_order[mode_index]);
		
		// Transpose modes back to original order
		print_status("Transposing back to original mode order");
		std::vector<uint8_t> default_mode_order(modes);
		std::iota(default_mode_order.begin(), default_mode_order.end(), 0);
		transpose(default_mode_order, decompressed_sizes, "decompress_st_hosvd_transpose");
		force_clear(temp_core); // Clear temporary buffer
		// We do NOT resize the core here to product(decompressed_sizes) because the core size may still need to increase for the T-HODCT
		
		stop_timer();
		
	}
	
	template<typename ftype>
	void Compressor<ftype>::st_hosvd_decompress_mode(uint8_t mode) {
		
		const Slice& slice = options.cutout[mode];
		if (options.verbose) {
			print_status("Decompressing mode " + mode_to_string(mode), false);
			if (!slice.is_standard())
				std::cout << " with cutout " << slice;
			std::cout << std::endl;
		}
		
		// Calculate dimensions
		size_t rows = quantized_sizes[mode];
		size_t cols = product(decompressed_sizes)/rows; // decompressed_sizes equals compressed_sizes in all modes which are still compressed
		Eigen::Map<mtype> core_matrix(core.data(), rows, cols);
		decompressed_sizes[mode] = factors[mode].rows();
		
		// Apply factor
		multiply_core(factors[mode], core_matrix, "decompress_st_hosvd_unproject");
		
	}
	
	template<typename ftype>
	void Compressor<ftype>::multiply_core(const mtype& factor, const Eigen::Map<mtype>& core_matrix, const std::string& timer_label) {
		
		// Because of the circular mode shift, we combine the matrix multiplication and tensor transposition here
		
		auto time = time_now();
		size_t cols = core_matrix.cols();
		size_t new_rows = factor.rows();
		Eigen::Map<mtype> temp_transposed_core_matrix(temp_core.data(), cols, new_rows);
		// Compute new_transposed_core = (factor * core)^T
		#ifdef ATC_OPENMP
		if (options.explicitly_parallellized_core_multiplications) {
			// We manually split up the matrix multiplications in column-based blocks (which Eigen seems to do as well internally, but while allocating large block caches)
			size_t old_rows = core_matrix.rows();
			omp_set_num_threads(options.threads);
			Eigen::setNbThreads(1);
			#pragma omp parallel
			{
				int thread = omp_get_thread_num();
				int threads = omp_get_num_threads();
				size_t min_col = (thread*cols)/threads;
				size_t max_col = ((thread + 1)*cols)/threads;
				size_t block_cols = max_col - min_col;
				temp_transposed_core_matrix.transpose().block(0, min_col, new_rows, block_cols).noalias() = factor*core_matrix.block(0, min_col, old_rows, block_cols);
			}
		} else
		#endif
		{
			temp_transposed_core_matrix.transpose().noalias() = factor*core_matrix;
		}
		Eigen::setNbThreads(options.threads);
		
		core.swap(temp_core);
		if (options.timings)
			results.timings[timer_label] += elapsed_time(time);
		
	}
	
	template<typename ftype>
	Eigen::Matrix<ftype, Eigen::Dynamic, Eigen::Dynamic> Compressor<ftype>::householder_to_factor(const mtype& reflectors) {
		size_t rows = reflectors.rows();
		size_t cols = reflectors.cols();
		// Scale reflectors back to the LAPACK/Eigen convention
		mtype scaled_reflectors = reflectors;
		vtype h = vtype::Zero(cols);
		for (size_t i = cols; i-- > 0;) {
			ftype sq_norm = scaled_reflectors.block(i + 1, i, rows - i - 1, 1).squaredNorm();
			if (sq_norm >= 1)
				throw std::runtime_error("Decoded reflector with norm of trailing components >= 1");
			ftype first_component = sqrt(1 - sq_norm);
			// Divide reflector by first component, so that implicit first component becomes 1
			scaled_reflectors.block(i + 1, i, rows - i - 1, 1) /= first_component;
			// Multiply default factor 2 by first component twice to compensate for scaling of v_i and v_i^T
			h(i) = 2*first_component*first_component;
		}
		return Eigen::HouseholderSequence<mtype, vtype>(scaled_reflectors, h) * mtype::Identity(rows, cols);
	}
	
	template<typename ftype>
	void Compressor<ftype>::transpose(const std::vector<uint8_t>& new_mode_order, const std::vector<size_t>& current_sizes, const std::string& timer_label) {
		
		std::chrono::high_resolution_clock::time_point time = time_now();
		
		if (current_mode_order != new_mode_order) {
			
			// Calculate transposed sizes
			std::vector<int> transposed_sizes(modes);
			for (uint8_t mode = 0; mode < modes; mode++)
				transposed_sizes[mode] = current_sizes[current_mode_order[mode]];
			
			// Determine permutation
			std::vector<int> perm(modes);
			for (uint8_t mode_index = 0; mode_index < modes; mode_index++)
				perm[mode_index] = std::distance(current_mode_order.begin(), std::find(current_mode_order.begin(), current_mode_order.end(), new_mode_order[mode_index]));
			
			// Execute transposition
			auto plan = hptt::create_plan(perm.data(), modes,
				1, core.data(), transposed_sizes.data(), NULL,
				0, temp_core.data(), NULL,
				hptt::ESTIMATE, std::max(1, options.threads)
			);
			plan->execute();
			core.swap(temp_core);
			
			current_mode_order = new_mode_order;
			
		}
		
		if (options.timings)
			results.timings[timer_label] += elapsed_time(time);
		
	}
	
	#include "explicit_instantiations/Compressor.txt"
	
}
