#include "wavelet.hpp"
#include <wavelet.hpp>
#include <utilities.hpp>
#include <image_io.hpp>

#include <ostream>
#include <cmath>

static const int size = 512;

rgb_wavelet_coefficients::pre_fingerprint rgb_wavelet_coefficients::pre_calculate(av::frame const & frame) {
	auto const image = to_raw_rgb_image(crop_to_square(frame), size, size);
	rgb_wavelet_coefficients::pre_fingerprint ret;

	// for every color
	for(unsigned int color = 0; color < 3; ++color){
		auto & vector = ret[color];
		vector.assign(make_u(image.width() * image.height()), 0);

		for(unsigned int n = 0; n < make_u(image.width() * image.height()); ++n){
			vector[n] = 2.0 * image.data[3*n + color] / double(255) - 1.0;
		}

		wvlt::wavelet_2D(vector.data(), make_u(image.width()), make_u(image.height()));
	}

	return ret;
}

rgb_wavelet_coefficients rgb_wavelet_coefficients::calculate(av::frame const & frame){
	auto const image = to_raw_rgb_image(crop_to_square(frame), size, size);
	rgb_wavelet_coefficients ret;

	std::vector<double> vector(make_u(image.width() * image.height()), 0);

	// for every color
	for(unsigned int color = 0; color < 3; ++color){
		auto& coefficient_array = color == 0 ? ret.reds : (color == 1 ? ret.greens : ret.blues);
		unsigned int array_index = 0;

		for(unsigned int n = 0; n < make_u(image.width() * image.height()); ++n){
			vector[n] = 2.0 * image.data[3*n + color] / double(255) - 1.0;
		}

		wvlt::wavelet_2D(vector.data(), make_u(image.width()), make_u(image.height()));

		auto copy = vector;
		for(auto & x : copy) x = std::abs(x);

		auto const n_coefficients = coefficient_array.size();
		std::nth_element(copy.begin(), copy.begin() + n_coefficients, copy.end(), std::greater<double>());
		auto const threshold = copy[n_coefficients-1];

		for(unsigned int n = 0; n < vector.size(); ++n){
			auto const x = vector[n];
			if(std::abs(x) >= threshold) {
				coefficient_array[array_index++] = std::make_pair(n, x);
			}
			if(array_index >= coefficient_array.size()) {
				break;
			}
		}
	}

	return ret;
}

static double square(double x){
	return x*x;
}

double rgb_wavelet_coefficients::distance_to(pre_fingerprint const & fingerprint) const {
	double distance = 0;

	for(unsigned int color = 0; color < 3; ++color){
		auto const & coefficients = color == 0 ? reds : (color == 1 ? greens : blues);

		for(auto&& p : coefficients){
			auto const x = p.second;
			auto const y = fingerprint[color][p.first];
			distance += square(x - y) - square(y);
		}
	}

	return distance;
}

namespace std {
	template <typename U, typename V>
	ostream& operator<<(ostream& out, pair<U, V> const & p){
		return out << '(' << p.first << ", " << p.second << ')';
	}
}

std::ostream& operator<<(std::ostream& out, rgb_wavelet_coefficients const & x){
	out << "rgb_wavelet_coefficients" << std::endl;

	for(int color = 0; color < 3; ++color){
		auto const & coefficient_array = color == 0 ? x.reds : (color == 1 ? x.greens : x.blues);

		out << '[' << coefficient_array[0];
		for(unsigned int i = 1; i < coefficient_array.size(); ++i) {
			out << ", " << coefficient_array[i];
		}
		out << ']' << std::endl;
	}
	return out;
}