#ifndef FINAL_TESTS_SHARED_GRID_HPP
#define FINAL_TESTS_SHARED_GRID_HPP

#include <vector>
#include <string>
#include <fstream>

#include "exception.hpp"

/**
 * Universal multi-dimensional grid of values of uniform type.
 * Data are stored in row-major order (first dimension is the compact one).
 */
template<typename T, int DIM>
class Grid
{
private:
	std::vector<T> mData;
	std::size_t mSize[DIM];

	template <int D, typename... Args>
	std::size_t flat_index(std::size_t index, Args... remainingIndices) const {
		static_assert(D < DIM, "Dimension out of bounds");

		if (index >= mSize[D]) {
			throw std::out_of_range("Index out of bounds");
		}

		if constexpr (D == DIM - 1) {
			return index;
		} else {
			return flat_index<D + 1>(remainingIndices...) * mSize[D] + index;
		}
	}

	template<int D, typename... Args>
	void save_size(std::size_t size, Args... remainingSizes)
	{
		static_assert(D >= 0 && D < DIM, "Dimension out of bounds");
		mSize[D] = size;
		if constexpr (D < DIM - 1) {
			save_size<D + 1>(remainingSizes...);
		}
	}

public:
	Grid()
	{
		static_assert(DIM > 0, "Grid must have at least one dimension");
	}

	/**
	 * Initialize emtpy grid with given sizes of dimensions.
	 * The number of arguments must match the number of dimensions.
	 */
	template<typename... Args> requires (sizeof...(Args) == DIM) && (std::convertible_to<Args, std::size_t> && ...)
	Grid(Args... sizes)
	{
		save_size<0>(sizes...);
		std::size_t totalSize = (static_cast<std::size_t>(sizes) * ...);
		mData.resize(totalSize);
	}

	/**
	 * Copy the entire contents of another grid of the same dimensions and type
	 * into this grid. The sizes of dimensions are also copied.
	 */
	void copyFrom(const Grid<T, DIM>& other)
	{
		for (std::size_t i = 0; i < DIM; ++i) {
			mSize[i] = other.mSize[i];
		}
		mData = other.mData;
	}

	/**
	 * Swap the entire contents of this grid with another grid of the same dimensions and type.
	 * Useful for double buffering.
	 */
	void swap(Grid<T, DIM>& other)
	{
		for (std::size_t i = 0; i < DIM; ++i) {
			std::swap(mSize[i], other.mSize[i]);
		}
		mData.swap(other.mData);
	}

	std::size_t dimensions() const
	{
		return DIM;
	}

	/**
	 * Get grid size in given dimension.
	 */
	template<int D>
	std::size_t size() const
	{
		static_assert(D >= 0 && D < DIM, "Dimension out of range.");
		return mSize[D];
	}

	/**
	 * Get reference to the grid element at given multi-dimensional index.
	 * The number of arguments must match the number of dimensions.
	 */
	template <typename... Args> requires (sizeof...(Args) == DIM) && (std::convertible_to<Args, std::size_t> && ...)
	T& operator()(Args... indices) {
		std::size_t offset = flat_index<0>(static_cast<size_t>(indices)...);
		return mData[offset];
	}

	template <typename... Args> requires (sizeof...(Args) == DIM) && (std::convertible_to<Args, std::size_t> && ...)
	const T& operator()(Args... indices) const {
		std::size_t offset = flat_index<0>(static_cast<size_t>(indices)...);
		return mData[offset];
	}

	/**
	 * Get raw pointer to the data stored in row-major format.
	 */
	T* data()
	{
		return mData.data();
	}

	const T* data() const
	{
		return mData.data();
	}

	/**
	 * Fill the grid with given value.
	 */
	void fill(T value)
	{
		std::fill(mData.begin(), mData.end(), value);
	}

	template<class GEN, class DIST>
	void generate(GEN &generator, DIST distribution)
	{
		for (auto&& x : mData) {
			x = (T)distribution(generator);
		}
	}

	/**
	 * Load grid data from binary file.
	 */
	void load(std::string filePath)
	{
		std::ifstream file(filePath, std::ios::binary);
		if (!file.is_open())
		{
			throw std::runtime_error("Could not open file: " + filePath);
		}

		// prepare reader lambdas
		static_assert(sizeof(std::size_t) == 8, "Assuming size_t is a 64bit number.");
		auto readNum = [&file]() -> std::size_t
		{
			std::size_t value;
			file.read(reinterpret_cast<char*>(&value), sizeof(std::size_t));
			return value;
		};
		auto readVal = [&file]() -> T
			{
				T value;
				file.read(reinterpret_cast<char*>(&value), sizeof(T));
				return value;
			};

		// get header (dimensions)
		std::size_t dims = readNum();
		if (dims != DIM) {
			throw std::runtime_error("Dimension mismatch: expected " + std::to_string(DIM) + ", got " + std::to_string(dims));
		}

		std::size_t totalSize = 1;
		for (std::size_t i = 0; i < DIM; ++i) {
			mSize[i] = readNum();
			totalSize *= mSize[i];
		}

		// get data
		mData.resize(totalSize);
		for (auto&& x : mData) {
			x = readVal();
		}
	}

	/**
	 * Save grid data to binary file.
	 */
	void save(std::string filePath)
	{
		std::ofstream file(filePath, std::ios::binary);
		if (!file.is_open())
		{
			throw std::runtime_error("Could not open file: " + filePath);
		}

		static_assert(sizeof(std::size_t) == 8, "Assuming size_t is a 64bit number.");
		auto writeNum = [&](std::size_t value) {
			file.write(reinterpret_cast<const char*>(&value), static_cast<std::streamsize>(sizeof(std::size_t)));
			if (!file) {
				throw (bpp::RuntimeError() << "Failed while writing to file: " << filePath);
			}
		};
		auto writeVal = [&](T value) {
			file.write(reinterpret_cast<const char*>(&value), static_cast<std::streamsize>(sizeof(T)));
			if (!file) {
				throw (bpp::RuntimeError() << "Failed while writing to file: " << filePath);
			}
		};

		// write header (dimensions)
		writeNum(DIM);
		for (std::size_t i = 0; i < DIM; ++i) {
			writeNum(mSize[i]);
		}

		// write data
		for (const auto& x : mData) {
			writeVal(x);
		}

		file.flush();
		if (!file) {
			throw (bpp::RuntimeError() << "Failed to flush file: " << filePath);
		}
	}
};

#endif