#ifndef GENS_H_INCLUDED
#define GENS_H_INCLUDED

#include <vector>
#include <thread>
#include <future>

#include "rnn.h"

#define FRC_NO_ICU 1
#include "FRC.h"

namespace generations {

int maxThreads;

template <typename Species, typename TestData>
struct Individual {
	Individual(pcg32_k64& R, long long f = 0)
		: ind(R), fitness(f) {;}
	template <typename T>
	Individual(T s, long long f = 0)
		: ind(s), fitness(f) {;}
	Individual()
		: ind(), fitness(0) {;}
	using TestData_t = TestData;
	Species ind;
	long long fitness;
	long long test(const TestData&, pcg32_k64& _r, std::mutex* m=nullptr);
	Individual mutate(int divergence, pcg32_k64& _r) const {
		return {ind.mutate(divergence, _r), fitness};
	}
};

using rnnTestData = std::tuple<
	std::vector<std::pair<std::string, long long>>,
	long long,
	path::spacialRepresentation*const>;
using defRNN = rnnStack<16, 16, 2>;
constexpr const static int rnnGenSize = 100;
using rnnInd = Individual<defRNN, rnnTestData>;

template <typename Individual>
inline bool operator<(const Individual& a, const Individual& b)
	{return a.fitness<b.fitness;}

template <typename Individual>
inline Individual breed(const Individual& p1, const Individual& p2, pcg32_k64& _r)
{
	return {breed(p1.ind, p2.ind, _r), (p1.fitness + p2.fitness)/2};
}

template<>
inline rnnInd breed(const rnnInd& p1, const rnnInd& p2, pcg32_k64& _r)
{
	return {
		p1.ind.mutate_toward(p2.ind, float(p2.fitness-p1.fitness)/p1.fitness, _r),
		// p1.fitness-(p1.fitness/p2.fitness)
		std::min(p1.fitness, p2.fitness)
	};
}

template <typename Individual, typename Container, int gensize>
class gen{
public:
	gen(pcg32_k64& R)
		: pool(build<Container>(
			[&R](){return Individual{R};}, gensize
		)), r(&R) {;}
	gen(Container&& c, pcg32_k64 *R)
		: pool(c), r(R) {;}
	
	//Create a new generation
	//Elitist selection: very best is kept unchanged
	//Very worst is killed, replaced by random immigrant (keeping its fitness)
	//The other n-2 are either bred or mutated, stochastically favoring the best
	gen advance(int d) const
	{
		static_assert(gensize > 1, "gensize must be >= 2");
		auto& _r = *r;
		auto npool = pool;
		std::sort(npool.begin(), npool.end());
		npool.back() = npool.back().mutate(d, _r);
		int bias = npool.back().fitness + std::sqrt(npool.back().fitness);
		std::array<int, 3> b_m = {gensize*2, gensize*1, gensize/2};
		auto fitnesses = build<std::array<long long, gensize>>(
			pool.begin(), pool.end(),
			[bias](const Individual& i) {return bias - i.fitness;}
		);
		for (int i = 1; i < gensize; ++i) {
			auto b = chooseCategorical(b_m, _r);
			--b_m[b];
			if (b==0) {
				npool[i] = pool[chooseCategorical(fitnesses, _r)].mutate(_r(d)+d, _r);
			} else if (b==1){
				int a = chooseCategorical(fitnesses, _r),
					b = chooseCategorical(fitnesses, _r);
				while (b == a)
					{b = chooseCategorical(fitnesses, _r);}
				npool[i] = breed(pool[a], pool[b], _r).mutate(_r(d)+d/2, _r);
			} else {
				npool[i] = Individual{_r, 0};
			}
		}
		return {std::move(npool), r};
	}
	
	template <bool print=false>
	long long test(typename Individual::TestData_t const& t) {
		auto& _r = *r;
		if constexpr (print) {
			for (size_t i = 0; i != gensize; ++i) {
				std::cout<<i*100/gensize<<"% "<<std::flush;
				pool[i].test(t, *r);
			}
		} else {
			auto worker_ranges = buildiota<std::vector<size_t>>(
				maxThreads, 0, gensize/maxThreads
			);
			worker_ranges.push_back(gensize);
			// for (auto i : worker_ranges) {
				// std::cout<<i<<' ';
			// }
			// std::cout<<std::endl;
			std::vector<std::future<void>> threads;
			for (int i = 0; i != maxThreads; ++i) {
				size_t start = worker_ranges[i], end = worker_ranges[i+1];
				if (start != end) {
					threads.push_back(std::async(
						// std::launch::async,
						[&,start,end]{
							auto R=_r;
							std::for_each(
								pool.begin()+start,
								pool.begin()+end,
								[&](Individual& i){i.test(t, R, nullptr);}
						);}
					));
					_r.advance(_r());
				}
			}
			for (auto& t : threads) {
				t.wait();
			}
			// for (auto& i : pool) {
				// std::cout<<i.fitness<<' ';
			// }
			// std::cout<<std::endl;
			// std::terminate();
		}
		std::sort(pool.begin(), pool.end());
		return pool.front().fitness;
	}
	
	//Only test those individuals which have no fitness score yet
	long long quickTest(typename Individual::TestData_t const& t) {
		auto& _r = *r;
		std::for_each(
			// std::execution::par_unseq,
			pool.begin(), pool.end(),
			[&](Individual& i){
				if (!i.fitness)
					i.test(t, _r);
			}
		);
		std::sort(pool.begin(), pool.end());
		return pool.front().fitness;
	}
	
	const Container& exportInds() const {return pool;}
	int size() const {return gensize;}
private:
	Container pool;
	pcg32_k64 *r;
};

template <>
long long rnnInd::test(const rnnTestData& t, pcg32_k64& r, std::mutex* m)
{
	long long accum = 0;
	for (auto& w : std::get<0>(t)) {
		auto path = path::wordToPath(std::get<2>(t), w.first);
		auto out = ind.stepOver(path);
		{
			std::unique_lock<std::mutex> l;
			if (m) {
				l = std::unique_lock<std::mutex>{*m};
			}
			auto test = path::wordFromPath(path, out, r);
			float error = editDistance(
				w.first,
				test
			);
			accum += (error / w.first.length()) * w.second/std::get<1>(t);
		}
	}
	return fitness = accum/std::get<0>(t).size();
}

using kTestData = std::tuple<
	int,
	long long,
	const std::vector<std::pair<std::string, long long>>&,
	int>;
constexpr const static int kgen_size = 50;
using kInd = Individual<sKeymap, kTestData>;
using kgen_array = std::array<kInd, kgen_size>;
using kgen = gen<kInd, kgen_array, kgen_size>;
using rnnGen = gen<
	rnnInd,
	std::array<rnnInd, rnnGenSize>,
	rnnGenSize>;
template <>
long long kInd::test(const kTestData& t, pcg32_k64& r, std::mutex* m)
{
	std::cout<<"\'"<<keymap(ind.exportKeys())<<"\' h="<<fitness<<std::endl;
	std::unique_lock<std::mutex> l;
	if (m) {
		l = std::unique_lock<std::mutex>{*m};
	}
	rnnGen rnns{r};
	auto k_s = path::keymapToSpace(ind);
	auto last = high_resolution_clock::now();
	float best = 1;
	for (int i = 0; i <= std::get<0>(t); ++i) {
		if (i)
			rnns = rnns.advance(20-(10*static_cast<float>(i)/std::get<0>(t)));
		rnns.test(std::make_tuple(
			buildsample(std::get<2>(t),std::get<3>(t), r),
			std::get<1>(t),
			k_s.get()
		));
		auto& rs = rnns.exportInds();
		if (i==std::get<0>(t)) {
			auto now = high_resolution_clock::now();
			std::cout
				<<std::setw(5)<<rs.front().fitness<<" (r = "
				<<std::fixed<<std::setprecision(2)<<float(rs.back().fitness)/rs.front().fitness
				<<"); ("<<std::setw(2)<<static_cast<long long>((1-(rs.front().fitness/best))*100)
				<<"% in "<<std::setw(3)<<duration_cast<seconds>(now - last).count()<<"s)"<<std::endl;
			last = now;
		} else if (i%2 == 0) {
			std::cout
				<<std::setw(5)<<rs.front().fitness<<" (r = "
				<<std::fixed<<std::setprecision(2)<<float(rs.back().fitness)/rs.front().fitness
				<<"); "<<std::flush;
			if (i == 0)
				best = rs.front().fitness;
		}//*/
	}
	return fitness = rnns.exportInds().front().fitness;
}

// long long nullTest(const rnnTestData& t, [[maybe_unused]] pcg32_k64& r)
// {
	// long long accum = 0;
	// for (auto& w : std::get<0>(t)) {
		// float error = editDistance(
			// w.first,
			// ""s
		// );
		// accum += (error / w.first.length()) * w.second/std::get<1>(t);
	// }
	// return accum/std::get<0>(t).size();
// }
template <int yes, int no>
long long nullTest(const rnnTestData& t, pcg32_k64& r)
{
	long long accum = 0;
	for (auto& w : std::get<0>(t)) {
		auto path = path::wordToPath(std::get<2>(t), w.first);
		Eigen::VectorXf all(2);
		all<<yes,no;
		auto test = path::wordFromPath(
			path, std::vector<Eigen::VectorXf>(path.size(), all), r
		);
		float error = editDistance(
			w.first,
			test
		);
		accum += (error / w.first.length()) * w.second/std::get<1>(t);
	}
	return accum/std::get<0>(t).size();
}

// long long nullTest3(const rnnTestData& t, pcg32_k64& r)
// {
	// long long accum = 0;
	// for (auto& w : std::get<0>(t)) {
		// auto path = path::wordToPath(std::get<2>(t), w.first);
		// Eigen::VectorXf all(2);
		// all<<1,1;
		// auto test = path::wordFromPath(
			// path, std::vector<Eigen::VectorXf>(path.size(), all), r
		// );
		// float error = editDistance(
			// w.first,
			// test
		// );
		// accum += (error / w.first.length()) * w.second/std::get<1>(t);
	// }
	// return accum/std::get<0>(t).size();
// }

}

#endif //GENS_H_INCLUDED