#pragma once

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

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>> 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){
		std::nth_element(begin(d), begin(d) + slack, end(d));
		return (begin(d) + rand())->second;
	});
}