/*
 * This file implements the ATC command-line interface executable.
 */

#include <iostream>
#include <fstream>
#include <string>
#include <vector>
#include <algorithm>
#include <type_traits>
#include <boost/program_options.hpp>

#include "atc.hpp"
#include "atc/version.hpp"

namespace po = boost::program_options;

/*
 * Parsers for special ATC data types
 * These are defined in the ATC namespace so Boost finds them when it needs them,
 * else we have problems with name resolution
 */

namespace atc {
	
	std::istream& operator>>(std::istream& in, atc::StorageOrder& order) {
		std::string token;
		in >> token;
		if (token == "msi")
			order = atc::StorageOrder::MSI;
		else if (token == "lsi")
			order = atc::StorageOrder::LSI;
		else
			in.setstate(std::ios_base::failbit);
		return in;
	}

	std::istream& operator>>(std::istream& in, atc::FloatType& float_type) {
		std::string token;
		in >> token;
		if (token == "float32")
			float_type = atc::FloatType::float32;
		else if (token == "float64")
			float_type = atc::FloatType::float64;
		else
			in.setstate(std::ios_base::failbit);
		return in;
	}

	std::istream& operator>>(std::istream& in, atc::DataType& data_type) {
		std::string token;
		in >> token;
		if (token == "uint8")
			data_type = atc::DataType::uint8;
		else if (token == "uint16")
			data_type = atc::DataType::uint16;
		else if (token == "uint32")
			data_type = atc::DataType::uint32;
		else if (token == "uint64")
			data_type = atc::DataType::uint64;
		else if (token == "int8")
			data_type = atc::DataType::int8;
		else if (token == "int16")
			data_type = atc::DataType::int16;
		else if (token == "int32")
			data_type = atc::DataType::int32;
		else if (token == "int64")
			data_type = atc::DataType::int64;
		else if (token == "float32")
			data_type = atc::DataType::float32;
		else if (token == "float64")
			data_type = atc::DataType::float64;
		else
			in.setstate(std::ios_base::failbit);
		return in;
	}

	std::istream& operator>>(std::istream& in, atc::IntType& int_type) {
		std::string token;
		in >> token;
		if (token == "uint32")
			int_type = atc::IntType::uint32;
		else if (token == "uint64")
			int_type = atc::IntType::uint64;
		else
			in.setstate(std::ios_base::failbit);
		return in;
	}

	std::istream& operator>>(std::istream& in, atc::CoreFlattening& order) {
		std::string token;
		in >> token;
		if (token == "unfold")
			order = atc::CoreFlattening::Unfold;
		else if (token == "zigzag")
			order = atc::CoreFlattening::Zigzag;
		else if (token == "z-order")
			order = atc::CoreFlattening::ZOrder;
		else
			in.setstate(std::ios_base::failbit);
		return in;
	}

	std::istream& operator>>(std::istream& in, atc::EntropyCodingMethod& method) {
		std::string token;
		in >> token;
		if (token == "ac")
			method = atc::EntropyCodingMethod::Arithmetic;
		else if (token == "ans")
			method = atc::EntropyCodingMethod::ANS;
		else
			in.setstate(std::ios_base::failbit);
		return in;
	}

	std::istream& operator>>(std::istream& in, atc::HouseholderWeightingMethod& method) {
		std::string token;
		in >> token;
		if (token == "approx")
			method = atc::HouseholderWeightingMethod::Approximation;
		else if (token == "csn")
			method = atc::HouseholderWeightingMethod::CoreSliceNorms;
		else if (token == "uniform")
			method = atc::HouseholderWeightingMethod::Uniform;
		else
			in.setstate(std::ios_base::failbit);
		return in;
	}

	std::istream& operator>>(std::istream& in, atc::Slice& slice) {
		std::string token;
		in >> token;
		slice = atc::Slice(token);
		return in;
	}
	
}

/**
 * Generates a function that when called with a vector option value string, parses the string and stores it in the pre-defined location.
 * This allows us to parse std::vector arguments ourselves, which aren't properly supported in config files.
 * Should only be called with targets that will stay alive until notify is called.
 */
template<typename T>
auto generate_notifier(std::vector<T>& target) {
	return [&] (const std::string& value) {
		if (value.empty())
			throw boost::program_options::validation_error(boost::program_options::validation_error::kind_t::at_least_one_value_required);
		std::istringstream iss(value);
		target.resize(0);
		while (!iss.eof()) {
			T val;
			iss >> val;
			if (iss.fail())
				throw boost::program_options::validation_error(boost::program_options::validation_error::kind_t::invalid_option_value, "", value);
			target.push_back(val);
		}
	};
}

/**
 * This class provides a generic way to add options to an options_description object similar to options_description_easy_init.
 * The main difference is a second overload of operator() for std::vector's, which adjusts the option to either a CLI or config file setting.
 */
class AddOptions {
public:
	AddOptions(po::options_description& desc, bool config_file) : desc(desc), config_file(config_file) { }
	AddOptions& operator()(const std::string& name, po::value_semantic* value_semantic, const std::string& option_description) {
		desc.add_options()(name.c_str(), value_semantic, option_description.c_str());
		return *this;
	}
	template<typename T>
	AddOptions& operator()(const std::string& name, std::vector<T>& target, const std::string& option_description) {
		if (config_file)
			desc.add_options()(name.c_str(), po::value<std::string>()->multitoken()->notifier(generate_notifier(target)), option_description.c_str());
		else
			desc.add_options()(name.c_str(), po::value<std::vector<T>>(&target)->multitoken(), option_description.c_str());
		return *this;
	}
private:
	po::options_description& desc;
	const bool config_file;
};

// Prints an error message based on the given message and terminates the program with exit code 1 (indicating an error).
void error(const std::string& message) {
	std::cout << "Error: " << message << std::endl;
	std::cout << "For usage information, use --help." << std::endl;
	exit(1);
}

/*
 * Converts a vector of integers in the range [0, 255].
 * This is handled separately because we cannot treat these integers as uint8_t, because then they would be equivalent to char
 * and the tokens read from the command line would be interpreted using their ASCII value instead of parsing them as integers.
 */
std::vector<uint8_t> convert_vector_of_bytes(const std::vector<uint16_t>& vector, const std::string& name) {
	if (vector.size() > 0 && *std::max_element(vector.begin(), vector.end()) > uint16_t(std::numeric_limits<uint8_t>::max()))
		error("Given " + name + " contains values over 255");
	return std::vector<uint8_t>(vector.begin(), vector.end());
}

int main(int argc, char** argv) {
	
	atc::Options options;
	bool help = false, print_results = false;
	std::string config_file, original_file, compressed_file, decompressed_file;
	atc::DataType data_type;
	std::vector<size_t> sizes;
	double target_error;
	std::vector<uint16_t> mode_processing_order, mode_storage_order; // We can't read directly into std::vector<uint8_t> because then tokens are interpreted as chars
	bool no_rank_truncation = !options.rank_truncation, no_householder = !options.householder, no_split_bit_plane_truncation = !options.split_bitplane_truncation, epcm = options.explicitly_parallellized_core_multiplications;
	
	// Default values are not specified for most options because we use the defaults from the ATC library itself (see Options.hpp)
	
	// To work around certain limitations of Boost.Program_options config files, we define this function to easily create two instances of the same options, one adapted to each setting.
	// We define it as a lambda instead of as a separate function due to the high coupling with local variables.
	auto generate_options_description = [&] (bool use_config_file) { 
		
		po::options_description general_options("General");
		AddOptions(general_options, use_config_file)
			("help", po::bool_switch(&help), "Produce this help message")
			("config", po::value<std::string>(&config_file)->default_value(""), "Configuration file, storing a set of options and their values, in the format \"<option> = <value>\" with each option on its own line (INI). Boolean options can be set using e.g. \"true\". If an option is provided both in the command line and the given configuration file, the command line will take priority.")
			("of", po::value<std::string>(&original_file)->default_value(""), "Original data file, used as input for compression")
			("cf", po::value<std::string>(&compressed_file)->default_value(""), "Compressed file, used as output for compression and/or input for decompression")
			("df", po::value<std::string>(&decompressed_file)->default_value(""), "Decompressed file, used as output for decompression")
			("print-results", po::bool_switch(&print_results), "Print compression/decompression time as well as compressed file size and compression factor in the case of compression")
			("verbose", po::bool_switch(&options.verbose), "Print progress/status updates as well as timings in case --print-results is used")
			("suppress", po::bool_switch(&options.suppress_warnings), "Suppress warnings")
			("threads", po::value<int>(&options.threads), "Maximal amount of threads to be used, 0 (default) means that ATC will attempt to use the amount of cores available on the system.")
			("storage-order", po::value<atc::StorageOrder>(&options.storage_order), "The ordering of the tensor modes which determines the tensor flattening. Possible values: msi (default): most-significant-index first (C order, row-major for matrices), lsi: least-significant-index first (Fortran order, column-major for matrices).")
			("float-type", po::value<atc::FloatType>(&options.float_type), "The floating-point type used for most computational work internally. Possible values: float32, float64 (default)")
		;
		
		po::options_description compression_options("Compression");
		AddOptions(compression_options, use_config_file)
			("type", po::value<atc::DataType>(&data_type), "Data type of the original file. Possible values: uint8, uint16, uint32, uint64, int8, int16, int32, int64, float32, float64")
			("sizes", sizes, "Sizes of the tensor modes, separated by spaces")
			("psnr", po::value<double>(&target_error), "Target error in terms of PSNR (in dB)")
			("re", po::value<double>(&target_error), "Target error in terms of relative error")
			("rmse", po::value<double>(&target_error), "Target error in terms of RMSE")
			("skip-bytes", po::value<size_t>(&options.skip_bytes), "The number of bytes to be skipped in the original data file (default: 0)")
			("rank-truncation", po::value<double>(&options.rank_truncation_max_sse_share), "Rank truncation maximal SSE share: upper bound to ratio of SSE introduced by rank truncation (by default only from the ST-HOSVD, optionally also from the T-HODCT) to target SSE, should be in the range [0, 1] (default: 0.5).")
			("dct", po::bool_switch(&options.apply_dct), "Apply T-HODCT preprocessing")
			("dct-truncation", po::value<double>(&options.dct_truncation_max_sse_share), "DCT truncation maximal SSE share: upper bound to ratio of SSE introduced in T-HODCT to SSE introduced in all rank truncation (both in the T-HOSVD and the ST-HOSVD), should be in the range [0, 1] (default: 0).")
			("mode-processing-order", mode_processing_order, "ST-HOSVD mode processing order, specified by enumerating the mode indices (0-based) in the desired order, separated by spaces. Default: ST-HOSVD heuristic, smallest modes first.")
			("precise-svd", po::bool_switch(&options.precise_svd), "Compute the SVD's using the significantly slower but more precise Eigen::BDCSVD solver instead of using the eigenvalue decomposition of the Gramian matrix. Disabled by default because this added precision normally has no compression effect in practice.")
			("print-sq-sv", po::bool_switch(&options.print_squared_singular_values), "If the verbose flag is set, prints the list of computed squared singular values (or eigenvalues from the Gramian matrix) for each mode during the ST-HOSVD.")
			("no-householder", po::bool_switch(&no_householder), "Disables Householder-reflector-based factor compression.")
			("int-type", po::value<atc::IntType>(&options.int_type), "The integer type used for quantization, a smaller type can be sufficient for high target errors. Possible values: uint32, uint64 (default)")
			("core-flattening", po::value<atc::CoreFlattening>(&options.core_flattening), "The ordering (permutation) of core coefficients that should be used to flatten the core for quantization. Possible values: unfold (default), zigzag, z-order")
			("mode-storage-order", mode_storage_order, "The order of the tensor modes in the internal representation of the core before quantization, can have a significant effect on compression efficiency in the case the default is core encoding order. Specified by enumerating the mode indices (0-based) in the desired order, separated by spaces. If not specified, ATC determines this order heuristically.")
			("qb", po::value<size_t>(&options.quantization_blocks), "Quantization blocks: the amount of blocks the quantized core should be split into in order to parallelize quantization and dequantization. Default: 0, which means the number of blocks will equal the number of requested threads.")
			("lpmt", po::value<size_t>(&options.last_pass_max_threads), "The maximal number of threads that should be used during the last pass of the quantization process. During this pass the SSE reduction is tracked and synchronized in between all threads. To limit this overhead it is best not to use arbitrarily many threads in this step. Default: 8")
			("qrem", po::value<double>(&options.quantization_rel_error_margin), "Sets the quantization relative error margin, should be a value > 0 and < 1 (default: " + std::to_string(options.quantization_rel_error_margin) + "). Decreasing this value will improve error control but slow down quantization, especially when the last_pass_max_threads parameter is high.")
			("no-split-bit-plane-truncation", po::bool_switch(&no_split_bit_plane_truncation), "Disables split bit plane truncation.")
			("entropy-coding-method", po::value<atc::EntropyCodingMethod>(&options.entropy_coding_method), "The entropy coding method used to encode the run-lengths obtained during quantization. Possible values: ac (arithmetic coding, default), ans (asymmetric numeral systems)")
		;
		
		po::options_description decompression_options("Decompression");
		AddOptions(decompression_options, use_config_file)
			("cutout", options.cutout, "Slices (at most one per mode) to be taken from the data during decompression, using NumPy-style notation (e.g. one slice could be \"5:-1:3\"), with the individual slices separated by spaces. If \"b\" or \"l\" is appended to a slice string, respectively box or Lanczos filtering will be used as the reduction for that slice. For more information check the ATC documentation.")
			("autocrop", po::value<double>(&options.autocrop_threshold), "Crop away the first/last slices with too little energy in the decompressed tensor. The argument is the autocrop threshold, first/last rows in factors whose root-mean-squared value is smaller than this value will be removed. This value should be in the range [0, 1].")
		;
		
		po::options_description compression_decompression_options("Compression + decompression (in case --of, --cf and --df are provided)");
		AddOptions(compression_decompression_options, use_config_file)
			("ce", po::bool_switch(&options.calculate_error), "Calculate and print the actual error of the decompressed data, using the same error metric as the target error if specified, otherwise PSNR by default.")
		;
		
		// Research options, these were necessary for experiments but are not meant for users so they are not printed as part of the help message
		// TODO: Clean up
		po::options_description research_options("Research");
		AddOptions(research_options, use_config_file)
			("ignore-dequantization-correction", po::bool_switch(&options.ignore_dequantization_correction), "Do not consider the dequantization correction for significant coefficients (unknown bits are estimated to be 1/2) during error control.")
			("tthresh-factor-stopping-criterion", po::bool_switch(&options.tthresh_factor_stop_criterion), "Use the old TTHRESH factor stopping criterion, without applying the normalization adjustment.")
			("no-rank-truncation", po::bool_switch(&no_rank_truncation), "Disables rank truncation all-together during the ST-HOSVD phase. Slightly different from setting RTMSS to 0, which still allows truncation of 0-slices (which can occur in real data). Still allows truncation of 0-slices of the quantized core.")
			("static-rank-truncation-thresholding", po::bool_switch(&options.static_rank_truncation_thresholding), "Disables dynamic rank truncation thresholding during the ST-HOSVD.")
			("epcm", po::bool_switch(&epcm), "Enables explicitly parallellized core multiplications.")
			("hwm", po::value<atc::HouseholderWeightingMethod>(&options.householder_weighting_method), "Householder weighting method. Options: approx (Approximation), csn (CoreSliceNorms), Uniform (uniform)")
		;
		
		po::options_description all_options("Options");
		all_options
			.add(general_options)
			.add(compression_options)
			.add(decompression_options)
			.add(compression_decompression_options)
			.add(research_options)
		;
		
		std::string usage_string;
		if (!use_config_file) {
			// Construct usage string already too
			std::stringstream ss;
			ss << "ATC version " << atc::version << "\n\n"
				<< "Usage:\n\n"
				<< "Compression:            ./atc --of <original-file> --type <data-type> --sizes <tensor-sizes> --<error-metric> <target-error> --cf <compressed-file> [OPTIONS]\n"
				<< "Decompression:          ./atc --cf <compressed-file> --df <decompressed-file> [OPTIONS]\n"
				<< "Comp. + decomp.:        ./atc --of <original-file> --type <data-type> --sizes <tensor-sizes> --<error-metric> <target-error> --cf <compressed-file> --df <decompressed-file> [OPTIONS]\n\n"
				<< "Example:                ./atc --of ../data/test_data_3D.raw --type uint8 --sizes 3 3 3 --psnr 80 --cf ../data/comp.atc --df ../data/decomp.raw\n\n"
				<< "Options:\n\n"
				<< general_options << "\n"
				<< compression_options << "\n"
				<< decompression_options << "\n"
				<< compression_decompression_options << "\n";
			usage_string = ss.str();
		}
		
		return std::pair<po::options_description, std::string>(all_options, usage_string);
		
	};
	
	auto [cli_desc, cli_usage_string] = generate_options_description(false);
	auto config_file_desc = generate_options_description(true).first;
	
	// Parse options
	po::variables_map vm;
	try {
		
		// Import values from command line
		po::store(po::parse_command_line(argc, argv, cli_desc), vm);
		po::notify(vm);
		
		// Also import values from the config file, if specified
		if (!config_file.empty()) {
			std::ifstream ifs;
			ifs.open(config_file, std::ifstream::in);
			po::store(po::parse_config_file(ifs, config_file_desc), vm);
			po::notify(vm);
		}
		
	} catch (const std::exception& e) {
		error(e.what());
	}
	
	// Extra option handling which wasn't handled by option specification above
	options.mode_processing_order = convert_vector_of_bytes(mode_processing_order, "mode processing order");
	options.mode_storage_order = convert_vector_of_bytes(mode_storage_order, "mode storage order");
	options.timings = print_results;
	options.autocrop = vm.count("autocrop");
	options.rank_truncation = !no_rank_truncation;
	options.householder = !no_householder;
	options.split_bitplane_truncation = !no_split_bit_plane_truncation;
	options.explicitly_parallellized_core_multiplications = epcm;
	
	// Handle help message
	if (help) {
		std::cout << cli_usage_string << std::flush;
		return 0;
	}
	
	atc::Results results;
	if (!original_file.empty() && !compressed_file.empty()) {
		
		// Check presence of all required options
		if (!vm.count("type"))
			error("The data type must be specified when compressing.");
		size_t target_errors_count = vm.count("psnr") + vm.count("re") + vm.count("rmse");
		if (target_errors_count < 1)
			error("The target error must be specified when compressing.");
		if (target_errors_count > 1)
			error("Only one target error may be specified.");
		if (!vm.count("sizes"))
			error("The tensor sizes must be specified when compressing.");
		
		// Compression-specific extra option handling
		if (vm.count("psnr"))
			options.error_metric = atc::ErrorMetric::PSNR;
		else if (vm.count("re"))
			options.error_metric = atc::ErrorMetric::RelativeError;
		else if (vm.count("rmse"))
			options.error_metric = atc::ErrorMetric::RMSE;
		
		if (decompressed_file.empty()) {
			// compress
			try {
				results = atc::compress(original_file, data_type, sizes, target_error, compressed_file, options);
			} catch (const std::exception& e) {
				error(e.what());
			}
		} else {
			// compress_decompress
			try {
				results = atc::compress_decompress(original_file, data_type, sizes, target_error, compressed_file, decompressed_file, options);
			} catch (const std::exception& e) {
				error(e.what());
			}
		}
		
		if (options.calculate_error && !decompressed_file.empty()) {
			// Print regardless of print-results
			switch (options.error_metric) {
			case atc::ErrorMetric::PSNR:
				std::cout << "PSNR (dB)";
				break;
			case atc::ErrorMetric::RelativeError:
				std::cout << "Relative error";
				break;
			case atc::ErrorMetric::RMSE:
				std::cout << "RMSE";
				break;
			}
			std::cout << ": " << results.error << std::endl;
		}
		if (print_results) {
			std::cout << "Compressed size (bytes): " << results.compressed_size << std::endl;
			std::cout << "Compression factor: " << results.compression_factor << std::endl;
			std::cout << "Compression time (s): " << results.timings["compress"] << std::endl;
		}
		
	} else if (!compressed_file.empty() && !decompressed_file.empty()) {
		// decompress
		try {
			results = atc::decompress(compressed_file, decompressed_file, options);
		} catch (const std::exception& e) {
			error(e.what());
		}
	} else {
		error("Either the original and compressed file or compressed and decompressed file (or all three) must be specified.");
	}
	
	// Handles both compress_decompress and decompress
	if (print_results && results.timings.count("decompress"))
		std::cout << "Decompression time (s): " << results.timings["decompress"] << std::endl;
	
	return 0;
	
}
