#include <thread>
#include <mutex>
#include <tuple>
#include <cassert>
#include "compute_host.h"
#include "host_aberth.h"
#include "polygen.h"
#include "util.h"

using namespace std;


vector<int> computeHost(int W, int H, Com bottomLeft, Com topRight) {
	constexpr const int numThreads = 4;
	static_assert(ispow2(numThreads));

	cerr << "Finding roots for " << PolyGen::numPolys() << " polynomials" << endl;

	vector<int> counts(W * H);
	mutex countsMutex;

	vector<PolyGen::Job> jobs = PolyGen::genJobs(numThreads);
	assert(jobs.size() == numThreads);

	vector<thread> threads(jobs.size());
	for (int i = 0; i < (int)jobs.size(); i++) {
		threads[i] = thread([W, H, &counts, &countsMutex, job = jobs[i], bottomLeft, topRight]() {
			auto calcIndex = [](double value, double left, double right, int steps) -> int {
				return (value - left) / (right - left) * (steps - 1) + 0.5;
			};
			auto calcPos = [W, H, bottomLeft, topRight, &calcIndex](Com z) -> pair<int, int> {
				return make_pair(
					calcIndex(z.real(), bottomLeft.real(), topRight.real(), W),
					calcIndex(z.imag(), topRight.imag(), bottomLeft.imag(), H)
				);
			};

			vector<int> localCounts(W * H);

			Poly poly = job.init;
			for (int i = 0; i < job.numItems; i++) {
				for (Com z : aberth(poly)) {
					int x, y;
					tie(x, y) = calcPos(z);
					if (0 <= x && x < W && 0 <= y && y < H) {
						localCounts[W * y + x]++;
					}
				}
				PolyGen::next(poly);
			}

			lock_guard<mutex> guard(countsMutex);
			for (int i = 0; i < W * H; i++) counts[i] += localCounts[i];
		});
	}

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

	return counts;
}