#pragma once

#include <algorithm>
#include <array>
#include <cstddef>
#include <fstream>
#include <iostream>
#include <mpi.h>
#include <string>
#include <tuple>
#include <vector>

#include "gsl/multi_span"
#include "gsl/span"

#include "MpiEnvironment.hpp"
#include "state.hpp"

struct Size {
	std::size_t Cols{};
	std::size_t Rows{};
};

struct Coord {
	int X{};
	int Y{};
};

class MpiWireworld {
	const MpiEnvironment& _env;
	const Configuration& _cfg;
	Size _tileSize;
	std::vector<State> _memoryA;
	std::vector<State> _memoryB;
	gsl::multi_span<State, -1, -1> _model;
	gsl::multi_span<State, -1, -1> _nextModel;

	std::vector<int> _myNeighbors;
	std::vector<int> _mySizes;
	std::vector<MPI_Datatype> _mySendTypes;
	const std::vector<MPI_Datatype>& _myRecvTypes{_mySendTypes};
	std::vector<MPI_Aint> _mySendDisplacements;
	std::vector<MPI_Aint> _myRecvDisplacements;
	MPI_Comm _commDistGraph{MPI_COMM_NULL};

	// begin IO functions
	static auto ReadHeader(std::istream& input) {
		std::size_t noCols{};
		std::size_t noRows{};
		input >> noCols >> noRows;
		if (noCols < 1 || noRows < 1)
			throw std::logic_error("File header corrupt");
		std::string dummy; // skip line break
		std::getline(input, dummy);
		const auto& headerSize = input.tellg();
		return Size{noCols, noRows};
	}

	static auto GetTileSize(Size globalSize, Size gridSize) {
		const auto& tileSizeCols = globalSize.Cols / gridSize.Cols;
		const auto& tileSizeRows = globalSize.Rows / gridSize.Rows;
		return Size{tileSizeCols, tileSizeRows};
	}

	static void StridedRead(std::istream& input, std::size_t inputChunkLength,
	                        std::size_t count, std::size_t inputStride,
	                        gsl::span<State> buf, std::size_t bufStride) {

		std::size_t writeOffset{0};
		for (size_t i{0}; i < count; ++i) {
			input.read(reinterpret_cast<char*>(buf.data()) + writeOffset,
			           inputChunkLength);
			input.seekg(inputStride - inputChunkLength, std::ios::cur);
			writeOffset += bufStride;
		}
	}

	static void ReadTiles(std::istream& input, Size srcSize, Size gridSize,
	                      std::size_t rank, gsl::span<State> buf,
	                      std::size_t bufStride) {
		constexpr auto LF = 1; // linefeed chars

		const auto& tileSize = GetTileSize(srcSize, gridSize);

		const auto& tileX = rank % gridSize.Cols;
		const auto& tileY = rank / gridSize.Cols;

		// seek to the begin of the tile
		const auto& displacement =
		    (srcSize.Cols + LF) * (tileSize.Rows * tileY) +
		    (tileSize.Cols * tileX);
		input.seekg(displacement, std::ios::cur);

		const auto& tileStride = srcSize.Cols + LF;

		StridedRead(input, tileSize.Cols, tileSize.Rows, tileStride, buf,
		            bufStride);
	}

	void ReadFile(const std::string& path, Size gridSize) {
		std::ifstream input(path);

		const auto& globalSize = ReadHeader(input);
		_tileSize = GetTileSize(globalSize, gridSize);

		_memoryA.resize((_tileSize.Cols + 2) * (_tileSize.Rows + 2));
		_memoryB.resize((_tileSize.Cols + 2) * (_tileSize.Rows + 2));
		_model =
		    gsl::as_multi_span(_memoryA.data(), gsl::dim(_tileSize.Rows + 2),
		                       gsl::dim(_tileSize.Cols + 2));
		_nextModel =
		    gsl::as_multi_span(_memoryB.data(), gsl::dim(_tileSize.Rows + 2),
		                       gsl::dim(_tileSize.Cols + 2));

		ReadTiles(input, globalSize, gridSize, _env.worldRank(),
		          gsl::span<State>(_memoryA.data() + _tileSize.Cols + 2 + 1,
		                           _memoryA.size() - 1),
		          _tileSize.Cols + 2);

		/*

		// Trying to acomplish this with MPI


		// if (env.isMaster()) {
		//	std::cout << "noCols: " << _noCols << '\n';
		//	std::cout << "noRows: " << _noRows << '\n';
		//}

		// const auto& noDims = 2;

		// const auto& gsizes = std::array<int, 2>{
		//	_noCols + 1,  //
		//    _noRows, //
		//};

		// const auto& distribs = std::array<int, 2>{
		//    MPI_DISTRIBUTE_BLOCK, //
		//    MPI_DISTRIBUTE_BLOCK  //
		//};

		// const auto& dargs = std::array<int, 2>{
		//    MPI_DISTRIBUTE_DFLT_DARG, //
		//    MPI_DISTRIBUTE_DFLT_DARG  //
		//};

		// const auto& psizes = std::array<int, 2>{
		//    2, //
		//    4  //
		//};

		//// const auto& noDims = 1;

		//// const auto& gsizes = std::array<int, 1>{_noCols};

		//// const auto& distribs = std::array<int, 1>{MPI_DISTRIBUTE_BLOCK};

		//// const auto& dargs = std::array<int, 1>{MPI_DISTRIBUTE_DFLT_DARG};

		//// const auto& psizes = std::array<int, 1>{8};

		////MPI_Datatatype tiletype;
		////int global_sizes[2] = { _noCols + 1, _noRows };
		////int tile_sizes[2] = { 10, 5 };
		////int tile_start[2] = { }

		////MPI_Type_create_subarray(2, global_sizes, tile_sizes, )

		// MPI_Datatype tiletype;
		////MPI_Datatype linetype;
		////MPI_Datatype linetype_ext;

		////MPI_Type_contiguous(_noCols, MPI_CHAR, &linetype);
		////MPI_Type_create_resized(linetype, 0, _noCols + 1, &linetype_ext);
		////MPI_Type_commit(&linetype_ext);
		////MPI_Type_free(&linetype);

		//
		// MPI_Type_create_darray(env.worldSize(), env.worldRank(), noDims,
		//                       gsizes.data(), distribs.data(), dargs.data(),
		//                       psizes.data(), MPI_ORDER_C, MPI_CHAR,
		//                       &tiletype);
		// MPI_Type_commit(&tiletype);
		// int tiletypeSize;
		// MPI_Type_size(tiletype, &tiletypeSize);

		// if (env.isMaster())
		//	std::cout << "tiletypeSize: " << tiletypeSize << '\n';

		// MPI_File fh;
		// MPI_File_open(MPI_COMM_WORLD, path.c_str(), MPI_MODE_RDONLY,
		//              MPI_INFO_NULL, &fh);
		// MPI_File_set_view(fh, _headerSize, MPI_CHAR, MPI_CHAR, "native",
		//                  MPI_INFO_NULL);

		// const auto& typeCols = 10; // how to get those?
		// const auto& typeRows = 5;



		MPI_File_read_all(fh, buf.data(), 1, tiletype, MPI_STATUS_IGNORE);

		MPI_File_close(&fh);

		std::replace(std::begin(buf), std::end(buf), '\n', 'L');

		*/
	}
	// end IO functions

  public:
	MpiWireworld(const MpiEnvironment& env, const Configuration& cfg)
	    : _env(env), _cfg(cfg) {
		const auto& gridSize = Size{cfg.GridColumns, cfg.GridRows};
		ReadFile(_cfg.InputFilePath, gridSize);

		// types begin
		MPI_Datatype haloRowType;
		MPI_Type_contiguous(_tileSize.Cols, MPI_CHAR, &haloRowType);
		MPI_Type_commit(&haloRowType);

		MPI_Datatype haloColumnType;
		MPI_Type_vector(_tileSize.Rows, 1, _tileSize.Cols + 2, MPI_CHAR,
		                &haloColumnType);
		MPI_Type_commit(&haloColumnType);

		MPI_Datatype haloCornerType = MPI_CHAR;
		// types end

		const auto& generalSendTypes = std::array<MPI_Datatype, 8>{
		    haloCornerType, haloRowType,    haloCornerType, //
		    haloColumnType, haloColumnType,                 //
		    haloCornerType, haloRowType,    haloCornerType  //
		};
		const auto& generalRecvTypes = generalSendTypes; // same

		const auto& cols = _tileSize.Cols;
		const auto& rows = _tileSize.Rows;

		auto idx = [&](std::size_t x, std::size_t y) {
			return static_cast<MPI_Aint>(y * (cols + 2) + x);
		};

		std::array<MPI_Aint, 8> generalSendDisplacements{
		    idx(1, 1),    idx(1, 1),    idx(cols, 1),   //
		    idx(1, 1),    idx(cols, 1),                 //
		    idx(1, rows), idx(1, rows), idx(cols, rows) //
		};

		std::array<MPI_Aint, 8> generalRecvDisplacements{
		    idx(0, 0),        idx(1, 0),        idx(cols + 1, 0),       //
		    idx(0, 1),        idx(cols + 1, 1),                         //
		    idx(0, rows + 1), idx(1, rows + 1), idx(cols + 1, rows + 1) //
		};

		std::array<int, 8> generalSizes{
		    // TODO: real better name
		    1, 1, 1, //
		    1, 1,    //
		    1, 1, 1  //
		};

		auto rank2coord = [&](int rank) {
			return Coord{
			    static_cast<int>(rank % gridSize.Cols), //
			    static_cast<int>(rank / gridSize.Cols)  //
			};
		};

		constexpr auto DebugRank = 0;

		auto coord2rank = [&](Coord c) {
			// if (env.worldRank() == DebugRank) std::cout << "c2r input X" <<
			// c.X << " Y" << c.Y << '\n';
			const auto& x = (c.X + gridSize.Cols) % gridSize.Cols;
			const auto& y = (c.Y + gridSize.Rows) % gridSize.Rows;
			// if (env.worldRank() == DebugRank) std::cout << "c2r after mod X"
			// << x << " Y" << y << '\n';
			return static_cast<int>(gridSize.Cols * y + x);
		};

		const auto& myCo = rank2coord(env.worldRank());
		std::array<Coord, 8> virtualNeighborCoords{{
		    {myCo.X - 1, myCo.Y - 1}, //
		    {myCo.X + 0, myCo.Y - 1}, //
		    {myCo.X + 1, myCo.Y - 1}, //
		    {myCo.X - 1, myCo.Y + 0}, //
		    {myCo.X + 1, myCo.Y + 0}, //
		    {myCo.X - 1, myCo.Y + 1}, //
		    {myCo.X + 0, myCo.Y + 1}, //
		    {myCo.X + 1, myCo.Y + 1}  //
		}};

		auto isInsideGrid = [&](Coord c) {
			return 0 <= c.X && c.X < gridSize.Cols && 0 <= c.Y &&
			       c.Y < gridSize.Rows;
		};

		// interate over all neighbors
		for (std::size_t i{0}; i < 8; ++i) {
			const auto& c = virtualNeighborCoords[i];
			if (isInsideGrid(c)) { // neighbor
				_myNeighbors.push_back(coord2rank(c));
				_mySendTypes.push_back(generalSendTypes[i]);
				_mySendDisplacements.push_back(generalSendDisplacements[i]);
				_myRecvDisplacements.push_back(generalRecvDisplacements[i]);
				_mySizes.push_back(generalSizes[i]);
			}
		}

		// if (env.worldRank() == DebugRank) {
		//	std::cout << "neighbors:\n";
		//	for (const auto i : neighbors) std::cout << i << ",\n";
		//	std::cout << '\n';
		//}

		MPI_Dist_graph_create_adjacent(
		    MPI_COMM_WORLD,                         // comm_old
		    _myNeighbors.size(),                    // indegree
		    _myNeighbors.data(),                    // sources
		    reinterpret_cast<int*>(MPI_UNWEIGHTED), // sourceweights
		    _myNeighbors.size(),                    // outdegree
		    _myNeighbors.data(),                    // destinations
		    reinterpret_cast<int*>(MPI_UNWEIGHTED), // destweights
		    MPI_INFO_NULL,                          // info
		    0,                                      // reorder
		    &_commDistGraph                         // comm_dist_graph
		    );

		MPI_Neighbor_alltoallw(_model.data(),               // sendbuf
		                       _mySizes.data(),             // sendcounts
		                       _mySendDisplacements.data(), // sdispl
		                       _mySendTypes.data(),         // sendtypes
		                       _model.data(),               // recvbuf
		                       _mySizes.data(),             // recvcounts
		                       _myRecvDisplacements.data(), // rdispls
		                       _myRecvTypes.data(),         // recvtypes
		                       _commDistGraph               // comm
		                       );
	}

	friend std::ostream& operator<<(std::ostream& out, const MpiWireworld& g) {
		for (std::size_t x{0}; x < g._tileSize.Rows + 2; ++x) {
			out << '"';
			for (std::size_t y{0}; y < g._tileSize.Cols + 2; ++y) {
				out << to_integral(g._model[x][y]);
			}
			out << "\"\n";
		}
		return out;
	}

	void simulateStep() {
		for (std::size_t x{1}; x <= _tileSize.Rows; ++x) {
			for (std::size_t y{1}; y <= _tileSize.Cols; ++y) {
				auto nextState = _model[x][y];
				switch (_model[x][y]) {
				case State::ElectronHead:
					nextState = State::ElectronTail;
					break;
				case State::ElectronTail:
					nextState = State::Conductor;
					break;
				case State::Conductor: {
					const std::array<State, 9> mooreNeighborhood = {
					    _model[x - 1][y - 1], //
					    _model[x + 0][y - 1], //
					    _model[x + 1][y - 1], //
					    _model[x - 1][y + 0], //
					    _model[x + 1][y + 0], //
					    _model[x - 1][y + 1], //
					    _model[x + 0][y + 1], //
					    _model[x + 1][y + 1]  //
					};

					const auto& headCount = std::count(
					    std::begin(mooreNeighborhood),
					    std::end(mooreNeighborhood), State::ElectronHead);

					nextState = (1 == headCount || headCount == 2)
					                ? State::ElectronHead
					                : State::Conductor;

				} break;
				default:
					break;
				}
				_nextModel[x][y] = nextState;
			}
		}
		std::swap(_model, _nextModel);
	}
};