#include "simulation.h"
#include "SDL3/SDL_timer.h"
#include "input.h"
#include "thread_pool.h"
#include <SDL3/SDL_log.h>
#include <SDL3/SDL_render.h>
#include <algorithm>
#include <clay/clay.h>
#include <cstdint>
#include <mutex>
#include <set>
#include <vector>

#ifdef __glibc_likely
#define likely(cond_) __glibc_likely(cond_)
#else
#define likely(cond_) (cond_)
#endif

#define SIM_MULTITHREADING 1

namespace simulation {
bool drawDebugInfo{ true };

struct ThreadWorkload {
	size_t seg_idx{0}, seg_len{1};
	std::mutex mtx{};
	std::vector<Cell> changes{};
};

static std::set<Cell> living{};

static std::vector<std::shared_ptr<ThreadWorkload>> overpopulated{};
static std::vector<std::shared_ptr<ThreadWorkload>> underpopulated{};
static std::vector<std::shared_ptr<ThreadWorkload>> born{};
static uint64_t generationStartTime{ 0 };
static unsigned tasks{ 0 }; std::mutex tasksMutex{};
static std::condition_variable tasksChanged{};

CellIterator::CellIterator(Cell begin, Cell end)
: state{ begin } , begin{ begin }, end{ end } {}

CellIterator::CellIterator(Cell begin, Cell end, Cell state)
: state{ state } , begin{ begin }, end{ end } {}

static inline bool CellIsAlive(Cell const &cell) {
	return living.contains(cell);
}

CellIterator &CellIterator::operator++() {
	++(this->state.x);
	if (this->state.x == this->end.x) {
		this->state.x = this->begin.x;
		this->state.y = SDL_min(this->state.y + 1, this->end.y);
	}
	return *this;
}

bool CellIterator::operator==(CellIterator const &src) const {
	return src.begin == this->begin && src.end == this->end && (src.state == this->state || src.at_end() == this->at_end());
}

bool CellIterator::operator!=(CellIterator const &src) const {
	return src.begin != this->begin || src.end != this->end || (src.state != this->state && src.at_end() != this->at_end());
}

Cell const &CellIterator::operator*() const {
	return state;
}

bool CellIterator::at_end() const {
	return this->state.y == this->end.y;
}

static size_t CountNeighbors(Cell const &cell) {
	size_t count{ 0 };
	for (Cell const &c : CellRange{ {cell.x - 1, cell.y - 1}, { cell.x + 2, cell.y + 2} }) {
		if (c != cell && CellIsAlive(c)) {
			++count;
		}
	}
	return count;
}

static void BlockUntilTasksDone() {
	std::unique_lock lock{ tasksMutex };
	while (tasks > 0) {
		tasksChanged.wait(lock);
	}
}

static void TaskBegin() {
	std::scoped_lock lock{ tasksMutex };
	tasks++;
}

static void TaskComplete() {
	std::scoped_lock lock{ tasksMutex };
	if (--tasks == 0) {
		SDL_Log("Generation Complete %lfs", (double)(SDL_GetTicksNS() - generationStartTime) * 0.000000001);
	}
	tasksChanged.notify_all();
}

static void FindOverpopulated(std::shared_ptr<ThreadWorkload> wl) {
	TaskBegin();
	std::scoped_lock lock{ wl->mtx };
	wl->changes.clear();

	std::set<Cell>::iterator begin{ living.begin() };
	std::set<Cell>::iterator end{ living.begin() };

	std::advance(begin, wl->seg_idx * wl->seg_len);
	std::advance(end, wl->seg_idx * wl->seg_len + wl->seg_len);

	std::copy_if(begin, end, std::back_inserter(wl->changes),
	[&](Cell const &c) -> bool {
		return CountNeighbors(c) > 3;
	});
	TaskComplete();
}

static void FindUnderpopulated(std::shared_ptr<ThreadWorkload> wl) {
	TaskBegin();
	std::scoped_lock lock{ wl->mtx };
	wl->changes.clear();

	std::set<Cell>::iterator begin{ living.begin() };
	std::set<Cell>::iterator end{ living.begin() };

	std::advance(begin, wl->seg_idx * wl->seg_len);
	std::advance(end, wl->seg_idx * wl->seg_len + wl->seg_len);

	std::copy_if(begin, end, std::back_inserter(wl->changes),
	[&](Cell const &c) -> bool {
		return CountNeighbors(c) < 2;
	});
	TaskComplete();
}

static void FindBorn(std::shared_ptr<ThreadWorkload> wl) {
	TaskBegin();
	std::scoped_lock lock{ wl->mtx };
	wl->changes.clear();

	std::set<Cell>::iterator itr{ living.begin() };

	std::advance(itr, wl->seg_idx * wl->seg_len);

	std::for_each_n(itr, wl->seg_len, [&wl](Cell const &center){
		CellRange range{{ center.x - 1, center.y - 1 }, {center.x + 2, center.y + 2 }};
		std::copy_if(range.begin(), range.end(), std::back_inserter(wl->changes),
		[&](Cell const &c) -> bool {
			return !CellIsAlive(c) && CountNeighbors(c) == 3;
		});
	});
	TaskComplete();
}

static void PopulateChanges() {
	static bool first_run{ true };
#if SIM_MULTITHREADING
	BlockUntilTasksDone();
	constexpr size_t split{ 4 };
	if (first_run) {
		first_run = false;
		born.resize(split);
		underpopulated.resize(split);
		overpopulated.resize(split);
	}

	SDL_Log("Multithreading ON");
	generationStartTime = SDL_GetTicksNS();
	size_t const seg_length{ living.size() / split };
	for (size_t i{ 0 }; i < split; ++i) {
		if (overpopulated[i] == nullptr)
			overpopulated[i] = std::make_shared<ThreadWorkload>();
		{ std::scoped_lock lock{ overpopulated[i]->mtx };
			overpopulated[i]->seg_idx = i; overpopulated[i]->seg_len = seg_length;
			threading::tasks.ScheduleTask(std::bind(FindOverpopulated, overpopulated[i]));
		}
		if (underpopulated[i] == nullptr)
			underpopulated[i] = std::make_shared<ThreadWorkload>();
		{ std::scoped_lock lock{ underpopulated[i]->mtx };
			underpopulated[i]->seg_idx = i; underpopulated[i]->seg_len = seg_length;
			threading::tasks.ScheduleTask(std::bind(FindUnderpopulated, underpopulated[i]));
		}
		if (born[i] == nullptr)
			born[i] = std::make_shared<ThreadWorkload>();
		{ std::scoped_lock lock{ born[i]->mtx };
			born[i]->seg_idx = i; born[i]->seg_len = seg_length;
			threading::tasks.ScheduleTask(std::bind(FindBorn, born[i]));
		}
	}
#else
	if (first_run) {
		first_run = false;
		overpopulated.emplace_back(std::make_shared<ThreadWorkload>());
		overpopulated[0]->seg_idx = 0;
		overpopulated[0]->seg_len = living.size();
		underpopulated.emplace_back(std::make_shared<ThreadWorkload>());
		underpopulated[0]->seg_idx = 0;
		underpopulated[0]->seg_len = living.size();
		born.emplace_back(std::make_shared<ThreadWorkload>());
		born[0]->seg_idx = 0;
		born[0]->seg_len = living.size();
	}
	SDL_Log("Multithreading OFF");
	{ std::scoped_lock lock{ tasksMutex };
		tasks = 3; }
	generationStartTime = SDL_GetTicksNS();
	FindOverpopulated(overpopulated[0]);
	FindUnderpopulated(underpopulated[0]);
	FindBorn(born[0]);
#endif
}

void InitializeRandom(size_t livingChance, int64_t fillArea) {
	BlockUntilTasksDone();
	living.clear();

	Cell itr{ 0, 0 };
	while (itr.y < fillArea) {
		if (std::rand() % livingChance == 0	) {
			living.insert(itr);
		}
		++itr.x;
		if (itr.x == fillArea) {
			itr.x = 0;
			++itr.y;
		}
	}
	PopulateChanges();
}

void Step() {
	BlockUntilTasksDone();
	for (std::shared_ptr<ThreadWorkload> wl : underpopulated) {
		if (wl == nullptr) continue;
		std::scoped_lock l{ wl->mtx };
		for (Cell const &cell : wl->changes)
			living.erase(cell);
	}
	for (std::shared_ptr<ThreadWorkload> wl : overpopulated) {
		if (wl == nullptr) continue;
		std::scoped_lock l{ wl->mtx };
		for (Cell const &cell : wl->changes)
			living.erase(cell);
	}
	for (std::shared_ptr<ThreadWorkload> wl : born) {
		if (wl == nullptr) continue;
		std::scoped_lock l{ wl->mtx };
		for (Cell const &cell : wl->changes)
			living.insert(cell);
	}

	PopulateChanges();
}

static bool simulationHovered{ false };
static SDL_FPoint viewOffset{ 0, 0 };

void Draw(SDL_Renderer *renderer, double cellSizePercent) {
	if (simulationHovered) {
		viewOffset.x += input::scrollMotion.x;
		viewOffset.y += input::scrollMotion.y;
	}
	int w;
	SDL_GetCurrentRenderOutputSize(renderer, &w, nullptr);
	float const cellWidth = static_cast<float>(w) * cellSizePercent;
	SDL_FRect cellRect{
		0, 0, cellWidth, cellWidth
	};
	SDL_SetRenderDrawColor(renderer, 255, 255, 255, 255);
	for (Cell const &cell : living) {
		cellRect.x = (viewOffset.x + cell.x) * cellRect.w;
		cellRect.y = (viewOffset.y + cell.y) * cellRect.h;
		SDL_RenderFillRect(renderer, &cellRect);
	}
	if (!drawDebugInfo) {
		return;
	}
	BlockUntilTasksDone();
	for (std::shared_ptr<ThreadWorkload> wl : overpopulated) {
		std::scoped_lock l{ wl->mtx };
		SDL_SetRenderDrawColor(renderer, 255, 0, 0, 255);
		for (Cell const &cell : wl->changes) {
			cellRect.x = (viewOffset.x + cell.x) * cellRect.w;
			cellRect.y = (viewOffset.y + cell.y) * cellRect.h;
			SDL_RenderRect(renderer, &cellRect);
		}
	}
	for (std::shared_ptr<ThreadWorkload> wl : underpopulated) {
		std::scoped_lock l{ wl->mtx };
		SDL_SetRenderDrawColor(renderer, 255, 0, 255, 255);
		for (Cell const &cell : wl->changes) {
			cellRect.x = (viewOffset.x + cell.x) * cellRect.w;
			cellRect.y = (viewOffset.y + cell.y) * cellRect.h;
			SDL_RenderRect(renderer, &cellRect);
		}
	}
	for (std::shared_ptr<ThreadWorkload> wl : born) {
		std::scoped_lock l{ wl->mtx };
		SDL_SetRenderDrawColor(renderer, 0, 255, 0, 255);
		for (Cell const &cell : wl->changes) {
			cellRect.x = (viewOffset.x + cell.x) * cellRect.w;
			cellRect.y = (viewOffset.y + cell.y) * cellRect.h;
			SDL_RenderRect(renderer, &cellRect);
		}
	}
}

void SetSimulationHovered(bool value) {
	simulationHovered = value;
}
bool operator==(Cell const &lhs, Cell const &rhs) {
	return lhs.x == rhs.x && lhs.y == rhs.y;
}
bool operator!=(Cell const &lhs, Cell const &rhs) {
	return lhs.x != rhs.x || lhs.y != rhs.y;
}

bool operator<(Cell const &lhs, Cell const &rhs) {
	if (lhs.y == rhs.y) return lhs.x < rhs.x;
	else return lhs.y < rhs.y;
}
}