/*
 * This file implements the methods of atc::CompressorBase.
 */

#include <vector>
#include <cstdint>
#include <memory>
#include <Eigen/Core>

#include "CompressorBase.hpp"
#include "util.hpp"
#include "version.hpp"
#include "Breakpoint.hpp"

namespace atc {
	
	// data_type_id needs to be passed like this (stored in a member variable) because it RawIO depends on ftype and cannot be passed as an argument
	void CompressorBase::write_header(CompressedIO& io) {
		
		// Open IO and write header
		io.write(uint8_t(atc::version.length()));
		io.write(std::vector<char>(atc::version.begin(), atc::version.end()));
		
		// Construct flags bitfield
		flag_type flags = options.rank_truncation*flag_truncate_st_hosvd
			| (options.int_type == IntType::uint64)*flag_long_bitplane_encoding
			| (options.storage_order == StorageOrder::MSI)*flag_msi_storage_order
			| options.use_quantized_core_slice_norms*flag_use_quantized_core_slice_norms
			| options.householder*flag_householder
			| (options.entropy_coding_method == EntropyCodingMethod::ANS)*flag_ANS
			| options.apply_dct*flag_dct
			| (options.core_flattening == CoreFlattening::Zigzag)*flag_zigzag_flattening
			| (options.core_flattening == CoreFlattening::ZOrder)*flag_z_order_flattening
			| (options.householder_weighting_method == HouseholderWeightingMethod::CoreSliceNorms)*flag_hwm_core_slice_norms
		;
		io.write(flags);
		
		io.write(data_type_id);
		io.write(modes);
		write_sizes(io, original_sizes);
		if (options.rank_truncation) {
			if (options.apply_dct)
				write_sizes(io, dct_sizes);
			write_sizes(io, svd_sizes);
		}
		io.write(current_mode_order);
		
	}
	
	void CompressorBase::read_header(CompressedIO& io) {
		
		// Read version
		size_t input_version_length = io.read<uint8_t>();
		std::vector<char> input_version_vector = io.read<char>(input_version_length);
		std::string input_version(input_version_vector.begin(), input_version_vector.end());
		if (input_version != atc::version)
			warn("Compressed file was compressed with version " + input_version + " of ATC and may not be interpreted correctly by this version (" + atc::version + ")");
		
		// Read flags
		flag_type flags = io.read<flag_type>();
		options.rank_truncation = flags & flag_truncate_st_hosvd;
		options.int_type = (flags & flag_long_bitplane_encoding) ? IntType::uint64 : IntType::uint32;
		options.storage_order = (flags & flag_msi_storage_order) ? StorageOrder::MSI : StorageOrder::LSI;
		options.use_quantized_core_slice_norms = flags & flag_use_quantized_core_slice_norms;
		options.householder = flags & flag_householder;
		options.entropy_coding_method = (flags & flag_ANS) ? EntropyCodingMethod::ANS : EntropyCodingMethod::Arithmetic;
		options.apply_dct = flags & flag_dct;
		options.core_flattening = (flags & flag_z_order_flattening) ? CoreFlattening::ZOrder : ((flags & flag_zigzag_flattening) ? CoreFlattening::Zigzag : CoreFlattening::Unfold);
		options.householder_weighting_method = ((flags & flag_hwm_core_slice_norms) ? HouseholderWeightingMethod::CoreSliceNorms : HouseholderWeightingMethod::Approximation);
		
		// Read rest of header
		data_type_id = io.read<uint8_t>();
		modes = io.read<uint8_t>();
		original_sizes = read_sizes(io);
		
		// Read truncated sizes if necessary
		dct_sizes = (options.apply_dct && options.rank_truncation) ? read_sizes(io) : original_sizes;
		svd_sizes = options.rank_truncation ? read_sizes(io) : original_sizes;
		
		// Read current mode order of core
		current_mode_order = io.read<uint8_t>(modes);
		
	}
	
 	DataType CompressorBase::read_data_type(CompressedIO& compressed_io) {
		read_header(compressed_io);
		return id_to_data_type(data_type_id);
	}
	
	size_t CompressorBase::read_decompressed_size(CompressedIO& compressed_io, const Options& options, bool suppress_autocrop_warnings) {
		read_header(compressed_io);
		return calculate_decompressed_size(original_sizes, options, suppress_autocrop_warnings);
	}
	
	size_t CompressorBase::calculate_decompressed_size(const std::vector<size_t>& sizes, const Options& options, bool suppress_autocrop_warnings) {
		if (options.autocrop ) {
			if (options.cutout.size() > 0)
				throw std::invalid_argument("Cannot use autocrop when manually slicing the tensor too");
			if (!suppress_autocrop_warnings)
				warn("Decompressed sizes cannot be pre-computed exactly with autocrop, using upper bound (full decompressed data size) instead.");
		}
		size_t res = 1;
		std::vector<Slice> cutout = options.cutout;
		cutout.resize(sizes.size());
		if (options.storage_order == StorageOrder::MSI)
			std::reverse(cutout.begin(), cutout.end());
		for (uint8_t mode = 0; mode < sizes.size(); mode++) {
			cutout[mode].update_mode_size(sizes[mode]);
			res *= cutout[mode].get_size();
		}
		return res;
	}
	
	void CompressorBase::verify_slices() {
		for(uint8_t mode = 0; mode < modes; mode++) {
			const Slice& slice = options.cutout[mode];
			if (!slice.is_standard()) {
				if (options.calculate_error)
					throw std::invalid_argument("Cannot calculate error if decompressed tensor is sliced");
				if (slice.get_size() == 0)
					throw std::invalid_argument("Slice has size 0");
			}
		}
	}
	
	// Approximates the computational cost of performing a tensor-matrix mode-k product along each mode, using a specific mode processing order
	uint64_t mode_processing_order_cost(const std::vector<size_t>& current_sizes, const std::vector<size_t>& new_sizes, const std::vector<uint8_t>& mode_processing_order) {
		size_t size = product(current_sizes);
		uint64_t cost = 0;
		uint8_t modes = current_sizes.size();
		for (uint8_t mode_index = 0; mode_index < modes; mode_index++) {
			uint8_t mode = mode_processing_order[mode_index];
			// Matrix multiplication of sizes (new_sizes[mode], current_sizes[mode]) and (current_sizes[mode], size/current_sizes[mode])
			cost += size*new_sizes[mode];
			size /= current_sizes[mode];
			size *= new_sizes[mode];
		}
		return cost;
	}
	
	/**
	 * We perform an exhaustive search for simplicity.
	 * The optimal ordering can actually by found in O(d*log(d)) time (with d = number of modes),
	 * as described in section 8 of the TuckerMPI paper (Ballard et al. 2020),
	 * however this is not really important because the amount of modes should be very low.
	 */
	void CompressorBase::update_decompression_mode_processing_order() {
		std::vector<uint8_t> mode_processing_order(modes);
		std::iota(mode_processing_order.begin(), mode_processing_order.end(), 0);
		if (decompression_mode_processing_order.size() == 0) {
			// Decompression mode processing order has not been computed yet, initialize minimum cost to the highest value possible
			decompression_mode_processing_order_cost = std::numeric_limits<size_t>::max();
		}
		std::vector<size_t> new_sizes(modes);
		for (uint8_t mode = 0; mode < modes; mode++)
			new_sizes[mode] = (options.cutout[mode].is_standard() ? dct_sizes[mode] : options.cutout[mode].get_size());
		do {
			size_t cost = mode_processing_order_cost(quantized_sizes, new_sizes, mode_processing_order);
			if (cost < decompression_mode_processing_order_cost) {
				decompression_mode_processing_order = mode_processing_order;
				decompression_mode_processing_order_cost = cost;
			}
		} while (std::next_permutation(mode_processing_order.begin(), mode_processing_order.end()));
	}
	
	size_t CompressorBase::calculate_required_size_core(bool include_hodct) {
		
		// Simulate ST-HOSVD expansion/shrinking due to slicing
		size_t current_size = product(quantized_sizes);
		size_t max_size = current_size;
		update_decompression_mode_processing_order();
		for(uint8_t mode_index = 0; mode_index < modes; mode_index++) {
			uint8_t mode = decompression_mode_processing_order[mode_index];
			const Slice& slice = options.cutout[mode];
			current_size = current_size/quantized_sizes[mode]*(slice.is_standard() ? dct_sizes[mode] : slice.get_size());
			max_size = std::max(max_size, current_size);
		}
		
		// Simulate HODCT padding
		if (include_hodct) {
			for(uint8_t mode = 0; mode < modes; mode++) {
				const Slice& slice = options.cutout[mode];
				// Only apply DCT's during HODCT to modes with no special slicing
				// If the slice is non-standard, the DCT was already included in the ST-HOSVD factor convolution
				if (slice.is_standard())
					current_size = current_size/dct_sizes[mode]*original_sizes[mode];
			}
			max_size = std::max(max_size, current_size);
		}
		
		return max_size;
		
	}
	
	// General methods
	
	void CompressorBase::print_banner(const std::string& message) {
		if (options.verbose) {
			const uint8_t total_space = 130;
			uint8_t spacing = total_space - message.length();
			uint8_t spacing_left = spacing/2;
			std::cout << std::string(spacing_left - 1, '=') << " " << message << " " << std::string(spacing - spacing_left - 1, '=') << std::endl;
		}
	}
	
	void CompressorBase::print_info(const std::string& message, bool newline) {
		print_tagged_message("Info", message, newline);
	}
	
	void CompressorBase::print_status(const std::string& message, bool newline) {
		print_tagged_message("Status", message, newline);
	}
	
	void CompressorBase::warn(const std::string& message) {
		if (!options.suppress_warnings)
			print_tagged_message("Warning", message);
	}
	
	void CompressorBase::print_tagged_message(const std::string& tag, const std::string& message, bool newline) {
		if (options.verbose || tag == "Warning") {
			const uint8_t tag_length = 12;
			std::cout << "[" << tag << "]" << std::string(tag_length - tag.length() - 2, ' ') << message;
			if (newline)
				std::cout << std::endl;
		}
	}
	
	void CompressorBase::print_breakpoint_info(const SplitBreakpoint& breakpoint, const std::vector<size_t>& block_bounds, bool encoding, double current_sse, double target_sse) {
		size_t blocks = breakpoint.blocks.size();
		size_t size = block_bounds[blocks];
		BlockBreakpoint cumulative_indices;
		for (size_t block = 0; block < blocks; block++) {
			cumulative_indices.leading_index += breakpoint.blocks[block].leading_index - block_bounds[block];
			cumulative_indices.trailing_index += breakpoint.blocks[block].trailing_index - block_bounds[block];
		}
		auto print_cumulative_index = [this, &size] (const std::string& name, size_t cumulative_index) {
			std::stringstream ss;
			ss << "- " << name << " bits index (cumulative): " << cumulative_index << "/" << size
				<< " (" << std::fixed << std::setprecision(1) << (100.0*cumulative_index/size) << "%)";
			this->print_info(ss.str());
		};
		print_info("Core " + std::string(encoding ? "en" : "de") + "coding breakpoint:");
		print_info("- Bitplane: " + std::to_string(uint64_t(breakpoint.bitplane)));
		print_cumulative_index("Leading", cumulative_indices.leading_index);
		print_cumulative_index("Trailing", cumulative_indices.trailing_index);
		if (current_sse != -1)
			print_info("- Quantization SSE/target quantization SSE: " + std::to_string(current_sse) + "/" + std::to_string(target_sse));
	}
	
	void CompressorBase::print_compressed_io_pos([[maybe_unused]] const CompressedIO& io) {
		#ifndef NDEBUG
			print_info("Current CompressedIO position (bytes, bits): (" + std::to_string(io.get_byte_pos()) + ", " + std::to_string(io.get_bit_pos()) + ")");
		#endif
	}
	
	void CompressorBase::verify_mode_permutation(const std::vector<uint8_t>& permutation, const std::string& name) {
		std::vector<uint8_t> default_mode_permutation(modes);
		std::iota(default_mode_permutation.begin(), default_mode_permutation.end(), 0);
		if (permutation.size() > 0) {
			if (permutation.size() != size_t(modes)) {
				throw std::invalid_argument("Given mode processing order is of length " + std::to_string(permutation.size())
					+ " which is different from the number of modes " + std::to_string(size_t(modes)));
			}
			if (!std::is_permutation(permutation.begin(), permutation.end(), default_mode_permutation.begin())) {
				// Print mode processing order in the same order as given
				std::stringstream ss;
				ss << "Given " << name << " (" << size_t(permutation[0]);
				for (uint8_t mode = 1; mode < modes; mode++)
					ss << ", " << size_t(permutation[mode]);
				ss << ") is not a valid permutation";
				throw std::invalid_argument(ss.str());
			}
		}
	}
	
	std::string CompressorBase::mode_to_string(uint8_t mode) {
		return std::to_string(int(options.storage_order == StorageOrder::MSI ? (modes - 1 - mode) : mode));
	}
	
	// Maps a tuple of T's to a string representing the tuple
	// F(T) must map T to a string
	template<typename T, typename F>
	std::string tuple_to_string(const std::vector<T>& tuple, bool reverse, const F& mapping) {
		std::stringstream ss;
		ss << "(";
		for (uint8_t i = (reverse ? tuple.size() : 0); reverse ? (i-- > 0) : i < tuple.size(); reverse ? 0 : i++) {
			ss << mapping(tuple[i]);
			if (reverse ? (i > 0) : (i < tuple.size() - 1))
				ss << ", ";
		}
		ss << ")";
		return ss.str();
	}
	
	// A wrapper around the above function which sets the mapping function to the default value std::to_string,
	// because it is very complicated in practice to set a default function argument for a templated argument.
	template<typename T>
	std::string tuple_to_string(const std::vector<T>& tuple, bool reverse) {
		return tuple_to_string(tuple, reverse, [] (T val) {return std::to_string(val);});
	}
	
	std::string CompressorBase::sizes_to_string(const std::vector<size_t>& sizes) {
		return tuple_to_string(sizes, options.storage_order == StorageOrder::MSI);
	}
	
	std::string CompressorBase::mode_order_to_string(const std::vector<uint8_t>& mode_order) {
		return tuple_to_string(mode_order, options.storage_order == StorageOrder::MSI, [this] (uint8_t mode) {return this->mode_to_string(mode);});
	}
	
	void CompressorBase::write_sizes(CompressedIO& io, const std::vector<size_t>& sizes) {
		for(const size_t size : sizes)
			io.write(uint32_t(size));
	}
	
	std::vector<size_t> CompressorBase::read_sizes(CompressedIO& io) {
		std::vector<size_t> sizes(modes);
		for(uint8_t mode = 0; mode < modes; mode++) {
			sizes[mode] = size_t(io.read<uint32_t>());
			if (sizes[mode] == 0)
				throw std::runtime_error("Read a mode size of 0 in compressed file");
			if (sizes[mode] > std::numeric_limits<int>::max())
				throw std::runtime_error("Read a mode of size > 2^31 in compressed file");
		}
		return sizes;
	}
	
	Results CompressorBase::get_and_clear_results() {
		return std::move(results);
	}
	
	// Starts labeled timer
	void CompressorBase::start_timer(std::string label) {
		print_status("Started  " + label); // Double space to align with "Finished "
		if (options.verbose || options.timings)
			timers.push(std::make_pair(label, time_now()));
	}
	
	// Takes last labeled timer from stack and prints elapsed time
	void CompressorBase::stop_timer() {
		if (options.verbose || options.timings) {
			
			// Pop timer off stack and calculate time period
			if (timers.size() < 1)
				throw std::out_of_range("No timers left on stack");
			std::string label = timers.top().first;
			double time = elapsed_time(timers.top().second);
			timers.pop();
			
			// Print label and timing
			print_status("Finished " + label, false);
			if (options.timings) {
				results.timings[label] = time;
				if (options.verbose)
					std::cout << ", elapsed time: " << time << " s";
			}
			if (options.verbose)
				std::cout << std::endl;
			
		}
	}
	
}
