/*
 * This file defines the class atc::CompressedIO, declares its members
 * and also define its template methods, because they cannot be implemented in the source file.
 */

#pragma once

#include <cstdint>
#include <cstring>
#include <string>
#include <iostream>
#include <fstream>
#include <vector>
#include <map>

#include "IOMode.hpp"
#include "Options.hpp"
#include "LittleEndianIO.hpp"

namespace atc {
	
	/*
	 * This class manages the compressed output stream and provides various utility functions to easily encode/decode certain information.
	 */
	class CompressedIO {
	public:
		
		// General
		
		void enable_bit_mode(); // Changes the state of the object to bit mode, i.e. individual bits can be written/read to be packed together in the output/input stream.
		/*
		 * Changes the state of the object back to byte mode.
		 * In the case of writing, the buffer containing the last stored bits will be written,
		 * in the case of reading the buffer containing the last read bits will be discarded.
		 */
		void disable_bit_mode(); 
		void close(); // Close the current read/write file/memory stream.
		size_t get_bytes_written(); // Returns the size of the written compressed stream in bytes. Should be used after writing and calling close.
		
		// Debugging
		// These methods only exist when NDEBUG is defined, because implementing them causes a performance overhead
		#ifndef NDEBUG
			int8_t get_bit_pos() const; // Gets the current read/write position in the bit buffer
			size_t get_byte_pos() const; // Gets the current amount of bytes read/written
		#endif
		
		// Entropy coding helpers
		
		// Converts an array of symbols to a map containing the key-value pairs (symbol: frequency (i.e. amount of occurrences)).
		static std::map<uint64_t, uint64_t> calculate_frequency_table(const std::vector<uint64_t>& data);
		// Merges a series of frequency tables (maps) into a single frequency table, containing key-value pairs (symbol: sum of frequencies for this symbol over all frequency tables).
		static std::map<uint64_t, uint64_t> merge_frequency_tables(const std::vector<std::map<uint64_t, uint64_t>>& frequency_tables);
		/*
		 * Constructs a map (symbol: (range_start, scaled_frequency)),
		 * where the frequencies are scaled such that their sum equals a particular value, depending on the entropy coding method.
		 */
		static std::map<uint64_t, std::pair<uint64_t, uint64_t>> calculate_encoding_frequency_map(const std::map<uint64_t, uint64_t>& frequencies, EntropyCodingMethod entropy_coding_method);
		/*
		 * Constructs a map (range_start -> (symbol, scaled_frequency)),
		 * where the frequencies are scaled such that their sum equals a particular value, depending on the entropy coding method.
		 * Also adds a sentinel element at the end to make queries for upper bounds easier.
		 */
		static std::map<uint64_t, std::pair<uint64_t, uint64_t>> calculate_decoding_frequency_map(const std::map<uint64_t, uint64_t>& frequencies, EntropyCodingMethod entropy_coding_method);
		
		// WRITING
		
		void open_write(const std::string& output_file); // Opens a writing stream in File mode.
		void open_write(); // Opens a writing stream in Memory mode.
		/*
		 * Moves the stored compressed buffer and returns it.
		 * Used when done with in-memory compression to retrieve the compressed output.
		 */
		std::vector<char> get_and_clear_buffer();
		
		// Byte mode
		
		/*
		 * Writes a value while in byte mode.
		 * Is also used internally while in bit mode to write the full bit buffer when it is full,
		 * but should not be used externally during bit mode.
		 */
		template<typename T>
		void write(const T val);
		// Writes a sequence of values while in byte mode.
		template<typename T>
		void write(const std::vector<T>& data);
		
		// Bit mode
		
		// Writes a single bit in bit mode.
		void write_bit(bool bit);
		// Writes a primitive value (1-8 bytes) while in bit mode.
		template<typename T>
		void write_bits(const T bits);
		// Encodes a frequency table into the compressed output stream.
		size_t write_frequency_table(const std::map<uint64_t, uint64_t>& frequencies);
		/*
		 * Encodes a sequence of symbols using a custom frequency map (based on the actual symbol frequencies) and a given entropy coding method.
		 * Returns the number of bits used.
		 */
		size_t entropy_encode(const std::map<uint64_t, std::pair<uint64_t, uint64_t>>& encoding_frequency_map, const std::vector<uint64_t>& data, EntropyCodingMethod entropy_coding_method);
		
		// READING
		
		void open_read(const std::string& input_file); // Opens a reading stream in File mode.
		void open_read(const char* input_data, size_t size); // Opens a reading stream in Memory mode. Ensures that data is never read beyond the indicated bound.
		
		// Byte mode
		
		/*
		 * Reads a value while in byte mode.
		 * Is also used internally while in bit mode when the bit buffer is empty,
		 * but should not be used externally during bit mode.
		 */
		template<typename T>
		T read();
		// Reads a sequence of values while in byte mode
		template<typename T>
		std::vector<T> read(size_t amount);
		
		// Bit mode
		
		// Reads a single bit while in bit mode.
		bool read_bit();
		// Reads a primitive value (1-8 bytes) while in bit mode
		template<typename T>
		T read_bits();
		// Decodes a frequency table from the compressed output stream.
		std::map<uint64_t, uint64_t> read_frequency_table();
		// Decodes a custom frequencies map and entropy-coded bitstream and returns the array of decompressed symbols.
		std::vector<uint64_t> entropy_decode(const std::map<uint64_t, std::pair<uint64_t, uint64_t>>& decoding_frequency_map, EntropyCodingMethod entropy_coding_method);
		
	private:
		
		// General
		enum class State {Neutral, Write, Read}; // Represents the current state of the CompressedIO object
		IOMode mode; // Stores whether the data source/destination is in a file or in-memory
		State state = State::Neutral; // Stores the current read/write/neutral state of this object
		std::fstream stream; // The file input/output stream, used in File mode
		size_t bytes_written; // The amount of bytes written to the compressed output so far
		std::string output_file; // The output file when writing in File mode, needs to be stored to retrieve bytes_written.
		size_t byte_pos; // Only tracked used if NDEBUG is not defined, to avoid performance loss
		
		// Buffering
		std::vector<char> write_buffer; // Used to store the compressed output buffer when writing in Memory mode.
		const char* read_buffer; // Stores the data pointer given in open_read(const char*, size_t) when reading in Memory mode.
		size_t read_buffer_size; // Stores the given size of the read buffer when reading in Memory mode.
		size_t read_buffer_index; // Stores the current reading index in the read buffer when reading in Memory mode.
		
		// Bits
		bool bit_mode = false; // When bit mode is true, we read/write at the level of bis. Otherwise, at the level of bytes.
		uint64_t bit_buffer; // Used to pack bits together for reading/writing in bit mode, always 64-bit.
		int8_t bit_pos; // The next position in the bit buffer to be written to/read from
		void write_bits_internal(uint64_t bits, uint8_t amount); // Writes the data stored in the amount least-significant bits of bits
		uint64_t read_bits_internal(uint8_t amount); // Reads amount bits from the bit stream and returns them in the least-significant positions of the returned integer.
		
		// Entropy coding
		
		// AC constants
		static const uint8_t AC_CODE_VALUE_BITS = 32;
		static const uint64_t AC_MAX_CODE = (1UL << AC_CODE_VALUE_BITS) - 1;
		static const uint64_t AC_ONE_FOURTH = (AC_MAX_CODE + 1UL)/4;
		static const uint64_t AC_ONE_HALF = AC_ONE_FOURTH*2;
		static const uint64_t AC_THREE_FOURTHS = AC_ONE_FOURTH*3;
		
		// ANS constants
		static const uint8_t ANS_LOG_L = 32; // log2 of ANS_L
		static const uint64_t ANS_L = uint64_t(1) << ANS_LOG_L; // Lower bound on the valid state interval
		static const uint8_t ANS_LOG_B = 8; // log2 of ANS_B
		static const uint64_t ANS_B = uint64_t(1) << ANS_LOG_B; // Ratio in between upper and lower bound of the valid state interval
		
		size_t encoding_bits; // Counts the amount of bits used for entropy coding during one entropy_encode call.
		size_t pending_bits; // Stores the amount of pending bits, which are used in the implementation of arithmetic coding.
		
		/*
		 * Encodes the given value using a simple compression scheme (storing length in bits + bits without leading zeroes).
		 * Used e.g. when encoding frequency maps to save bits (involves storing a lot of small values).
		 */
		void write_bits_dynamic_length(uint64_t val);
		// Decodes the given value using a simple decompression scheme (storing length in bits + bits without leading zeroes).
		uint64_t read_bits_dynamic_length();
		void write_bit_and_pending(bool bit); // Writes bit, followed by pending_bits times ~bit and then resets pending_bits to 0. Used in the AC implementation.
		void arithmetic_encode(const std::map<uint64_t, std::pair<uint64_t, uint64_t>>& encoding_frequency_map, const std::vector<uint64_t>& data); // Perform entropy coding with AC.
		std::vector<uint64_t> arithmetic_decode(const std::map<uint64_t, std::pair<uint64_t, uint64_t>>& decoding_frequency_map, uint64_t symbols_amount); // Perform entropy decoding with AC.
		void ANS_encode(const std::map<uint64_t, std::pair<uint64_t, uint64_t>>& encoding_frequency_map, const std::vector<uint64_t>& data); // Perform entropy coding with ANS.
		std::vector<uint64_t> ANS_decode(const std::map<uint64_t, std::pair<uint64_t, uint64_t>>& decoding_frequency_map, uint64_t symbols_amount); // Perform entropy decoding with ANS.
		
	};
	
	// Implementation of template methods
	
	template<typename T>
	void CompressedIO::write(T val) {
		size_t size = sizeof(T);
		switch (mode) {
			case IOMode::File:
				little_endian_file_write(stream, &val, 1);
				break;
			case IOMode::Memory:
				size_t current_size = write_buffer.size();
				write_buffer.resize(current_size + size);
				little_endian_memory_write(write_buffer.data(), current_size, &val, 1);
				break;
		}
		#ifndef NDEBUG
			byte_pos += size;
		#endif
	}
	
	template<typename T>
	void CompressedIO::write(const std::vector<T>& data) {
		size_t size = sizeof(T)*data.size();
		switch (mode) {
			case IOMode::File:
				little_endian_file_write(stream, data.data(), data.size());
				break;
			case IOMode::Memory:
				size_t current_size = write_buffer.size();
				write_buffer.resize(current_size + size);
				little_endian_memory_write(write_buffer.data(), current_size, data.data(), data.size());
				break;
		}
		#ifndef NDEBUG
			byte_pos += size;
		#endif
	}
	
	template<typename T>
	void CompressedIO::write_bits(const T val) {
		// Iterate over bytes as type and reinterpret each byte as an unsigned char
		#ifdef ATC_BIG_ENDIAN
		for (uint8_t pos = sizeof(T); pos-- > 0; )
		#else
		for (uint8_t pos = 0; pos < sizeof(T); pos++)
		#endif
			write_bits_internal(uint64_t(reinterpret_cast<const unsigned char*>(&val)[pos]), 8);
	}
	
	template<typename T>
	T CompressedIO::read() {
		T val;
		size_t size = sizeof(T);
		switch (mode) {
			case IOMode::File:
				little_endian_file_read(stream, &val, 1);
				if (stream.eof() && stream.fail())
					throw std::runtime_error("Read past the end of compressed file");
				break;
			case IOMode::Memory:
				if (read_buffer_index + size > read_buffer_size)
					throw std::runtime_error("Read past the end of compressed data");
				little_endian_memory_read(read_buffer, read_buffer_index, &val, 1);
				read_buffer_index += size;
				break;
		}
		#ifndef NDEBUG
			byte_pos += size;
		#endif
		return val;
	}
	
	template<typename T>
	std::vector<T> CompressedIO::read(size_t amount) {
		std::vector<T> vals(amount);
		size_t size = sizeof(T)*amount;
		switch (mode) {
			case IOMode::File:
				little_endian_file_read(stream, vals.data(), amount);
				if (stream.eof() && stream.fail())
					throw std::runtime_error("Read past the end of compressed file");
				break;
			case IOMode::Memory:
				if (read_buffer_index + size > read_buffer_size)
					throw std::runtime_error("Read past the end of compressed data");
				little_endian_memory_read(read_buffer, read_buffer_index, vals.data(), amount);
				read_buffer_index += size;
				break;
		}
		#ifndef NDEBUG
			byte_pos += size;
		#endif
		return vals;
	}
	
	template<typename T>
	T CompressedIO::read_bits() {
		T val;
		unsigned char* ptr = reinterpret_cast<unsigned char*>(&val);
		#ifdef ATC_BIG_ENDIAN
		for (uint8_t pos = sizeof(T); pos-- > 0; )
		#else
		for (uint8_t pos = 0; pos < sizeof(T); pos++)
		#endif
			ptr[pos] = static_cast<unsigned char>(read_bits_internal(8));
		return val;
	}
	
}
