/*
 * This file implements atc.hpp.
 * A large section of this file simply consists of mapping interface functions to a smaller set of fundamental functions.
 */

#include <vector>
#include <utility>
#include <memory>

#include "atc.hpp"
#include "Results.hpp"
#include "Options.hpp"
#include "Compressor.hpp"
#include "RawIO.hpp"
#include "CompressedIO.hpp"
#include "IOMode.hpp"
#include "DataType.hpp"
#include "util.hpp"

namespace atc {
	
	std::vector<size_t> dummy_sizes; // Default value, only used to make the OriginalMetadata default constructor valid
	
	/*
	 * The following 3 structs are used as a generic way to store metadata about the original/compressed/decompressed data.
	 * Mainly they provide a unified way of passing metadata both when that target data is in a file or in-memory,
	 * which simplifies the unifying of the interface functions into one general compressor_operation function.
	 */
	
	template<typename dtype = uint8_t>
	struct OriginalMetadata {
		OriginalMetadata() : sizes(dummy_sizes) { }; // Default constructor, should only be used when there is no original data (i.e. for decompress calls).
		OriginalMetadata(const std::string& file, DataType data_type, const std::vector<size_t>& sizes, double target_error)
			: mode(IOMode::File), file(file), data_type(data_type), sizes(sizes), target_error(target_error) { };
		OriginalMetadata(const dtype* buffer, const std::vector<size_t>& sizes, const double target_error)
			: mode(IOMode::Memory), buffer(buffer), sizes(sizes), target_error(target_error) { };
		IOMode mode;
		std::string file;
		DataType data_type;
		const dtype* buffer;
		const std::vector<size_t>& sizes;
		double target_error;
	};
	
	struct CompressedMetadata {
		CompressedMetadata(const std::string& file) : mode(IOMode::File), file(file) { };
		CompressedMetadata() : mode(IOMode::Memory) { }; // Used when the compressed data is an output, i.e. stored in an std::vector.
		CompressedMetadata(const char* buffer, size_t size) : mode(IOMode::Memory), buffer(buffer), size(size) { }; // Used when compressed data is an input, i.e. given 
		IOMode mode;
		std::string file;
		const char* buffer;
		size_t size;
	};
	
	template<typename dtype = uint8_t>
	struct DecompressedMetaData {
		DecompressedMetaData() = default; // Default constructor, should only be used when there is no decompressed data (i.e. for compress calls).
		DecompressedMetaData(const std::string& file) : mode(IOMode::File), file(file) { };
		DecompressedMetaData(dtype* buffer, size_t size) : mode(IOMode::Memory), buffer(buffer), size(size) { };
		IOMode mode;
		std::string file;
		dtype* buffer;
		size_t size;
	};
	
	/*
	 * This general function handles all ATC compress, decompress and compress_decompress functions.
	 * The original/compressed/decompressed data targets are wrapped into generic structs to facilitate this.
	 */
	template<typename dtype, typename ftype>
	Results compressor_operation(bool compress, bool decompress, const OriginalMetadata<dtype>& original, const CompressedMetadata& compressed, const DecompressedMetaData<dtype>& decompressed, const Options& options) {
		
		if (options.calculate_error && (!compress || !decompress))
			throw std::invalid_argument("Cannot calculate error using only compress or decompress, try compress_decompress instead");
		
		Compressor<ftype> compressor(options);
		
		// Set-up RawIO object for both compression/decompression
		// The following two variables are defined at function-level scope to make sure that, whichever container is used, the corresponding RawIO object is only destroyed at the end of this function
		std::unique_ptr<RawIOBase<ftype>> abstract_raw_io;
		RawIO<ftype, dtype> concrete_raw_io;
		RawIOBase<ftype>* raw_io_ptr; // Here we actually store the pointer regardless of which container is used
		if (!(compress && original.mode == IOMode::Memory) && !(decompress && decompressed.mode == IOMode::Memory)) {
			// Actual dtype is not known at compile time, current dtype is just a dummy type which can be ignored in this case
			// If we're compressing, data_type is given
			// If we're decompressing, we can extract it from the compressed data
			abstract_raw_io = RawIOBase<ftype>::create(compress ? original.data_type : (compressed.mode == IOMode::File ? read_data_type(compressed.file) : read_data_type(compressed.size, compressed.buffer)));
			raw_io_ptr = abstract_raw_io.get();
		} else {
			// We compress from or decompress to memory, so dtype is known
			raw_io_ptr = &concrete_raw_io;
		}
		
		CompressedIO compressed_io;
		
		std::vector<char> compressed_data;
		if (compress) {
			
			// Set up RawIO
			switch (original.mode) {
				case IOMode::File:
					raw_io_ptr->set_read_target(original.file, options.skip_bytes);
					break;
				case IOMode::Memory:
					concrete_raw_io.set_read_target(original.buffer, product(original.sizes));
					break;
			}
			
			// Set-up CompressedIO
			switch (compressed.mode) {
				case IOMode::File:
					compressed_io.open_write(compressed.file);
					break;
				case IOMode::Memory:
					compressed_io.open_write();
					break;
			}
			
			compressor.compress(raw_io_ptr, original.sizes, original.target_error, compressed_io);
			
			if (compressed.mode == IOMode::Memory)
				compressed_data = compressed_io.get_and_clear_buffer();
			
		}
		
		if (decompress) {
			
			// Set-up CompressedIO
			switch (compressed.mode) {
				case IOMode::File:
					compressed_io.open_read(compressed.file);
					break;
				case IOMode::Memory:
					compressed_io.open_read(compress ? compressed_data.data() : compressed.buffer, compress ? compressed_data.size() : compressed.size);
					break;
			}
			
			// Set up RawIO
			switch (decompressed.mode) {
				case IOMode::File:
					raw_io_ptr->set_write_target(decompressed.file);
					break;
				case IOMode::Memory:
					concrete_raw_io.set_write_target(decompressed.buffer, decompressed.size);
					break;
			}
			
			compressor.decompress(compressed_io, raw_io_ptr);
			
		}
		
		Results results = compressor.get_and_clear_results();
		results.compressed_data.swap(compressed_data);
		return results;
		
	}
	
	// Handles dynamic dispatch in terms of float type
	template<typename dtype = uint8_t>
	Results compressor_operation(bool compress, bool decompress, const OriginalMetadata<dtype>& original, const CompressedMetadata& compressed, const DecompressedMetaData<dtype>& decompressed, const Options& options) {
		switch(options.float_type) {
			case FloatType::float32:
				return compressor_operation<dtype, float>(compress, decompress, original, compressed, decompressed, options);
			case FloatType::float64:
				return compressor_operation<dtype, double>(compress, decompress, original, compressed, decompressed, options);
		}
		return Results(); // To make compiler happy, should never be reached
	}
	
	// Wrappers for default arguments
	template<typename dtype = uint8_t>
	Results compressor_operation(const OriginalMetadata<dtype>& original, const CompressedMetadata& compressed, const Options& options) {
		return compressor_operation<dtype>(true, false, original, compressed, DecompressedMetaData<dtype>(), options);
	}
	template<typename dtype = uint8_t>
	Results compressor_operation(const OriginalMetadata<dtype>& original, const Options& options) {
		return compressor_operation<dtype>(original, CompressedMetadata(), options);
	}
	template<typename dtype = uint8_t>
	Results compressor_operation(const CompressedMetadata& compressed, const DecompressedMetaData<dtype>& decompressed, const Options& options) {
		return compressor_operation<dtype>(false, true, OriginalMetadata<dtype>(), compressed, decompressed, options);
	}
	template<typename dtype = uint8_t>
	Results compressor_operation(const OriginalMetadata<dtype>& original, const CompressedMetadata& compressed, const DecompressedMetaData<dtype>& decompressed, const Options& options) {
		return compressor_operation<dtype>(true, true, original, compressed, decompressed, options);
	}
	template<typename dtype = uint8_t>
	Results compressor_operation(const OriginalMetadata<dtype>& original,  const DecompressedMetaData<dtype>& decompressed, const Options& options) {
		return compressor_operation<dtype>(original, CompressedMetadata(), decompressed, options);
	}
	
	// compress implementations
	
	Results compress(const std::string& original_file, DataType data_type, const std::vector<size_t>& sizes, double target_error, const std::string& compressed_file, const Options& options) {
		return compressor_operation(OriginalMetadata<>(original_file, data_type, sizes, target_error), CompressedMetadata(compressed_file), options);
	}
	
	Results compress(const std::string& original_file, DataType data_type, const std::vector<size_t>& sizes, double target_error, const Options& options) {
		return compressor_operation(OriginalMetadata<>(original_file, data_type, sizes, target_error), options);
	}
	
	template<typename dtype>
	Results compress(const dtype* original_data, const std::vector<size_t>& sizes, double target_error, const std::string& compressed_file, const Options& options) {
		return compressor_operation(OriginalMetadata<dtype>(original_data, sizes, target_error), CompressedMetadata(compressed_file), options);
	}
	
	template<typename dtype>
	Results compress(const dtype* original_data, const std::vector<size_t>& sizes, double target_error, const Options& options) {
		return compressor_operation(OriginalMetadata<dtype>(original_data, sizes, target_error), options);
	}
	
	// decompress implementations
	
	Results decompress(const std::string& compressed_file, const std::string& decompressed_file, const Options& options) {
		return compressor_operation(CompressedMetadata(compressed_file), DecompressedMetaData<>(decompressed_file), options);
	}
	
	template<typename dtype>
	Results decompress(const std::string& compressed_file, size_t decompressed_size, dtype* decompressed_data, const Options& options) {
		return compressor_operation(CompressedMetadata(compressed_file), DecompressedMetaData<dtype>(decompressed_data, decompressed_size), options);
	}
	
	Results decompress(size_t compressed_size, const char* compressed_data, const std::string& decompressed_file, const Options& options) {
		return compressor_operation(CompressedMetadata(compressed_data, compressed_size), DecompressedMetaData<>(decompressed_file), options);
	}
	
	template<typename dtype>
	Results decompress(size_t compressed_size, const char* compressed_data, size_t decompressed_size, dtype* decompressed_data, const Options& options) {
		return compressor_operation(CompressedMetadata(compressed_data, compressed_size), DecompressedMetaData<dtype>(decompressed_data, decompressed_size), options);
	}
	
	// decompress implementations
	
	Results compress_decompress(const std::string& original_file, DataType data_type, const std::vector<size_t>& sizes, double target_error, const std::string& compressed_file, const std::string& decompressed_file, const Options& options) {
		return compressor_operation(OriginalMetadata<>(original_file, data_type, sizes, target_error), CompressedMetadata(compressed_file), DecompressedMetaData<>(decompressed_file), options);
	}

	template<typename dtype>
	Results compress_decompress(const std::string& original_file, const std::vector<size_t>& sizes, double target_error, const std::string& compressed_file, dtype* decompressed_data, const Options& options) {
		return compressor_operation(OriginalMetadata<dtype>(original_file, get_data_type<dtype>(), sizes, target_error), CompressedMetadata(compressed_file), DecompressedMetaData<dtype>(decompressed_data, product(sizes)), options);
	}

	Results compress_decompress(const std::string& original_file, DataType data_type, const std::vector<size_t>& sizes, double target_error, const std::string& decompressed_file, const Options& options) {
		return compressor_operation(OriginalMetadata<>(original_file, data_type, sizes, target_error), DecompressedMetaData<>(decompressed_file), options);
	}

	template<typename dtype>
	Results compress_decompress(const std::string& original_file, const std::vector<size_t>& sizes, double target_error, dtype* decompressed_data, const Options& options) {
		return compressor_operation(OriginalMetadata<dtype>(original_file, get_data_type<dtype>(), sizes, target_error), DecompressedMetaData<dtype>(decompressed_data, product(sizes)), options);
	}

	template<typename dtype>
	Results compress_decompress(const dtype* original_data, const std::vector<size_t>& sizes, double target_error, const std::string& compressed_file, const std::string& decompressed_file, const Options& options) {
		return compressor_operation(OriginalMetadata<dtype>(original_data, sizes, target_error), CompressedMetadata(compressed_file), DecompressedMetaData<dtype>(decompressed_file), options);
	}

	template<typename dtype>
	Results compress_decompress(const dtype* original_data, const std::vector<size_t>& sizes, double target_error, const std::string& compressed_file, dtype* decompressed_data, const Options& options) {
		return compressor_operation(OriginalMetadata<dtype>(original_data, sizes, target_error), CompressedMetadata(compressed_file), DecompressedMetaData<dtype>(decompressed_data, product(sizes)), options);
	}

	template<typename dtype>
	Results compress_decompress(const dtype* original_data, const std::vector<size_t>& sizes, double target_error, const std::string& decompressed_file, const Options& options) {
		return compressor_operation(OriginalMetadata<dtype>(original_data, sizes, target_error), DecompressedMetaData<dtype>(decompressed_file), options);
	}

	template<typename dtype>
	Results compress_decompress(const dtype* original_data, const std::vector<size_t>& sizes, double target_error, dtype* decompressed_data, const Options& options) {
		return compressor_operation(OriginalMetadata<dtype>(original_data, sizes, target_error), DecompressedMetaData<dtype>(decompressed_data, product(sizes)), options);
	}

	
	// Wrappers
	// Most of these just decompose an std::vector input into a pointer and a size
	
	// compress wrappers: convert original data from std::vector<dtype> to (dtype*, size_t)
	// Because the size of the data buffer is assumed to be product(sizes) afterwards, we initially verify that the passed buffer really is of that size
	
	// Verifies that, when passing the original data as an std::vector in-memory, the size of the vector actually matches the product of the given tensor sizes.
	template<typename dtype>
	void verify_original_data_size(const std::vector<dtype>& original_data, const std::vector<size_t>& sizes) {
		size_t size = product(sizes);
		if (original_data.size() != size)
			throw std::invalid_argument("Original data buffer size (" + std::to_string(original_data.size()) + " does not match product of tensor sizes (" + std::to_string(size) + ")");
	}
	
	template<typename dtype>
	Results compress(const std::vector<dtype>& original_data, const std::vector<size_t>& sizes, double target_error, const std::string& compressed_file, const Options& options) {
		verify_original_data_size(original_data, sizes);
		return compress(original_data.data(), sizes, target_error, compressed_file, options);
	}
	
	template<typename dtype>
	Results compress(const std::vector<dtype>& original_data, const std::vector<size_t>& sizes, double target_error, const Options& options) {
		verify_original_data_size(original_data, sizes);
		return compress(original_data.data(), sizes, target_error, options);
	}
	
	// decompress wrappers
	
	// Part 1: convert compressed data from std::vector<char> to (char*, size_t)
	
	Results decompress(const std::vector<char>& compressed_data, const std::string& decompressed_file, const Options& options) {
		return decompress(compressed_data.size(), compressed_data.data(), decompressed_file, options);
	}
	
	template<typename dtype>
	Results decompress(const std::vector<char>& compressed_data, size_t decompressed_size, dtype* decompressed_data, const Options& options) {
		return decompress(compressed_data.size(), compressed_data.data(), decompressed_size, decompressed_data, options);
	}
	
	template<typename dtype>
	TypedResults<dtype> decompress(const std::vector<char>& compressed_data, const Options& options) {
		return decompress<dtype>(compressed_data.size(), compressed_data.data(), options);
	}
	
	// Part 2: convert decompressed_data from (dtype*, size_t) to std::vector<dtype>
	
	// The following three functions let us read/calculate the decompressed tensor size without necessarily printing warnings when using autocrop.
	// This way these internal functions can be used both internally for vectorizing decompressed output and externally to implement read_decompressed_size and calculate_decompressed_size.
	
	size_t read_decompressed_size_internal(const std::string& compressed_file, const Options& options, bool suppress_autocrop_warnings) {
		CompressedIO io;
		io.open_read(compressed_file);
		return CompressorBase().read_decompressed_size(io, options, suppress_autocrop_warnings);
	}
	
	size_t read_decompressed_size_internal(size_t compressed_size, const char* compressed_data, const Options& options, bool suppress_autocrop_warnings) {
		CompressedIO io;
		io.open_read(compressed_data, compressed_size);
		return CompressorBase().read_decompressed_size(io, options, suppress_autocrop_warnings);
	}
	
	size_t calculate_decompressed_size_internal(const std::vector<size_t>& sizes, const Options& options, bool suppress_autocrop_warnings) {
		std::vector<size_t> sizes_copy(sizes);
		if (options.storage_order == StorageOrder::MSI)
			std::reverse(sizes_copy.begin(), sizes_copy.end());
		return CompressorBase().calculate_decompressed_size(sizes_copy, options, suppress_autocrop_warnings);
	}
	
	// This function simplifies a lot of the wrappers that follow
	// decompression_function is a function taking the arguments (decompressed_data_ptr, decompressed_data_size) and returning decompression results
	template<typename dtype, typename F>
	TypedResults<dtype> vectorize_decompress(const F& decompression_function, size_t size) {
		std::vector<dtype> decompressed_data(size);
		atc::TypedResults<dtype> results(decompression_function(size, decompressed_data.data()));
		results.decompressed_data.swap(decompressed_data);
		// Actual size might be smaller due to autocrop, so resize.
		// This might lead to a reallocation + copy, but only if autocrop is used, otherwise decompressed_data.size() is unchanged after all.
		results.decompressed_data.resize(product(results.decompressed_sizes)); 
		return results;
	}
	
	template<typename dtype>
	TypedResults<dtype> decompress(const std::string& compressed_file, const Options& options) {
		return vectorize_decompress<dtype>([&] (size_t size, dtype* ptr) {return decompress(compressed_file, size, ptr, options);}, read_decompressed_size_internal(compressed_file, options, true));
	}
	
	template<typename dtype>
	TypedResults<dtype> decompress(size_t compressed_size, const char* compressed_data, const Options& options) {
		return vectorize_decompress<dtype>([&] (size_t size, dtype* ptr) {return decompress(compressed_size, compressed_data, size, ptr, options);}, read_decompressed_size_internal(compressed_size, compressed_data, options, true));
	}
	
	// compress_decompress wrappers
	
	// Part 1: convert original data from std::vector<dtype> to (dtype*, size_t)
	
	template<typename dtype>
	Results compress_decompress(const std::vector<dtype>& original_data, const std::vector<size_t>& sizes, double target_error, const std::string& compressed_file, const std::string& decompressed_file, const Options& options) {
		verify_original_data_size(original_data, sizes);
		return compress_decompress(original_data.data(), sizes, target_error, compressed_file, decompressed_file, options);
	}
	
	template<typename dtype>
	Results compress_decompress(const std::vector<dtype>& original_data, const std::vector<size_t>& sizes, double target_error, const std::string& compressed_file, dtype* decompressed_data, const Options& options) {
		verify_original_data_size(original_data, sizes);
		return compress_decompress(original_data.data(), sizes, target_error, compressed_file, decompressed_data, options);
	}
	
	template<typename dtype>
	TypedResults<dtype> compress_decompress_to_memory(const std::vector<dtype>& original_data, const std::vector<size_t>& sizes, double target_error, const std::string& compressed_file, const Options& options) {
		verify_original_data_size(original_data, sizes);
		return compress_decompress(original_data.data(), sizes, target_error, compressed_file, options);
	}
	
	template<typename dtype>
	Results compress_decompress(const std::vector<dtype>& original_data, const std::vector<size_t>& sizes, double target_error, const std::string& decompressed_file, const Options& options) {
		verify_original_data_size(original_data, sizes);
		return compress_decompress(original_data.data(), sizes, target_error, decompressed_file, options);
	}
	
	template<typename dtype>
	Results compress_decompress(const std::vector<dtype>& original_data, const std::vector<size_t>& sizes, double target_error, dtype* decompressed_data, const Options& options) {
		verify_original_data_size(original_data, sizes);
		return compress_decompress(original_data.data(), sizes, target_error, decompressed_data, options);
	}
	
	template<typename dtype>
	TypedResults<dtype> compress_decompress(const std::vector<dtype>& original_data, const std::vector<size_t>& sizes, double target_error, const Options& options) {
		verify_original_data_size(original_data, sizes);
		return compress_decompress(original_data.data(), sizes, target_error, options);
	}
	
	// Part 2: convert decompressed_data from (size_t, dtype*) to std::vector<dtype>
	
	// This function simplifies a lot of the wrappers that follow
	// compress_decompress_function is a function taking the argument decompressed_data_ptr and returning decompression results
	template<typename dtype, typename F>
	TypedResults<dtype> vectorize_compress_decompress(const F& compress_decompress_function, const std::vector<size_t>& sizes, const Options& options) {
		return vectorize_decompress<dtype>([&] (size_t, dtype* ptr) {return compress_decompress_function(ptr);}, calculate_decompressed_size_internal(sizes, options, true));
	}
	
	template<typename dtype>
	TypedResults<dtype> compress_decompress(const std::string& original_file, const std::vector<size_t>& sizes, double target_error, const std::string& compressed_file, const Options& options) {
		return vectorize_compress_decompress<dtype>([&] (dtype* ptr) {return compress_decompress(original_file, sizes, target_error, compressed_file, ptr, options);}, sizes, options);
	}
	
	template<typename dtype>
	TypedResults<dtype> compress_decompress(const std::string& original_file, const std::vector<size_t>& sizes, double target_error, const Options& options) {
		return vectorize_compress_decompress<dtype>([&] (dtype* ptr) {return compress_decompress(original_file, sizes, target_error, ptr, options);}, sizes, options);
	}
	
	template<typename dtype>
	TypedResults<dtype> compress_decompress_to_memory(const dtype* original_data, const std::vector<size_t>& sizes, double target_error, const std::string& compressed_file, const Options& options) {
		return vectorize_compress_decompress<dtype>([&] (dtype* ptr) {return compress_decompress(original_data, sizes, target_error, compressed_file, ptr, options);}, sizes, options);
	}
	
	template<typename dtype>
	TypedResults<dtype> compress_decompress(const dtype* original_data, const std::vector<size_t>& sizes, double target_error, const Options& options) {
		return vectorize_compress_decompress<dtype>([&] (dtype* ptr) {return compress_decompress(original_data, sizes, target_error, ptr, options);}, sizes, options);
	}
	
	// read_data_type_implementations
	
	DataType read_data_type(const std::string& compressed_file) {
		CompressedIO io;
		io.open_read(compressed_file);
		return CompressorBase().read_data_type(io);
	}
	
	DataType read_data_type(size_t compressed_size, const char* compressed_data) {
		CompressedIO io;
		io.open_read(compressed_data, compressed_size);
		return CompressorBase().read_data_type(io);
	}
	
	DataType read_data_type(const std::vector<char>& compressed_data) {
		return read_data_type(compressed_data.size(), compressed_data.data());
	}
	
	// read_decompressed_size implementations
	
	size_t read_decompressed_size(const std::string& compressed_file, const Options& options) {
		return read_decompressed_size_internal(compressed_file, options, false);
	}
	
	size_t read_decompressed_size(size_t compressed_size, const char* compressed_data, const Options& options) {
		return read_decompressed_size_internal(compressed_size, compressed_data, options, false);
	}
	
	size_t read_decompressed_size(const std::vector<char>& compressed_data, const Options& options) {
		return read_decompressed_size(compressed_data.size(), compressed_data.data(), options);
	}
	
	// calculate_decompressed_size implementation
	
	size_t calculate_decompressed_size(const std::vector<size_t>& sizes, const Options& options) {
		return calculate_decompressed_size_internal(sizes, options, false);
	}
	
	#include "explicit_instantiations/atc.txt"
	
}
