/*
 * This file implements the class atc::CompressedIO.
 */

#include <fstream>
#include <stdexcept>
#include <utility>
#include <map>
#include <algorithm>
#include <iterator>
#include <cassert>
#include <cmath>
#include <stack>
#include <math.h>
#include <experimental/filesystem>

#include "CompressedIO.hpp"

namespace atc {
	
	// General methods
	
	void CompressedIO::close() {
		if (bit_mode)
			disable_bit_mode();
		if (state == State::Write) {
			switch (mode) {
				case IOMode::File:
					stream.close(); // Makes sure the file is written to disk, so the next call gives the right output
					bytes_written = std::experimental::filesystem::file_size(output_file);
					break;
				case IOMode::Memory:
					bytes_written = write_buffer.size();
					break;
			}
		}
		state = State::Neutral;
	}
	
	void CompressedIO::enable_bit_mode() {
		assert(!bit_mode);
		switch (state) {
		case State::Neutral:
			assert(false);
			break;
		case State::Write:
			bit_buffer = 0;
			bit_pos = 63;
			break;
		case State::Read:
			bit_buffer = 0;
			bit_pos = -1;
			break;
		}
		bit_mode = true;
	}
	
	void CompressedIO::disable_bit_mode() {
		assert(bit_mode);
		if (state == State::Write && bit_pos < 63) {
			// Write any remaining bits
			write(bit_buffer);
		}
		bit_mode = false;
	}
	
	size_t CompressedIO::get_bytes_written() {
		return bytes_written;
	}
	
	#ifndef NDEBUG
		int8_t CompressedIO::get_bit_pos() const {
			return (bit_pos + 64)%64; // Hacky way to get the "normal" modulo operator for negative numbers ...
		}
		
		size_t CompressedIO::get_byte_pos() const {
			// When reading in bit mode, 8 bytes will be extracted from the stream into the read buffer,
			// These should be compensated for in order to get the same positions as when writing
			return byte_pos - 8*int(bit_mode && state == State::Read && bit_pos >= 0);
		}
	#endif
	
	// Writing methods
	
	void CompressedIO::open_write(const std::string& output_file) {
		assert(state == State::Neutral);
		stream.open(output_file, std::ios::out | std::ios::binary);
		if (!stream.is_open())
			throw std::ios_base::failure("Could not open compressed file \"" + output_file + "\" for writing");
		mode = IOMode::File;
		state = State::Write;
		bit_mode = false;
		this->output_file = output_file;
		byte_pos = 0;
	}
	
	void CompressedIO::open_write() {
		assert(state == State::Neutral);
		mode = IOMode::Memory;
		state = State::Write;
		bit_mode = false;
		byte_pos = 0;
	}
	
	std::vector<char> CompressedIO::get_and_clear_buffer() {
		return std::move(write_buffer);
	}
	
	void CompressedIO::write_bit(bool bit) {
		if (bit_pos >= 1) {
			bit_buffer |= static_cast<uint64_t>(bit) << bit_pos;
			bit_pos--;
		} else {
			bit_buffer |= static_cast<uint64_t>(bit);
			write(bit_buffer);
			bit_buffer = 0;
			bit_pos = 63;
		}
	}
	
 	void CompressedIO::write_bits_internal(uint64_t bits, uint8_t amount) {
		assert(amount <= 64);
		if (amount <= bit_pos + 1) {
			bit_buffer |= bits << (bit_pos + 1 - amount);
			bit_pos -= amount;
			if (bit_pos < 0) {
				write(bit_buffer);
				bit_buffer = 0;
				bit_pos = 63;
			}
		} else {
			if (bit_pos > -1)
				bit_buffer |= bits >> (amount - (bit_pos + 1));
			write(bit_buffer);
			bit_buffer = 0;
			amount -= bit_pos + 1;
			bit_buffer |= bits << (64 - amount);
			bit_pos = 63 - amount;
		}
	}
	
	// Reading methods
	
	void CompressedIO::open_read(const std::string& input_file) {
		assert(state == State::Neutral);
		stream.open(input_file, std::ios::in | std::ios::binary);
		if (!stream.is_open())
			throw std::ios_base::failure("Could not open compressed file \"" + input_file + "\" for reading");
		mode = IOMode::File;
		state = State::Read;
		bit_mode = false;
		byte_pos = 0;
	}
	
	void CompressedIO::open_read(const char* input_data, size_t size) {
		assert(state == State::Neutral);
		if (input_data == nullptr)
			throw std::invalid_argument("Pointer to compressed data buffer cannot be null");
		mode = IOMode::Memory;
		state = State::Read;
		bit_mode = false;
		read_buffer = input_data;
		read_buffer_size = size;
		read_buffer_index = 0;
		byte_pos = 0;
	}
	
	bool CompressedIO::read_bit() {
		if (bit_pos >= 0) {
			bool res = bool(bit_buffer << (63 - bit_pos) >> 63);
			bit_pos--;
			return res;
		} else {
			bit_buffer = read<uint64_t>();
			bit_pos = 62;
			return bool(bit_buffer >> 63);
		}
	}
	
	uint64_t CompressedIO::read_bits_internal(uint8_t amount) {
		assert(amount <= 64);
		uint64_t result = 0;
		if (amount <= bit_pos + 1) {
			result = bit_buffer << (63 - bit_pos) >> (64 - amount);
			bit_pos -= amount;
		} else {
			if (bit_pos > -1)
				result = bit_buffer << (64 - bit_pos - 1) >> (64 - amount);
			bit_buffer = read<uint64_t>();
			amount -= bit_pos + 1;
			result |= bit_buffer >> (64 - amount);
			bit_pos = 63 - amount;
		}
		return result;
	}
	
	// Entropy coding
	
	void CompressedIO::write_bits_dynamic_length(uint64_t val) {
		
		// First, the value's length
		uint8_t length = 0;
		uint64_t temp = val;
		while (temp) {
			temp >>= 1;
			length++;
		}
		length = std::max(uint8_t(1), length); // A 0 still requires 1 bit to encode
		write_bits_internal(length, 6);
		
		// Next, the value itself
		write_bits_internal(val, length);
		
		encoding_bits += 6 + length;
		
	}
	
	uint64_t CompressedIO::read_bits_dynamic_length() {
		// First, the value's length
		uint8_t length = read_bits_internal(6);
		// Next, the value itself
		uint64_t val = read_bits_internal(length);
		return val;
	}
	
	std::map<uint64_t, uint64_t> CompressedIO::calculate_frequency_table(const std::vector<uint64_t>& data) {
		std::map<uint64_t, uint64_t> frequencies;
		for (const auto& val : data)
			frequencies[val]++;
		return frequencies;
	}
	
	std::map<uint64_t, uint64_t> CompressedIO::merge_frequency_tables(const std::vector<std::map<uint64_t, uint64_t>>& frequency_tables) {
		std::map<uint64_t, uint64_t> frequencies;
		for (const auto& frequency_table : frequency_tables)
			for (const auto& entry : frequency_table)
				frequencies[entry.first] += entry.second;
		return frequencies;
	}
	
	size_t CompressedIO::write_frequency_table(const std::map<uint64_t, uint64_t>& frequencies) {
		
		uint64_t old_encoding_bits = encoding_bits;
		encoding_bits = 0;
		
		// Write table size
		uint64_t unique_symbols_amount = frequencies.size();
		write_bits_dynamic_length(unique_symbols_amount);
		encoding_bits += sizeof(uint64_t)*8;
		
		// Write table itself
		for (auto& entry : frequencies) {
			write_bits_dynamic_length(entry.first); // Symbol
			write_bits_dynamic_length(entry.second); // Frequency
		}
		
		std::swap(old_encoding_bits, encoding_bits);
		return old_encoding_bits;
		
	}
	
	std::map<uint64_t, uint64_t> CompressedIO::read_frequency_table() {
		
		// Read table size
		uint64_t unique_symbols_amount = read_bits_dynamic_length();
		if (unique_symbols_amount == 0)
			throw std::runtime_error("Read frequency table with 0 symbols");
		
		// Table itself
		std::map<uint64_t, uint64_t> frequencies;
		for (uint64_t i = 0; i < unique_symbols_amount; i++) {
			uint64_t symbol = read_bits_dynamic_length();
			uint64_t frequency = read_bits_dynamic_length();
			frequencies[symbol] = frequency; // Symbol, frequency
		}
		
		return frequencies;
		
	}
	
	std::map<uint64_t, std::pair<uint64_t, uint64_t>> CompressedIO::calculate_encoding_frequency_map(const std::map<uint64_t, uint64_t>& frequencies, EntropyCodingMethod entropy_coding_method) {
		uint64_t total_frequency = 0;
		for (auto& entry : frequencies)
			total_frequency += entry.second;
		uint64_t new_full_range = (entropy_coding_method == EntropyCodingMethod::ANS ? ANS_L : total_frequency);
		uint64_t cum_frequency = 0; // Tracks the cumulative unscaled frequency
		uint64_t cum_scaled_frequency = 0; // Tracks the cumulative scaled frequency
		std::map<uint64_t, std::pair<uint64_t, uint64_t>> res;
		for (auto& entry : frequencies) {
			uint64_t next_cum_scaled_frequency = (cum_frequency + entry.second)*new_full_range/total_frequency;
			assert(next_cum_scaled_frequency > cum_scaled_frequency);
			res[entry.first] = {cum_scaled_frequency, next_cum_scaled_frequency - cum_scaled_frequency};
			cum_scaled_frequency = next_cum_scaled_frequency;
			cum_frequency += entry.second;
		}
		return res;
	}
	
	std::map<uint64_t, std::pair<uint64_t, uint64_t>> CompressedIO::calculate_decoding_frequency_map(const std::map<uint64_t, uint64_t>& frequencies, EntropyCodingMethod entropy_coding_method) {
		std::map<uint64_t, std::pair<uint64_t, uint64_t>> res;
		uint64_t full_range = 0;
		for (auto& entry : calculate_encoding_frequency_map(frequencies, entropy_coding_method)) {
			uint64_t symbol = entry.first;
			uint64_t range_start = entry.second.first;
			uint64_t frequency = entry.second.second;
			res[range_start] = {symbol, frequency};
			full_range = std::max(full_range, range_start + frequency); // Calculate extent of full range
		}
		res[full_range] = {0, 0}; // Add a sentinel at the back to make sure any query for an upper bound can be decremented safely
		return res;
	}
	
	size_t CompressedIO::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) {
		
		encoding_bits = 0;
		write_bits_dynamic_length(data.size());
		if (encoding_frequency_map.size() > 1) {
			// Only perform entropy coding if there is more than one kind of symbol
			// Else encoding is trivial anyway and probabilities of 1 can cause problems for the ANS implementation
			switch (entropy_coding_method) {
				case EntropyCodingMethod::Arithmetic:
					arithmetic_encode(encoding_frequency_map, data);
					break;
				case EntropyCodingMethod::ANS:
					ANS_encode(encoding_frequency_map, data);
					break;
			}
		}
		
		return encoding_bits;
		
	}
	
	std::vector<uint64_t> CompressedIO::entropy_decode(const std::map<uint64_t, std::pair<uint64_t, uint64_t>>& decoding_frequency_map, EntropyCodingMethod entropy_coding_method) {
		
		uint64_t symbols_amount = read_bits_dynamic_length();
		if (symbols_amount > 0) {
			if (decoding_frequency_map.size() >= 3) {
				switch (entropy_coding_method) {
					case EntropyCodingMethod::Arithmetic:
						return arithmetic_decode(decoding_frequency_map, symbols_amount);
					case EntropyCodingMethod::ANS:
						return ANS_decode(decoding_frequency_map, symbols_amount);
				}
			} else {
				// If there's only one kind of symbol, skip entropy decoding and just return that symbol symbols_amount times
				return std::vector<uint64_t>(symbols_amount, decoding_frequency_map.begin()->second.first);
			}
		}
		
		// In case there are zero symbols
		return std::vector<uint64_t>();
		
	}
	
	/*
	 * Arithmetic coding
	 * 
	 * The following arithmetic coding code is a slightly modified version of code from
	 * https://marknelson.us/posts/2014/10/19/data-compression-with-arithmetic-coding.html
	 * distributed under the following MIT License:
	 * 
	 * Copyright (c) 2014 Mark Thomas Nelson
	 * 
	 * Permission is hereby granted, free of charge, to any person obtaining a copy
	 * of this software and associated documentation files (the "Software"), to deal
	 * in the Software without restriction, including without limitation the rights
	 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
	 * copies of the Software, and to permit persons to whom the Software is
	 * furnished to do so, subject to the following conditions:
	 * 
	 * The above copyright notice and this permission notice shall be included in
	 * all copies or substantial portions of the Software.
	 * 
	 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
	 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
	 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
	 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
	 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
	 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
	 * THE SOFTWARE.
	 */
	
	void CompressedIO::write_bit_and_pending(bool bit) {
		write_bit(bit);
		for (size_t i = 0; i < pending_bits; i++)
			write_bit(!bit);
		encoding_bits += 1 + pending_bits;
		pending_bits = 0;
	}
	
	void CompressedIO::arithmetic_encode(const std::map<uint64_t, std::pair<uint64_t, uint64_t>>& encoding_frequency_map, const std::vector<uint64_t>& data) {
		
		uint64_t symbols_amount = data.size();
		const auto& last_value_pair = encoding_frequency_map.rbegin()->second;
		uint64_t total_frequency = last_value_pair.first + last_value_pair.second;
		pending_bits = 0;
		uint64_t low = 0;
		uint64_t high = AC_MAX_CODE;
		
		for (size_t i = 0; i < symbols_amount; i++) {
			
			uint64_t symbol = data[i];
			auto& entry = encoding_frequency_map.at(symbol);
			uint64_t plow = entry.first;
			uint64_t phigh = entry.first + entry.second;
			
			uint64_t range = high - low + 1;
			high = low + (range * phigh / total_frequency) - 1;
			low = low + (range * plow / total_frequency);

			while (true) {
				
				if (high < AC_ONE_HALF) {
					write_bit_and_pending(0);
				} else if (low >= AC_ONE_HALF) {
					write_bit_and_pending(1);
				} else if (low >= AC_ONE_FOURTH && high < AC_THREE_FOURTHS) {
					pending_bits++;
					low -= AC_ONE_FOURTH;
					high -= AC_ONE_FOURTH;
				} else {
					break;
				}
				
				high <<= 1;
				high++;
				low <<= 1;
				high &= AC_MAX_CODE;
				low &= AC_MAX_CODE;
				
			}
			
		}
		pending_bits++;
		if (low < AC_ONE_FOURTH)
			write_bit_and_pending(0);
		else
			write_bit_and_pending(1);
		
		write_bits_internal(0UL, AC_CODE_VALUE_BITS - 2); // Trailing zeros
		encoding_bits += AC_CODE_VALUE_BITS - 2;
		
	}
	
	std::vector<uint64_t> CompressedIO::arithmetic_decode(const std::map<uint64_t, std::pair<uint64_t, uint64_t>>& decoding_frequency_map, uint64_t symbols_amount) {
		
		uint64_t total_frequency = decoding_frequency_map.rbegin()->first;
		uint64_t high = AC_MAX_CODE;
		uint64_t low = 0;
		uint64_t state = read_bits_internal(AC_CODE_VALUE_BITS);
		std::vector<size_t> data(symbols_amount);
		for (size_t i = 0; i < symbols_amount; i++) {
			
			uint64_t range = high - low + 1;
			uint64_t scaled_value =  ((state - low + 1) * total_frequency - 1) / range;
			
			auto it = decoding_frequency_map.upper_bound(scaled_value);
			uint64_t phigh = it->first;
			it--; // it now points to the greatest bound <= scaled_value
			data[i] = (it->second).first;
			uint64_t plow = it->first;
			
			high = low + (range*phigh)/total_frequency - 1;
			low = low + (range*plow)/total_frequency;
			
			while (true) {
				
				if (high < AC_ONE_HALF) {
					// Do nothing, bit is a zero
				} else if (low >= AC_ONE_HALF) {
					state -= AC_ONE_HALF;  // Subtract one half from all three code values
					low -= AC_ONE_HALF;
					high -= AC_ONE_HALF;
				} else if (low >= AC_ONE_FOURTH && high < AC_THREE_FOURTHS) {
					state -= AC_ONE_FOURTH;
					low -= AC_ONE_FOURTH;
					high -= AC_ONE_FOURTH;
				} else {
					break;
				}
				
				low <<= 1;
				high <<= 1;
				high++;
				state <<= 1;
				bool bit = read_bit();
				state |= uint64_t(bit);
				
			}
			
		}
		
		return data;
		
	}
	
	/*
	 * Asymmetric numeral systems
	 * This implementation encodes with almost the same efficiency as the entropy limit/arithmetic coding (overhead: ~0.25%)
	 * But it encodes faster and decodes much faster
	 * Based on:
	 * - https://arxiv.org/abs/1311.2540
	 * - https://en.wikipedia.org/wiki/Asymmetric_numeral_systems#Range_variants_(rANS)_and_streaming
	 * - https://github.com/rygorous/ryg_rans
	 * This code could be optimized further for speed, for example by pre-computing the integer reciprocals of the frequencies,
	 * which can replace expensive integer divisions by integer multiplications.
	 * However, this would have little effect on the total runtime of ATC compression/decompression.
	 */
	
	void CompressedIO::ANS_encode(const std::map<uint64_t, std::pair<uint64_t, uint64_t>>& encoding_frequency_map, const std::vector<uint64_t>& data) {
		
		// Encode symbols
		std::stack<bool> bit_stack; // Bits are sent here to be reversed later during writing
		uint64_t state = ANS_L; // Initialize state
		for (size_t i = data.size(); i-- > 0;) {
			uint64_t symbol = data[i];
			auto& entry = encoding_frequency_map.at(symbol);
			uint64_t range_start = entry.first;
			uint64_t frequency = entry.second;
			// Normalize state such that C(s, x) will be in the valid state interval again
			uint64_t max_state = ANS_B*frequency;
			while (state >= max_state) {
				for (uint8_t j = 0; j < ANS_LOG_B; j++) {
					bit_stack.push(state%2);
					state >>= 1;
				}
			}
			// Apply x_new = C(s, x)
			state = ((state/frequency) << ANS_LOG_L) + (state%frequency) + range_start;
		}
		write_bits_internal(state, ANS_LOG_L + ANS_LOG_B);
		encoding_bits += ANS_LOG_L + ANS_LOG_B;
		
		// Push bit stack to output stream
		encoding_bits += bit_stack.size();
		while (!bit_stack.empty()) {
			write_bit(bit_stack.top());
			bit_stack.pop();
		}
		
	}
	
	std::vector<uint64_t> CompressedIO::ANS_decode(const std::map<uint64_t, std::pair<uint64_t, uint64_t>>& decoding_frequency_map, uint64_t symbols_amount) {
		
		// Read symbols
		const uint64_t mask = ANS_L - 1;
		uint64_t state = read_bits_internal(ANS_LOG_L + ANS_LOG_B);
		std::vector<size_t> data(symbols_amount);
		for (size_t i = 0; i < symbols_amount; i++) {
			// Extract symbol and apply (s, x_new) = D(x)
			auto it = (--decoding_frequency_map.upper_bound(state & mask)); // it now points to the greatest bound <= (state & mask)
			uint64_t range_start = it->first;
			uint64_t symbol = (it->second).first;
			uint64_t frequency = (it->second).second;
			data[i] = symbol;
			state = frequency*(state >> ANS_LOG_L) + (state & mask) - range_start;
			// Normalize state into valid state interval
			while (state < ANS_L) {
				for (uint8_t j = 0; j < ANS_LOG_B; j++) {
					state <<= 1;
					state |= read_bits_internal(1);
				}
			}
		}
		
		return data;
		
	}
	
}
