#include <includes.hpp>
#include <utilities.hpp>

#include "jcmp.hpp"
#include "wavelet.hpp"
#include "remap.hpp"

#include <boost/filesystem.hpp>
#include <boost/program_options.hpp>
#include <png.hpp>

#include <vector>
#include <iostream>

using namespace wvlt;
using namespace jcmp;

// note: we take a copy, because we will modify it in place
static image compress(std::vector<double> const & img0, uint16_t width, double threshold, unsigned int& nzeros){
	uint16_t height = img0.size() / width;
	assert(is_pow_of_two(width));
	assert(is_pow_of_two(height));
	assert(width == height);

	std::vector<double> img(img0.size(), 0);

	for(uint16_t y = 0; y < height; y++){
		for(uint16_t x = 0; x < width; x++) {
			img[remap::to_hilbert(width, x, y)] = img0[x + y*width];
		}
	}

	wavelet(&img[0], width * height, 1);

#if 0
	// wavelet transform in x-direction
	for(unsigned int i = 0; i < height; ++i){
		wavelet(&img[i*width], width, 1);
	}

	// wavelet transform in y-direction
	for(unsigned int i = 0; i < width; ++i){
		wavelet(&img[i], height, width);
	}
#endif

	double min_abs = 10000.0;
	double max_abs = 0.0;

	for(auto& el : img){
		auto absel = std::abs(el);
		if(absel > threshold) {
			min_abs = std::min(min_abs, absel);
			max_abs = std::max(max_abs, absel);
		} else {
			el = 0;
		}
	}

	auto q = quantization::get(quantization::type::logarithmic, max_abs, min_abs);

	// save the principal coefficients
	std::vector<jcmp::coefficient> v;
	for(uint16_t y = 0; y < height; ++y){
		for(uint16_t x = 0; x < width; ++x){
			auto&& el = img[x + width*y];
			if(el != 0) v.push_back({q.forwards(el), uint(x), uint(y)});
		}
	}

	nzeros = v.size();
	return image(width, height, q.p, std::move(v));
}

static std::vector<double> decompress(image in){
	const auto width = in.header.width;
	const auto height = in.header.height;
	assert(is_pow_of_two(width));
	assert(is_pow_of_two(height));

	auto q = in.header.get_quantization(quantization::type::logarithmic);

	std::vector<double> image(width * height, 0.0);

	// read in coefficient on coordinates
	for(auto it = in.data.begin(); it != in.data.end(); ++it){
		auto&& coef = *it;
		image[coef.x + width*coef.y] = q.backwards(coef.c);
	}

	in.clear();

#if 0
	// inverse wavelet transform in y-direction
	for(unsigned int i = 0; i < width; ++i){
		unwavelet(&image[i], height, width);
	}

	// inverse wavelet transform in x-direction
	for(unsigned int i = 0; i < height; ++i){
		unwavelet(&image[i*width], width, 1);
	}
#endif

	unwavelet(&image[0], width*height, 1);

	std::vector<double> image2(width * height, 0.0);

	for(uint32_t d = 0; d < width * height; ++d) {
			const auto p = remap::from_hilbert(width, d);
			image2[p.first + width * p.second] = image[d];
	}

	return image2;
}

int main(int argc, char* argv[]){
	using namespace std;
	namespace fs = boost::filesystem;

	struct {
		vector<fs::path> filenames = {};
		double threshold = 5.0;
	} options;

	namespace po = boost::program_options;
	// Describe program options
	po::options_description opts;
	opts.add_options()
		("threshold", po::value<double>(), "")
		("help", po::bool_switch(), "show this help")
		("file", po::value< vector<fs::path> >(), "input file(s)");
	// and inputs
	po::positional_options_description p;
	p.add("file", -1);

	// Parse and set options
	try {
		po::variables_map vm;
		po::store(po::command_line_parser(argc, argv).
				  options(opts).positional(p).run(), vm);
		po::notify(vm);

		if(vm["help"].as<bool>()){
			std::cout << "jcmp" << std::endl;
			std::cout << opts << std::endl;
			return 0;
		}

		if(vm.count("file")){
			options.filenames = vm["file"].as< vector<fs::path> >();
		}

		if(vm.count("threshold")) {
			options.threshold = vm["threshold"].as<double>();
		}
	} catch(std::exception& e){
		std::cout << colors::red("ERROR: ") << e.what() << std::endl;
		std::cout << opts << std::endl;
		return 1;
	}

	for(const auto & path : options.filenames) {
		if(path.extension() != ".png") {
			cerr << "WARNING: can only read .png, but " << path << " is not" << endl;
		}

		// open file
		std::string filename = path.string();
		std::cout << field("file") << filename << std::endl;
		png::istream image(filename);

		auto width = image.get_width();
		auto height = image.get_height();

		// read into vector
		std::vector<double> image_vec;
		image_vec.reserve(width * height);
		for(unsigned char c = 0; image >> c;) image_vec.push_back(c/255.0);

		// compress and decompress to see how we lost information
		unsigned int nzeros = 0;
		auto compressed_vec = decompress(compress(image_vec, width, options.threshold, nzeros));

		// output some information
		std::cout << field("raw") << width*height << std::endl;
		std::cout << field("compressed") << nzeros/double(width*height) << std::endl;
		std::cout << field("nz") << nzeros << std::endl;

		// save to output file
		std::string cfilename = "hil_compressed_" + std::to_string(nzeros) + path.filename().string();
		png::gray_ostream compressed_image(width, height, cfilename);
		for(unsigned int i = 0; i < compressed_vec.size(); ++i) compressed_image << compressed_vec[i];
	}
}