/*
 * This file defines the class atc::RawIO and then declares and defines its members,
 * because the class is templated and not explicitly instantiated,
 * so all methods must be implemented in this header file.
 */

#pragma once

#include <cstdint>
#include <vector>
#include <stdexcept>
#include <iostream>
#include <sstream>
#include <fstream>
#include <numeric>
#include <functional>
#include <algorithm>
#include <memory>
#include <cmath>
#include <experimental/filesystem>

#include "RawIOBase.hpp"
#include "util.hpp"
#include "DataType.hpp"
#include "IOMode.hpp"
#include "LittleEndianIO.hpp"

namespace atc {
	
	/*
	 * This class manages the reading/writing of raw data (original/decompressed data),
	 * mapping in between the internally-used floating-point type and the external data type.
	 * Each atc::RawIO<ftype, dtype> inherits from atc::RawIOBase<ftype> such that data-type-independent calls can be made
	 * from atc::Compressor<ftype> using a (somewhat opaque) pointer.
	 */
	template<typename ftype, typename dtype>
	class RawIO : public RawIOBase<ftype> {
	public:
		size_t get_stored_data_size();
		uint8_t get_data_type_id();
		size_t get_data_type_size();
		void set_read_target(const std::string& data_file, size_t skip_bytes = 0);
		/*
		 * Sets a file target for reading. Ensures that we do not read past the given size.
		 * Not declared in atc::RawIOBase because of dtype dependency, but this works fine because it is only called from dtype-templated code.
		 */
		void set_read_target(const dtype* buffer, size_t size);
		std::vector<ftype> read(const std::vector<size_t>& sizes, bool store_data);
		void set_write_target(const std::string& file);
		/*
		 * Sets a file target for writing. Ensures that we do not read past the given size.
		 * Not declared in atc::RawIOBase because of dtype dependency, but this works fine because it is only called from dtype-templated code.
		 */
		void set_write_target(dtype* buffer, size_t size);
		double write(const std::vector<ftype>& data, bool calculate_sse, size_t size = 0);
	private:
		const size_t TEMP_BUFFER_SIZE = 1 << 20; // Upper bound to size of temporary buffer used for casting data. In bytes, not elements.
		const size_t TEMP_BUFFER_ELEMENTS = TEMP_BUFFER_SIZE/sizeof(dtype); // Temporary buffer size in elements.
		IOMode mode; // Indicates whether the data target is a file or buffer in-memory.
		std::string file; // The file target, used both for reading and writing in File mode.
		size_t skip_bytes; // The amount of bytes to be skipped at the start of the file, when reading in File mode.
		const dtype* read_buffer; // If the data is stored, this pointer will always point to the stored buffer. Also used for storing a pointer to the input buffer in the case of in-memory input.
		size_t memory_target_array_size; // Given size of reading/writing target buffer to prevent reading/writing out-of-bounds.
		size_t stored_buffer_size = 0; // Size of currently stored buffer, stored by calling read with store_data = true.
		bool store_data; // Stores whether or not data was stored in the last read call.
		std::vector<dtype> stored_data; // If the data is stored and the data was read from a file, the copy is stored here to provide for easy memory management.
		dtype* write_buffer; // If writing in Memory mode, this pointer stores the writing location given by the user.
	};
	
	template<typename ftype, typename dtype>
	size_t RawIO<ftype, dtype>::get_stored_data_size() {
		return stored_buffer_size;
	}
	
	template<typename ftype, typename dtype>
	uint8_t RawIO<ftype, dtype>::get_data_type_id() {
		return data_type_to_id(get_data_type<dtype>());
	}
	
	template<typename ftype, typename dtype>
	size_t RawIO<ftype, dtype>::get_data_type_size() {
		return sizeof(dtype);
	}
	
	template<typename ftype, typename dtype>
	void RawIO<ftype, dtype>::set_read_target(const std::string& file, size_t skip_bytes) {
		mode = IOMode::File;
		this->file = file;
		this->skip_bytes = skip_bytes;
	}
	
	template<typename ftype, typename dtype>
	void RawIO<ftype, dtype>::set_read_target(const dtype* buffer, size_t size) {
		if (buffer == nullptr)
			throw std::invalid_argument("Pointer to original data buffer cannot be null");
		mode = IOMode::Memory;
		this->read_buffer = buffer;
		memory_target_array_size = size;
	}
	
	template<typename ftype, typename dtype>
	std::vector<ftype> RawIO<ftype, dtype>::read(const std::vector<size_t>& sizes, bool store_data) {
		
		// Store arguments for later use
		size_t size = product(sizes);
		if (mode == IOMode::Memory && size > memory_target_array_size)
			throw std::invalid_argument("Given original data buffer (size " + std::to_string(memory_target_array_size) + ") is not large enough to extract tensor of given sizes (size " + std::to_string(size) + ")");
		stored_buffer_size = store_data ? size : 0;
		this->store_data = store_data;
		
		std::vector<ftype> res(size);
		switch (mode) {
		case IOMode::File: {
			
			// Check that file size matches expected size
			size_t expected_size = skip_bytes + size * sizeof(dtype);
			size_t file_size = std::experimental::filesystem::file_size(file);
			if (file_size != expected_size) {
				
				// Throw exception with detailed error message
				std::stringstream message;
				message << "Invalid file size: expected ";
				if (skip_bytes > 0)
					message << skip_bytes << " + ";
				message << "(" << sizes[0];
				for (uint8_t i = 1; i < sizes.size(); i++)
					message << "*" << sizes[i];
				message << ")*" << sizeof(dtype) << " = " << expected_size << " bytes, but found " << file_size << " bytes";
				if (file_size < expected_size)
					message << " (" << expected_size/double(file_size) << " times too small)";
				else
					message << " (" << file_size/double(expected_size) << " times too large)";
				if (skip_bytes == 0 && expected_size < file_size && file_size < expected_size + 1000)
					message << ". Perhaps the file has a header (use skip_bytes)?";
				throw std::invalid_argument(message.str());
				
			}
			
			// Check if file opens properly
			std::ifstream input_stream(file, std::ios::in | std::ios::binary);
			if (!input_stream.is_open())
				throw std::ios_base::failure("Could not open data file \"" + file + "\"");
			
			// Read data file
			input_stream.seekg(skip_bytes, std::ios::beg);
			if (store_data) {
				// Read full data file into original data buffer without casting, then copy to return buffer with casting
				stored_data.resize(size);
				little_endian_file_read(input_stream, stored_data.data(), size);
				std::copy(stored_data.begin(), stored_data.end(), res.begin()); // Copy also performs casting
				read_buffer = stored_data.data();
			} else {
				force_clear(stored_data); // If not storing data, stored_data should be emptied by this read call.
				if (std::is_same<ftype, dtype>::value) {
					// If original data does not need to be stored and the datatype and floating-point type are the same, no intermediary buffer is necessary
					little_endian_file_read(input_stream, res.data(), size);
				} else {
					// Read data from file in blocks, casting and storing each block in the output buffer
					std::vector<dtype> temp(TEMP_BUFFER_ELEMENTS);
					for(size_t i = 0; i < size; i += TEMP_BUFFER_ELEMENTS) {
						// Read, cast and store one block
						size_t current_block_size = std::min(TEMP_BUFFER_ELEMENTS, size - i);
						little_endian_file_read(input_stream, temp.data(), current_block_size);
						std::copy(temp.begin(), temp.begin() + current_block_size, res.begin() + i); // Copy also performs casting
					}
				}
			}
			input_stream.close();
			
			break;
			
		}
		case IOMode::Memory:
			force_clear(stored_data); // stored_data not needed, even if store_data is true we can directly store the user-given data pointer.
			res.assign(read_buffer, read_buffer + size); // assign also performs casting from dtype to ftype
			break;
		}
		
		return res;
		
	}
	
	template<typename ftype, typename dtype>
	void RawIO<ftype, dtype>::set_write_target(const std::string& file) {
		mode = IOMode::File;
		this->file = file;
	}
	
	template<typename ftype, typename dtype>
	void RawIO<ftype, dtype>::set_write_target(dtype* buffer, size_t size) {
		if (buffer == nullptr)
			throw std::invalid_argument("Pointer to decompressed data buffer cannot be null");
		mode = IOMode::Memory;
		this->write_buffer = buffer;
		memory_target_array_size = size;
	}
	
	// Calculates || v1[first1:last1] - v2[first2:first2 + (last1 - first1)] ||_F^2 using Kahan summation
	// Based on https://en.wikipedia.org/wiki/Kahan_summation_algorithm
	template<typename T1, typename T2>
	double sse_of_range(const T1* v1, const T2* v2, size_t first1, size_t last1, size_t first2) {
		// Calculates the SSE using Kahan summation
		double sse = 0;
		double c = 0;
		for(size_t i = 0; i < last1 - first1; i++) {
			double term = pow(double(v1[first1 + i]) - double(v2[first2 + i]), 2);
			double y = term - c;
			double t = sse + y;
			c = (t - sse) - y;
			sse = t;
		}
		return sse;
	}
	
	template<typename ftype, typename dtype>
	double RawIO<ftype, dtype>::write(const std::vector<ftype>& data, bool calculate_sse, size_t size) {
		
		// Checks
		assert(!calculate_sse || store_data);
		assert(size <= data.size());
		if (size == 0)
			size = data.size();
		if (mode == IOMode::Memory && size > memory_target_array_size)
			throw std::invalid_argument("Given decompressed data buffer (size " + std::to_string(memory_target_array_size) + ") is not large enough to fit decompressed data (size " + std::to_string(size) + ")");
		
		// Preliminaries
		std::ofstream output_stream;
		if (mode == IOMode::File) {
			output_stream = std::ofstream(file, std::ios::out | std::ios::binary);
			if (!output_stream.is_open()) {
				std::stringstream message;
				message << "Could not open output file \"" << file << "\"";
				throw std::ios_base::failure(message.str());
			}
		}
		
		// Write to data file
		double sse = 0; // Return 0 by default
		if (std::is_same<ftype, dtype>::value) {
			
			// No casting, so no intermediary buffer necessary
			if (calculate_sse)
				sse = sse_of_range(data.data(), read_buffer, 0, size, 0);
			switch (mode) {
			case IOMode::File:
				little_endian_file_write(output_stream, data.data(), data.size());
				break;
			case IOMode::Memory:
				std::copy(data.begin(), data.end(), write_buffer);
				break;
			}
			
		} else {
			
			// Casting required, so we'll work block by block using an intermediary buffer
			size_t actual_block_size = size;
			dtype* block_ptr = write_buffer;
			std::vector<dtype> temp;
			if (mode == IOMode::File) {
				actual_block_size = TEMP_BUFFER_ELEMENTS;
				temp.resize(actual_block_size);
				block_ptr = temp.data();
			}
			for(size_t i = 0; i < size; i += actual_block_size) {
				
				// Round, cast and copy data to temporary buffer
				size_t current_block_size = std::min(actual_block_size, size - i);
				if (std::is_floating_point<dtype>::value) {
					// Cast to another floating-point type, no need for rounding/clamping
					std::copy(data.begin() + i, data.begin() + i + current_block_size, block_ptr);
				} else {
					// Cast to an integer type, so round
					auto transform = [] (ftype x) {
						return dtype(std::min(std::max(std::round(x), ftype(std::numeric_limits<dtype>::min())), ftype(std::numeric_limits<dtype>::max())));
					};
					std::transform(data.begin() + i, data.begin() + i + current_block_size, block_ptr, transform);
				}
				
				if (mode == IOMode::File) {
					// Write temporary buffer to file
					little_endian_file_write(output_stream, block_ptr, current_block_size);
				}
				
				// Calculate SSE if necessary
				if (calculate_sse)
					sse += sse_of_range(block_ptr, read_buffer, 0, current_block_size, i);
				
			}
			
		}
		
		if (mode == IOMode::File)
			output_stream.close();
	
		return sse;
		
	}
	
}
