#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 "Communicator.hpp"
#include "FileIO.hpp"
#include "MpiEnvironment.hpp"
#include "Util.hpp"
#include "state.hpp"

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;
	Communicator _comm;

  public:
	auto getTileSize() { return _tileSize; }

	MpiWireworld(const MpiEnvironment& env, const Configuration& cfg)
	    : _env(env), _cfg(cfg) {
		const auto& gridSize = Size{cfg.GridColumns, cfg.GridRows};
				  
		const auto& header = FileIO::ReadHeader(cfg.InputFilePath);
		const auto& headerLength = header.HeaderLength;
		const auto& _globalSize = header.GlobalSize;
		_tileSize = FileIO::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));

		FileIO::ReadTile(cfg.InputFilePath, headerLength, _globalSize, gridSize,
		                 _env.worldRank(),
		                 gsl::span<State>(_memoryA));

		_comm = Communicator{env, gridSize, _tileSize};
	}

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

	void simulateStep() {
		_comm.Communicate(_model); // compute the border area first, then comm
		                           // and compute the rest async
		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);
	}
};