/*
 * This file handles the decompression aspect of the quantization and encoding phase of the pipeline.
 */

#include <sstream>
#include <bitset>
#include <algorithm>
#include <iomanip>
#include <numeric>

#ifdef ATC_OPENMP
	#include <omp.h>
#endif

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

namespace atc {
	
	// Wrapper method to call the corresponding template method
	template<typename ftype>
	void Compressor<ftype>::dequantize(CompressedIO& io, QuantizationTarget target, uint8_t mode) {
		switch (options.int_type) {
			case IntType::uint32:
				dequantize<uint32_t>(io, target, mode);
				break;
			case IntType::uint64:
				dequantize<uint64_t>(io, target, mode);
				break;
		}
	}
	
	template<typename ftype>
	template<typename itype>
	void Compressor<ftype>::dequantize(CompressedIO& io, QuantizationTarget target, uint8_t mode) {
		
		// Construct string used for timers
		std::stringstream ss;
		ss << "decompress_dequantize";
		if (target == QuantizationTarget::Core)
			ss << "_core";
		else if (target == QuantizationTarget::Factor)
			ss << "_factor_" << mode_to_string(mode);
		std::string timer_string = ss.str();
		
		start_timer(timer_string);
		
		// Calculate array size
		uint64_t size;
		if (target == QuantizationTarget::Core) {
			size = product(svd_sizes);
		} else {
			size_t rows = dct_sizes[mode];
			size_t cols = quantized_sizes[mode];
			size = options.householder ? cols*(2*rows - cols - 1)/2 : rows*cols;
		}
		
		// Read scale
		double scale = io.read<double>();
		
		// Start main decoding loop
		start_timer(timer_string + "_decode");
		std::vector<itype> decoded_data(size);
		double unscramble_time = 0;
		double raw_time = 0;
		double entropy_time = 0;
		bool all_raw = false;
		size_t blocks = (target == QuantizationTarget::Core ? std::max(size_t(io.read<uint64_t>()), size_t(1)) : 1);
		std::vector<size_t> block_bounds = calculate_block_bounds(size, blocks);
		SplitBreakpoint breakpoint(blocks);
		#ifdef ATC_OPENMP
			omp_set_num_threads(min(blocks, size_t(options.threads)));
		#endif
		print_compressed_io_pos(io);
		for(breakpoint.bitplane = 8*sizeof(itype) - 1; ; breakpoint.bitplane--) {
			
			// Preliminaries
			if (target != QuantizationTarget::Factor)
				print_status("Decoding core bit plane " + std::to_string(uint64_t(breakpoint.bitplane)));
			
			bool write_rle = io.read<uint8_t>();
			std::map<uint64_t, std::pair<uint64_t, uint64_t>> decoding_frequency_map;
			if (write_rle) {
				io.enable_bit_mode();
				auto frequencies = io.read_frequency_table();
				io.disable_bit_mode();
				decoding_frequency_map = CompressedIO::calculate_decoding_frequency_map(frequencies, options.entropy_coding_method);
			}
			
			// First sequentially extract all compressed block data from the stream
			std::vector<std::vector<char>> compressed_blocks(blocks);
			for (size_t block = 0; block < blocks; block++) {
				size_t compressed_block_size = io.read<uint64_t>();
				compressed_blocks[block] = io.read<char>(compressed_block_size);
			}
			
			// Decode each block
			#pragma omp parallel for reduction(+:unscramble_time,raw_time,entropy_time)
			for (size_t block = 0; block < blocks; block++) {
				
				CompressedIO block_io;
				block_io.open_read(compressed_blocks[block].data(), compressed_blocks[block].size());
				block_io.enable_bit_mode();
				size_t start_index = block_bounds[block];
				
				// Regardless of all_raw, the bitplane decoding will start with this
				uint64_t raw_size = block_io.read_bits<uint64_t>();
				if (raw_size > block_bounds[block + 1] - block_bounds[block])
					throw std::runtime_error("Decoded a raw bits array size larger than the corresponding part of the quantized array.");
				
				size_t leading_index, trailing_index;
				if (all_raw) {
					// All raw, so only read raw bits
					auto time = time_now();
					size_t index = start_index;
					for (; index < start_index + raw_size; index++)
						decoded_data[index] |= itype(block_io.read_bit()) << breakpoint.bitplane;
					leading_index = index;
					trailing_index = index;
					unscramble_time += elapsed_time(time);
				} else {
					
					auto time = time_now();
					std::vector<bool> raw(raw_size);
					for (uint64_t i = 0; i < raw_size; i++)
						raw[i] = block_io.read_bit();
					raw_time += elapsed_time(time);
					
					time = time_now();
					std::vector<uint64_t> runlengths;
					if (write_rle) {
						// RLE bits were encoded for this bitplane
						runlengths = block_io.entropy_decode(decoding_frequency_map, options.entropy_coding_method);
					}
					entropy_time += elapsed_time(time);
					
					// If rle_index < rle.size(), there are still run-lengths to read from the buffer
					// If rle_index == rle.size(), there are no more run-lengths left in the buffer but the last one is being processed
					// If rle_index > rle.size(), there are no more run-lengths to be processed
					size_t rle_index = 0;
					size_t rle_value = (rle_index < runlengths.size() ? runlengths[rle_index] : 0);
					rle_index++;
					size_t raw_index = 0;
					leading_index = start_index;
					trailing_index = start_index;
					
					// Unscramble raw and RLE data, which may cover different ranges
					time = time_now();
					for (size_t index = start_index; index < block_bounds[block + 1]; index++) {
						
						bool current_bit = false;
						
						if (decoded_data[index] == 0) {
							// Handle leading bit
							if (rle_value > 0) {
								rle_value--;
							} else {
								if (rle_index <= runlengths.size()) {
									current_bit = true;
									if (rle_index < runlengths.size()) {
										rle_value = runlengths[rle_index];
									} else {
										// Reached leading bit breakpoint
										leading_index = index + 1;
									}
									rle_index++;
								} else if (raw_index >= raw.size()) {
									// No trailing nor leading bit left to write, so stop
									break;
								}
							}
						} else {
							// Handle trailing bit
							if (raw_index < raw.size()) {
								// Still raw bits left, so write trailing bit
								current_bit = raw[raw_index++];
								if (raw_index == raw.size()) {
									// Reached trailing bit breakpoint
									trailing_index = index + 1;
								}
							} else if (rle_index > runlengths.size()) {
								// No trailing nor leading bit left to write, so stop
								break;
							}
						}
						
						if (current_bit)
							decoded_data[index] |= itype(1) << breakpoint.bitplane;
						
					}
					unscramble_time += elapsed_time(time);
				}
				
				breakpoint.blocks[block].leading_index = leading_index;
				breakpoint.blocks[block].trailing_index = trailing_index;
				
				block_io.disable_bit_mode();
				block_io.close();
				compressed_blocks[block].clear();
				
			}
			
			if (breakpoint.bitplane > 0) {
				// Check all_raw bit
				if (!all_raw)
					all_raw = io.read<bool>();
				// Check if done bit is true
				if (io.read<bool>())
					break;
			} else {
				break;
			}
			
		}
		stop_timer();
		print_compressed_io_pos(io);
		if (options.timings) {
			results.timings[timer_string + "_decode_unscramble"] = unscramble_time;
			results.timings[timer_string + "_decode_raw"] = raw_time;
			results.timings[timer_string + "_decode_entropy"] = entropy_time;
			if (options.verbose && target == QuantizationTarget::Core) {
				print_info("Core dequantization loop timings (sum of real-time for each thread):");
				print_info("- Unscrambling time: " + std::to_string(unscramble_time) + " s");
				print_info("- Raw decoding time: " + std::to_string(raw_time) + " s");
				print_info("- Entropy decoding time: " + std::to_string(entropy_time) + " s");
			}
		}
		
		if (target == QuantizationTarget::Core) {
			// The printed breakpoint indices may be smaller than the ones printed during compression,
			// since these are the tightest bounds such that the same dequantization correction is still applied
			print_breakpoint_info(breakpoint, block_bounds, false);
		} else {
			// No need to print breakpoint index for factors
			print_info("Factor " + mode_to_string(mode) + " decoding breakpoint: bitplane " + std::to_string(uint64_t(breakpoint.bitplane)));
		}
		
		// If we're dequantizing the core, first read core slice norms and compressed core ranks
		if (target == QuantizationTarget::Core) {
			quantized_sizes = options.use_quantized_core_slice_norms ? read_sizes(io) : svd_sizes;
			print_info("Quantized core sizes: " + sizes_to_string(quantized_sizes));
			core_slice_norms = std::vector<std::vector<double>>(modes);
			for (uint8_t mode = 0; mode < modes; mode++) {
				core_slice_norms[mode] = io.read<double>(quantized_sizes[mode]);
				if (*std::min(core_slice_norms[mode].begin(), core_slice_norms[mode].end()) < 0)
					throw std::runtime_error("Decoded a negative core slice norm");
			}
			print_compressed_io_pos(io);
		}
		
		// Dequantize to target while reading necessary signs from compressed IO
		start_timer(timer_string + "_correct_reorder_and_scale");
		io.enable_bit_mode();
		(target == QuantizationTarget::Factor ? weighted_factors[mode] : core)
			= reorder_and_dequantize_array(io, decoded_data, scale, breakpoint, block_bounds, target == QuantizationTarget::Factor, target == QuantizationTarget::Factor ? 0 : calculate_required_size_core(options.apply_dct));
		io.disable_bit_mode();
		print_compressed_io_pos(io);
		stop_timer();
		
		stop_timer();
		
	}
	
	template<typename ftype>
	template<typename itype>
	ftype Compressor<ftype>::dequantize_value(CompressedIO& io, itype val, double scale, size_t quantized_index, const SplitBreakpoint& breakpoint, size_t block) {
		if (val > 0) {
			if (quantized_index < breakpoint.blocks[block].trailing_index || ((val >> breakpoint.bitplane) >> 1) == 0) {
				val |= (itype(1) << breakpoint.bitplane) >> 1;
			} else {
				val |= itype(1) << breakpoint.bitplane;
			}
			return (io.read_bit() ? 1 : -1)*ftype(val)/scale;
		} else {
			return 0;
		}
	}
	
	template<typename ftype>
	template<typename itype>
	std::vector<ftype> Compressor<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, size_t size) {
		std::vector<ftype> dequantized_data(std::max(quantized_data.size(), size));
		QuantizedArrayIterator it(is_factor, svd_sizes, quantized_data.size(), current_mode_order, options.core_flattening);
		for (size_t block = 0; block < breakpoint.blocks.size(); block++) {
			for (; it.get_reordered_index() < block_bounds[block + 1]; it.advance())
				dequantized_data[it.get_original_index()] = dequantize_value(io, quantized_data[it.get_reordered_index()], scale, it.get_reordered_index(), breakpoint, block);
		}
		return dequantized_data;
	}
	
	template<typename ftype>
	void Compressor<ftype>::unweight_factors() {
		
		factors.resize(modes);
		for (uint8_t mode = 0; mode < modes; mode++) {
			size_t rows = dct_sizes[mode];
			size_t cols = quantized_sizes[mode];
			factors[mode] = mtype::Zero(rows, cols);
			if (options.householder) {
				// Arrange scaled coefficients into a single long vector
				size_t total_size = cols*(2*rows - cols - 1)/2;
				size_t index = total_size;
				Eigen::Map<vtype> weighted_factor_vector(weighted_factors[mode].data(), total_size);
				weighted_factor_vector.array() /= calculate_householder_weights(mode, false).array();
				for (size_t col = std::min(cols, rows - 1); col-- > 0;) {
					// Process columns from last to first, adding core slice norms as we go
					size_t reflector_length = rows - col - 1;
					index -= reflector_length;
					factors[mode].block(col + 1, col, reflector_length, 1) = weighted_factor_vector.segment(index, reflector_length);
				}
			} else {
				Eigen::Map<mtype> weighted_factor_matrix(weighted_factors[mode].data(), rows, cols);
				for (size_t col = 0; col < cols; col++) {
					double norm = core_slice_norms[mode][col];
					// We tolerate norms of zero because we only truncate away all contiguous zero norms at the end
					// In the case of non-monotonously decreasing norms this may not include all zero norms
					if (norm < 0)
						throw std::runtime_error("Decoded a negative core slice norm");
					factors[mode].col(col) = weighted_factor_matrix.col(col)/(norm > 0 ? norm : 1);
				}
			}
		}
		force_clear(weighted_factors);
		
	}
	
	#include "explicit_instantiations/Compressor.txt"
	
}

