/*
 * This file handles the optional T-HODCT phase of the pipeline, both during compression and decompression.
 */

#ifdef ATC_FFTW

#include <fftw3.h>

#include "Compressor.hpp"

namespace atc {
	
	template<typename ftype>
	void Compressor<ftype>::compress_t_hodct() {
		
		// Perform transfrom
		start_timer("compress_hodct_project");
		std::vector<bool> transform_mode(modes, true);
		hodct(original_sizes, transform_mode);
		correct_hodct_orthogonality(original_sizes, transform_mode);
		stop_timer();
		
		// Truncate if necessary
		if (options.rank_truncation && options.dct_truncation_max_sse_share > 0) {
			
			start_timer("compress_hodct_truncate");
			
			// Determine truncation ranks
			start_timer("compress_hodct_truncate_choose_ranks");
			double dct_truncation_target_sse = options.dct_truncation_max_sse_share*rank_truncation_target_sse;
			std::vector<size_t> strides(modes, 1);
			for (size_t i = 1; i < modes; i++)
				strides[i] = strides[i - 1]*original_sizes[i - 1];
			// We assume that the modes are currently stored in the default order
			for (uint8_t mode = 0; mode < modes; mode++) {
				
				// Calculate slice norms by iterating over every coefficient in order
				std::vector<size_t> indices(modes, 0);
				std::vector<double> sq_core_slice_norms(original_sizes[mode], 0);
				for (size_t read_index = 0; read_index < core.size(); read_index++) {
					sq_core_slice_norms[indices[mode]] += core[read_index]*core[read_index];
					indices[0]++;
					// Check if any index exceeds its range (dct_sizes, because some modes may already be truncated)
					uint8_t pos = 0;
					while (indices[pos] >= dct_sizes[pos] && pos < modes - 1) {
						// When an index is exceeded, increase the reading pointer by the appropriate amount in case dct_sizes != original_sizes
						read_index += (original_sizes[pos] - dct_sizes[pos])*strides[pos];
						indices[pos] = 0;
						indices[pos + 1]++;
						pos++;
					}
				}
				
				// Truncate mode based on core slice norms
				double target_mode_truncation_sse = (dct_truncation_target_sse - rank_truncation_current_sse)/(modes - mode);
				double current_sse = 0;
				size_t truncation_rank = original_sizes[mode];
				while (truncation_rank > 1) {
					double new_sse = current_sse + sq_core_slice_norms[truncation_rank - 1];
					if (new_sse > target_mode_truncation_sse) {
						// Exceeded target SSE, stop
						rank_truncation_current_sse += current_sse;
						break;
					}
					current_sse = new_sse;
					truncation_rank--;
				}
				dct_sizes[mode] = truncation_rank;
				
			}
			stop_timer();
			
			// Crop tensor in-place
			resize_core(original_sizes, dct_sizes, "");
			core.resize(product(dct_sizes));
			
			print_info("T-HODCT-truncated core sizes: " + sizes_to_string(dct_sizes));
			
			stop_timer();
		}
		
	}
	
	template<typename ftype>
	void Compressor<ftype>::decompress_t_hodct() {
		
		// Pad truncated modes which must still be transformed
		// Also updated decompressed_sizes for these modes
		start_timer("decompress_hodct_pad");
		std::vector<size_t> old_decompressed_sizes = decompressed_sizes;
		for (uint8_t mode = 0; mode < modes; mode++)
			if (apply_dct_to_mode[mode])
				decompressed_sizes[mode] = original_sizes[mode];
		resize_core(old_decompressed_sizes, decompressed_sizes, "");
		stop_timer();
		
		// Perform transform along certain modes for which standard slices were requested
		start_timer("decompress_hodct_unproject");
		correct_hodct_orthogonality(decompressed_sizes, apply_dct_to_mode, true);
		hodct(decompressed_sizes, apply_dct_to_mode, true);
		stop_timer();
		
	}
	
	template<typename ftype>
	void Compressor<ftype>::hodct(const std::vector<size_t>& sizes, const std::vector<bool>& transform_mode, bool inverse) {
		
		// Perform transforms sequentially
		size_t size = product(sizes);
		size_t stride = 1; // Tracks product of sizes of all earlier modes
		for (uint8_t mode_index = 0; mode_index < modes; mode_index++) {
			uint8_t mode = current_mode_order[mode_index];
			int transform_size = sizes[mode];
			size_t next_stride = stride*transform_size;
			if (transform_mode[mode]) {
				size_t amount = size/transform_size;
				if (mode == 0) {
					// First mode, can be handled with one call by collapsing all higher modes
					dct(1, &transform_size, amount, core.data(), NULL, 1, transform_size, inverse);
				} else if (mode < modes - 1) {
					// View the data as a 3D tensor by collapsing all lower modes and all higher modes
					// We then transform the data slice by slice (contiguous slices)
					// Each slice is transformed by applying the same DCT to each fibre along this mode
					for (size_t slice_index = 0; slice_index < amount/stride; slice_index++) {
						dct(1, &transform_size, stride, &core[next_stride*slice_index], NULL, stride, 1, inverse);
					}
				} else {
					// Last mode, can be handled with one call by collapsing all lower modes
					dct(1, &transform_size, amount, core.data(), NULL, amount, 1, inverse);
				}
			}
			stride = next_stride;
		}
		
	}
	
	template<typename ftype>
	void Compressor<ftype>::correct_hodct_orthogonality(const std::vector<size_t>& sizes, const std::vector<bool>& transform_mode, bool inverse) {
		
		// Scale the first row of each transform
		size_t normalization_factor = 1;
		size_t stride = 1; // Tracks product of sizes of all earlier modes
		ftype factor = inverse ? sqrt(2) : 1/sqrt(2);
		for (uint8_t mode_index = 0; mode_index < modes; mode_index++) {
			uint8_t mode = current_mode_order[mode_index];
			size_t size = sizes[mode];
			size_t next_stride = stride*size;
			// Check if mode was actually transformed and needs a correction
			if (transform_mode[mode]) {
				normalization_factor *= 2*size;
				// Iterate over all elements of the first slice along this mode
				for (size_t higher_index = 0; higher_index < core.size()/next_stride; higher_index++) {
					for (size_t lower_index = 0; lower_index < stride; lower_index++) {
						// Current multi-index is [lower_index, 0, higher_index],
						// with lower_index representing the lower mode_index modes
						// and higher_index representing the higher (modes - mode_index - 1) modes
						core[lower_index + higher_index*next_stride] *= factor;
					}
				}
			}
			stride = next_stride;
		}
		
 		// Scale all data with 1/sqrt(2*mode_size) for each transform
		Eigen::Map<vtype> mapped_core(core.data(), core.size());
		mapped_core /= sqrt(normalization_factor);
		
	}
	
	template<typename ftype>
	void Compressor<ftype>::orthogonal_dct(mtype& data, bool inverse) {
		if (inverse)
			data.row(0) *= sqrt(2);
		std::vector<int> transform_sizes{int(data.rows())};
		dct(1, transform_sizes.data(), data.cols(), data.data(), NULL, 1, data.rows(), inverse);
		if (!inverse)
			data.row(0) /= sqrt(2);
		data /= sqrt(2*data.rows());
	}
	
	// Wrapper methods for FFTW for 32-bit and 64-bit precision
	template<>
	void Compressor<float>::dct(int modes_to_transform, const int* transform_sizes, int amount, float* data, const int* transform_strides,
		int element_stride, int distance, bool inverse) {
		fftwf_r2r_kind kind = inverse ? FFTW_REDFT01 : FFTW_REDFT10;
		#ifdef ATC_OPENMP
			fftwf_plan_with_nthreads(std::max(1, options.threads));
		#endif
		fftwf_plan plan = fftwf_plan_many_r2r(modes_to_transform, transform_sizes, amount,
			data, transform_strides, element_stride, distance,
			data, transform_strides, element_stride, distance,
			&kind, FFTW_ESTIMATE
		);
		if (plan) {
			fftwf_execute(plan);
			fftwf_destroy_plan(plan);
		}
	}
	template<>
	void Compressor<double>::dct(int modes_to_transform, const int* transform_sizes, int amount, double* data, const int* transform_strides,
		int element_stride, int distance, bool inverse) {
		fftw_r2r_kind kind = inverse ? FFTW_REDFT01 : FFTW_REDFT10;
		#ifdef ATC_OPENMP
			fftw_plan_with_nthreads(std::max(1, options.threads));
		#endif
		fftw_plan plan = fftw_plan_many_r2r(modes_to_transform, transform_sizes, amount,
			data, transform_strides, element_stride, distance,
			data, transform_strides, element_stride, distance,
			&kind, FFTW_ESTIMATE
		);
		if (plan) {
			fftw_execute(plan);
			fftw_destroy_plan(plan);
		}
	}
	
	#include "explicit_instantiations/Compressor.txt"
	
}

#endif
