#include <algorithm>
#include <random>

#include <rosa/els.h>
#include <rosa/grasp.h>

#include <muscles/descent.h>
#include <muscles/move2opt.h>
#include <muscles/rgreedy.h>

#include <tsp/problem.h>
#include <tsp/solution.h>
#include <tsp/tsp.h>

#include "celero/Celero.h"
#include "inc/udm.h"

namespace {

namespace bench {

constexpr std::size_t SAMPLES     = 10;
constexpr std::size_t ITERATIONS  = 1;

tsp::Solution selectMin(tsp::Solution const&a, tsp::Solution const&b) { return a<b? a:b; }

}


class GraspElsFixture: public celero::TestFixture {
	std::shared_ptr<tsp::Problem> _problem;
	std::shared_ptr<GetRusageUDM> _getRusageUDM;

	GetRusage _getRusage;

public:
	static constexpr unsigned graspN      = 128;
	static constexpr unsigned elsIterMax  = 5;
	static constexpr unsigned elsGen      = 10;

	static auto rgreedy() { return RGreedy<tsp::Solution>{2}; }

public:
	GraspElsFixture(): _getRusageUDM{new GetRusageUDM} {}

	void setUp(ExperimentValue const&) override {
		tsp::Tsp tspData{"../data/dj38"};
		_problem.reset(new tsp::Problem{tspData.points()});

		_getRusage.start(bench::ITERATIONS);
	}

	void tearDown() override {
		_getRusage.stop();
		_getRusageUDM->addValue(_getRusage.get());
	}

	// std::vector<std::shared_ptr<celero::UserDefinedMeasurement>> getUserDefinedMeasurements() const override {
		// return {_getRusageUDM};
	// }

public:
	tsp::Problem const&problem() const { return *_problem; }

};


using RNG = std::mt19937;

auto hwElsGen(tsp::Solution const&solution, RNG &rng) {
	return Descent{}(Move2Opt{}(solution, rng));
}

auto hwElsInner(tsp::Solution const&solution, RNG &rng) {
	tsp::Solution best;

	if(GraspElsFixture::elsGen)
		best = hwElsGen(solution, rng);
	for(std::size_t i = 1; i < GraspElsFixture::elsGen; ++i) {
		tsp::Solution current = hwElsGen(solution, rng);
		best = bench::selectMin(std::move(best), std::move(current));
	}

	return best;
}

auto hwEls(tsp::Solution const&solution, RNG &rng) {
	tsp::Solution best = Descent{}(solution);

	for(std::size_t i = 0; i < GraspElsFixture::elsIterMax; ++i) {
		tsp::Solution current = hwElsInner(best, rng);
		best = bench::selectMin(std::move(best), std::move(current));
	}

	return best;
}

auto hwGraspGen(tsp::Problem const&problem, RNG &rng) {
	return hwEls(GraspElsFixture::rgreedy()(problem, rng), rng);
}

auto hwGraspEls(tsp::Problem const&problem, RNG &rng) {
	tsp::Solution best;

	if(GraspElsFixture::graspN)
		best = hwGraspGen(problem, rng);
	for(std::size_t i = 1; i < GraspElsFixture::graspN; ++i) {
		tsp::Solution current = hwGraspGen(problem, rng);
		best = bench::selectMin(std::move(best), std::move(current));
	}

	return best;
}

template<std::size_t K>
auto hwGraspElsPar(tsp::Problem const&problem, RNG &rng) {
	std::size_t const nThreads = std::min<std::size_t>(K, GraspElsFixture::graspN);
	std::size_t const step = GraspElsFixture::graspN/nThreads;
	std::size_t remain = GraspElsFixture::graspN - step*nThreads;

	tsp::Solution best;

	std::vector<std::thread> threads{nThreads-1};
	std::vector<tsp::Solution> solutions(nThreads);

	for(std::size_t i{}; i < (nThreads-1); ++i) {
		std::size_t offset = !!remain;
		remain -= offset;

		threads[i] = std::thread{
			[&,i,step=step+offset](auto const&problem, auto rng) {
				tsp::Solution &s = solutions[i];

				for(std::size_t j{}; j < step; ++j) {
					tsp::Solution cur = hwGraspGen(problem, rng);
					s = bench::selectMin(std::move(s), std::move(cur));
				}
			},
			std::cref(problem), rng
		};
	}

	{
		tsp::Solution &s = solutions[nThreads-1];

		for(std::size_t j{}; j < step; ++j) {
			tsp::Solution cur = hwGraspGen(problem, rng);
			s = bench::selectMin(std::move(s), std::move(cur));
		}
	}

	for(auto &thread: threads) thread.join();

	best = *std::min_element(std::begin(solutions), std::end(solutions));

	return best;
}

using ELS       = rosa::SkelEls<
	tsp::Solution,
	Descent,
	Move2Opt, Descent, FN(bench::selectMin)
>;

using GRASPxELS = rosa::SkelGrasp<
	tsp::Problem, tsp::Solution,
	RGreedy<tsp::Solution>, ELS,
	FN(bench::selectMin)
>;

}

BASELINE_F(GraspEls_Seq, Handwritten, GraspElsFixture, bench::SAMPLES, bench::ITERATIONS) {
	RNG rng;

	celero::DoNotOptimizeAway(
		hwGraspEls(problem(), rng)
	);
}

BENCHMARK_F(GraspEls_Seq, Skeleton, GraspElsFixture, bench::SAMPLES, bench::ITERATIONS) {
	auto graspEls = alsk::implement<alsk::exec::Sequential, GRASPxELS>();
	graspEls.skeleton.task.task<0>()                  = rgreedy();
	graspEls.skeleton.task.task<1>().task<1>().n      = elsIterMax;
	graspEls.skeleton.task.task<1>().task<1>().task.n = elsGen;
	graspEls.skeleton.n                               = graspN;

	celero::DoNotOptimizeAway(
		graspEls(problem())
	);
}

/* *** */

BASELINE_F(GraspEls_Par, Handwritten, GraspElsFixture, bench::SAMPLES, bench::ITERATIONS) {
	RNG rng;

	celero::DoNotOptimizeAway(
		hwGraspElsPar<2>(problem(), rng)
	);
}

BENCHMARK_F(GraspEls_Par, Skeleton, GraspElsFixture, bench::SAMPLES, bench::ITERATIONS) {
	auto graspEls = alsk::implement<alsk::exec::FirstLevelEqui, GRASPxELS>();
	graspEls.executor.cores                           = 4;
	graspEls.skeleton.task.task<0>()                  = rgreedy();
	graspEls.skeleton.task.task<1>().task<1>().n      = elsIterMax;
	graspEls.skeleton.task.task<1>().task<1>().task.n = elsGen;
	graspEls.skeleton.n                               = graspN;

	celero::DoNotOptimizeAway(
		graspEls(problem())
	);
}