/*
 * This file defines the atc::Compressor class and declares its members.
 */

#pragma once

#include <vector>
#include <stack>
#include <chrono>
#include <string>
#include <memory>
#include <cmath>
#include <functional>
#include <Eigen/Core>

#include "CompressorBase.hpp"
#include "DataType.hpp"
#include "Options.hpp"
#include "Results.hpp"
#include "IOMode.hpp"
#include "RawIO.hpp"
#include "CompressedIO.hpp"
#include "Breakpoint.hpp"

namespace atc {
	
	/*
	 * This class handles most of the compression work, only delegating the original/decompressed data IO (atc::RawIO) and compressed data IO (atc::CompressedIO).
	 * Ideally all this functionality should be decoupled more, however it is quite hard to do due to connections in between the T-HODCT, ST-HOSVD, quantization/encoding and other things.
	 * Because this class has to store the ST-HOSVD core, which depends on the floating-point type used, the entire class is templated in terms of this type parameter.
	 */
	template<typename ftype>
	class Compressor : public CompressorBase {
	public:
		
		Compressor(const Options& options); // Stores options and initializes multi-threading set-up
		~Compressor(); // Handles end of multi-threading set-up for FFTW and Eigen
		
		/*
		 * In the following two functions, raw_io is passed as a pointer to allow for polymorphism:
		 * The passed RawIOBase pointer is simply an interface to a RawIO object with a certain dtype.
		 * In this situation, polymorphism has a negligible performance cost due to the small amount of polymorphic calls,
		 * while it makes these functions a lot simpler because they do not have to be templated and explicitly instantiated for each valid <ftype, dtype> combination.
		 */
		
		// Compresses the original data managed by the raw_io object and writes the compressed output to the compressed_io object.
		void compress(RawIOBase<ftype>* raw_io, const std::vector<size_t>& sizes, double target_error, CompressedIO& compressed_io);
		// Decompresses the compressed data managed by the compressed_io object and writes the decompressed data to the raw_io object.
		void decompress(CompressedIO& compressed_io, RawIOBase<ftype>* raw_io);
		
	private:
		
		// Types
		enum class QuantizationTarget {Core, Factor}; // Indicates whether the core or a factor is currently being quantized when calling the quantization procedures.
		typedef Eigen::Matrix<ftype, Eigen::Dynamic, Eigen::Dynamic> mtype; // Short-hand for the Eigen matrix type corresponding to this floating-point type
		typedef Eigen::Matrix<ftype, Eigen::Dynamic, 1> vtype; // Short-hand for the Eigen vector type corresponding to this floating-point type
		
		std::vector<ftype> core, temp_core; // core stores the original data or the core after the ST-HOSVD, temp_core is simply a temporary buffer allocated for use during the ST-HOSVD.
		std::vector<mtype> factors; // ST-HOSVD factors (after Householder compression)
		std::vector<std::vector<ftype>> sq_singular_values; // The squared singular values obtained during the ST-HOSVD for each mode
		std::vector<std::vector<ftype>> weighted_factors; // Flattened ST-HOSVD factors after applying quantization weights
		
		// Only for research/debugging
		double original_target_sse;
		std::vector<ftype> data_copy, core_copy;
		std::vector<mtype> factors_copy;
		
		// T-HODCT pre-processing methods
		
		/*
		 * Applies a DCT to the stored data along all modes.
		 * Then optionally truncates each mode, following the same method as the ST-HOSVD.
		 */
		void compress_t_hodct();
		/*
		* Applies an inverse DCT to the stored data along all modes where apply_dct_to_mode[mode] is true.
		* The others already had an inverse DCT applied to them together with the corresponding factor and linear filter during the inverse HOSVD.
		* On entry the core should have the following mode sizes (stored in decompressed_sizes):
		*  - dct_sizes for modes where apply_dct_to_mode[mode] is true.
		*  - otherwise the corresponding slice size.
		*  The core should also be stored using the default mode order.
		*/
		void decompress_t_hodct();
		/*
		 * Applies a DCT-II or DCT-III (inverse) to all requested transform_modes of the tensor stored in core.
		 * This does not include any correction step or truncation.
		 */
		void hodct(const std::vector<size_t>& sizes, const std::vector<bool>& transform_mode, bool inverse = false);
		/*
		 * Scales the first slice of each transformed mode as well as the full tensor
		 * in such a way that the performed transforms are orthogonal.
		 * transform_mode indicate the modes along which an actual DCT was performed and which therefore need a correction.
		 */
		void correct_hodct_orthogonality(const std::vector<size_t>& sizes, const std::vector<bool>& transform_mode, bool inverse = false);
		// Performs a columnwise DCT-II or DCT-III (inverse) on a matrix but normalizes the output such that the transformation is orthogonal.
		void orthogonal_dct(mtype& data, bool inverse = false);
		// Applies a DCT-II or DCT-III (inverse) to a series of vectors.
		void dct(int modes_to_transform, const int* transform_sizes, int amount, ftype* data, const int* transform_strides, int element_stride, int distance, bool inverse = false);
		
		// ST-HOSVD methods
		
		/*
		 * Compresses the stored data using the ST-HOSVD.
		 * Both the input and output are stored in core, but temp_core is used as a temporary buffer.
		 * Specifically, temp_core is allocated and deallocated (simply through resizing) in this function.
		 */
		void compress_st_hosvd();
		/*
		 * Finds a suitable basis for the core (stored in data) along the current mode and projects the core using this basis.
		 * Projects the core using this basis.
		 * Both the input and output are stored in core, but temp_core is used as a temporary buffer and assumed to be sufficiently large.
		 * core should also be stored with a mode order such that the mode to be compressed is at the front.
		 * After this function, the mode order will have been cyclically shifted one position to the left, ready for the next compression step.
		 * However, because this function is called modes times anyway, we do not update current_mode_order.
		 */
		void st_hosvd_compress_mode(uint8_t mode_index);
		/*
		 * Decomposes the given factor (an orthogonal matrix, possibly with less columns than rows) into a series of Householder reflectors.
		 * The output matrix has the same dimensions, stores the reflectors in the lower Hessenberg part and stores the signs of R on the diagonal.
		 * The first component of each reflector is implicitly stored: it must be positive and make the norm of the full reflector equal 1.
		 */
		mtype factor_to_householder(const mtype& factor);
		/*
		 * Takes the stored core and factors and performs ST-HOSVD decompression.
		 * Allocates and deallocates sufficient space in temp_core to be used as a temporary buffer.
		 */
		void decompress_st_hosvd();
		/*
		 * ST-HOSVD-decompresses the core along the given mode.
		 * Both the input and output are stored in core, but temp_core is used as a temporary buffer and assumed to be sufficiently large.
		 * core should also be stored with a mode order such that the mode to be compressed is at the front.
		 * After this function, the mode order will have been cyclically shifted one position to the left, ready for the next compression step.
		 * However, because this function is called modes times anyway, we do not update current_mode_order.
		 */
		void st_hosvd_decompress_mode(uint8_t mode);
		/*
		 * Computes the operation core := (factor * core)^T, which effectively applies a cyclic shift to the core mode order.
		 * The method used here depends on options.explicitly_parallellized_core_multiplications.
		 * Both the input and output are stored in core, but temp_core is used as a temporary buffer and assumed to be sufficiently large.
		 */
		void multiply_core(const mtype& factor, const Eigen::Map<mtype>& core_matrix, const std::string& timer_label);
		/*
		 * Converts a sequence of Householder reflectors back to an orthogonal matrix.
		 * The reflectors are given in the lower Hessenberg part of the given matrix.
		 * The first component of each reflector is computed as the positive number such that the reflector has norm 1.
		 */
		mtype householder_to_factor(const mtype& reflectors);
		/*
		 * Transposes the tensor stored in core in the order current_mode_order to the order new_mode_order.
		 * Also updates current_mode_order.
		 * Assumes that temp_core is sufficiently large to be used as a temporary.
		 */
		void transpose(const std::vector<uint8_t>& new_mode_order, const std::vector<size_t>& current_sizes, const std::string& timer_label);
		
		// Quantization methods
		
		double core_price, core_eps; // Metrics used for factor quantization stopping criterion
		size_t core_total_bytes; // Metrics used for factor quantization stopping criterion
		std::vector<std::vector<double>> core_slice_norms; // The slice norms of the core after quantization, used for weighting the factor coefficients
		// Calls the corresponding quantize<itype> depending on options.int_type.
		void quantize(CompressedIO& io, QuantizationTarget target, uint8_t mode = 0);
		// Calls the corresponding quantize<itype> depending on options.dequantization_correction.
		template<typename itype>
		void quantize(CompressedIO& io, QuantizationTarget target, uint8_t mode = 0);
		/*
		 * Quantizes and encodes the core or factors
		 * When quantizing the core, also calculates the quantized core slice norms and ranks.
		 */
		template<typename itype, bool dequantization_correction>
		void quantize(CompressedIO& io, QuantizationTarget target, uint8_t mode = 0);
		// Weights factors coefficients according to their effect on the total error and stores the result in weighted_factors.
		void weight_factors();
		// Computes the quantization weights/scaling factors for all Householder-compressed factor coefficients and returns them as a vector.
		vtype calculate_householder_weights(uint8_t mode, bool allow_zero_weights = true);
		// Updates the stored core slice norms based on the stored floating-point core tensor.
		void update_core_slice_norms();
		/*
		* Updates the stored core slice norms and quantized sizes.
		* If use_quantized_data is true (default), this is done based on the given quantized data, scale and breakpoint,
		* else the stored floating-point core is used.
		* For indices < breakpoint.index all bitplanes < breakpoint.core_bitplane are ignored.
		* For indices >= breakpoint.index one more bitplane is ignored.
		*/
		template<typename itype>
		void update_core_slice_norms(const std::vector<itype>& quantized_core, double scale, const SplitBreakpoint& breakpoint, const std::vector<size_t>& block_bounds, bool use_quantized_data = true);
		// Prints various statistics about a quantized and encoded array (normally only used for the core).
		void print_encoded_array_statistics(QuantizationTarget target, size_t total_bits, size_t raw_bits, size_t sign_bits, size_t rle_bits, size_t leading_bits, uint8_t mode = 0);
		/*
		 * Writes leading bit positions to the output file (used for research).
		 * The values written are the positions (in uint8_t) of the last leading zero (so in the range 0, 1, ..., 64).
		 */
		template<typename itype>
		void write_leading_bit_positions(const std::vector<itype>& quantized_data, const SplitBreakpoint& breakpoint, const std::vector<size_t>& block_bounds);
		// Calls the corresponding dequantize<itype> depending on options.int_type.
		void dequantize(CompressedIO& io, QuantizationTarget target, uint8_t mode = 0);
		/*
		 * Decodes and dequantizes the core/factors
		 * If is_core is true, also reads compressed core sizes and core slice norms.
		 */
		template<typename itype>
		void dequantize(CompressedIO& io, QuantizationTarget target, uint8_t mode = 0);
		// Inverts weighting applied during compression and clear weighted factors.
		void unweight_factors();
		
		// Reordering (quantization) methods
		
		// Quantizes a single floating-point value
		template<typename itype>
		itype quantize_value(const ftype& val, double scale);
		/*
		 * Reorders and quantizes coefficients from floating-point array to quantized array.
		 * In the case of factors, no reordering happens. In the case of the core, reordering may happen (depends on core_flattening option).
		 */
		template<typename itype>
		std::vector<itype> reorder_and_quantize_array(const std::vector<ftype>& data, double scale, bool is_factor = false);
		/*
		 * Dequantizes a single value.
		 * quantized_index, breakpoint and block are needed to determine the amount of encoded bitplanes at this index.
		 * io is needed to extract a sign bit in case the coefficient is significant.
		 */
		template<typename itype>
		ftype dequantize_value(CompressedIO& io, itype val, double scale, size_t quantized_index, const SplitBreakpoint& breakpoint, size_t block);
		/*
		 * Reorders and dequantizes coefficients from quantized_data to a floating-point array and returns it.
		 * Also consumes bits from CompressedIO while reading signs.
		 */
		template<typename itype>
		std::vector<ftype> reorder_and_dequantize_array(CompressedIO& io, const std::vector<itype>& quantized_data, double scale, const SplitBreakpoint& breakpoint, const std::vector<size_t>& block_bounds, bool is_factor = false, size_t size = 0);
		
		// Other methods
		
		/*
		 * Resizes the core to the requested sizes (in-place if possible).
		 * We assume either all new sizes are larger or all new sizes are smaller than their current counterparts.
		 * We also assume that the allocated memory for the core is large enough to hold the new core.
		 * If the core increases in size, the new padded elements are set to 0.
		 * If the core decreases in size, the truncated elements are discarded (so they remain in memory but should not be used).
		 */
		void resize_core(const std::vector<size_t>& current_sizes, const std::vector<size_t>& new_sizes, const std::string& timer_label);
		/*
		 * Takes the factors after unweighting and merges them with an optional slicing filter matrix and an optional DCT matrix depending on the situation.
		 * The resulting factors can then be used during ST-HOSVD decompression to effectively perform all these operations simultaenously.
		 */
		void finalize_decompression_factors();
		// Prints flattened core, only used for debugging
		void print_core();
		
	};
	
}
