#pragma once
#include <utilities.hpp>

#include <algorithm>
#include <random>
#include <set>
#include <string>
#include <vector>

template <typename T>
struct Mozaic {
	Mozaic(int h_tiles_, int v_tiles_)
	: h_tiles(h_tiles_)
	, v_tiles(v_tiles_)
	, tiles(v_tiles, std::vector<T>(h_tiles))
	{}

	decltype(auto) operator[](size_t c) { return tiles[c]; }
	decltype(auto) operator[](size_t c) const { return tiles[c]; }

	// aka width & height
	int h_tiles = 0, v_tiles = 0;

	// tiles[row][column]
	std::vector<std::vector<T>> tiles;
};

template <typename F, typename T>
auto mozaic_fmap(Mozaic<T> const & mozaic, F&& f){
	using return_type = decltype(f(mozaic[0][0]));
	Mozaic<return_type> ret(mozaic.h_tiles, mozaic.v_tiles);

	for(auto r = 0; r < mozaic.v_tiles; ++r){
		for(auto c = 0; c < mozaic.h_tiles; ++c){
			ret[r][c] = f(mozaic[r][c]);
		}
	}

	return ret;
}

//! Currently takes one of the (slack+1) best fitting images
template <typename Distance>
auto optimize(Mozaic<std::vector<Distance>> const & distances, int slack){
	std::random_device rd;
	std::uniform_int_distribution<int> dist(0, slack);
	auto rand = [&]{ return dist(rd); };

	return mozaic_fmap(distances, [&](auto d){
		auto o = rand();
		std::nth_element(begin(d), begin(d) + o, end(d));
		return (begin(d) + o)->second;
	});
}

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

//! Chooses best tiles left, starting from the center
template <typename Distance>
auto optimize_unique(Mozaic<std::vector<Distance>> const & distances){
	std::vector<std::tuple<double, size_t, size_t>> queue;
	queue.reserve(make_u(distances.h_tiles * distances.v_tiles));
	for(auto r = 0; r < distances.v_tiles; ++r){
		for(auto c = 0; c < distances.h_tiles; ++c){
			auto const aspect = distances.h_tiles / double(distances.v_tiles);
			auto const d = square(0.5*(distances.h_tiles) - c) + square(aspect*(0.5*(distances.v_tiles) - r));
			queue.emplace_back(d, r, c);
		}
	}

	std::random_device random_device;
	std::mt19937 generator(random_device());
	std::shuffle(begin(queue), end(queue), generator);
	std::sort(begin(queue), end(queue));

	using index = decltype(distances[0][0][0].second);
	std::set<index> used;
	Mozaic<index> ret(distances.h_tiles, distances.v_tiles);
	for(auto&& t : queue){
		auto const r = std::get<1>(t);
		auto const c = std::get<2>(t);
		auto v = distances[r][c];

		std::partial_sort(begin(v), begin(v) + used.size() + 1, end(v));
		for(auto&& p : v){
			if(used.count(p.second)) continue;
			used.insert(p.second);
			ret[r][c] = p.second;
			break;
		}
	}

	return ret;
}