#include "MpiWireworld.hpp"

#include <algorithm>
#include <array>
#include <iostream>
#include <limits>
#include <mpi.h>
#include <string>
#include <vector>

#include "CollectiveCommunicator.hpp"
#include "FileIO.hpp"
#include "P2PCommunicator.hpp"
#include "State.hpp"
#include "Tile.hpp"
#include "Util.hpp"

void MpiWireworld::processArea(Coord start, Size size) {
	auto& _model = _tile.model();
	auto& _nextModel = _tile.nextModel();
	const auto modelWidth = _tile.modelWidth();

	// std::size_t is unsigned. modulo arithmetics is used for calculating the
	// index
	const std::size_t leftOffset = std::numeric_limits<std::size_t>::max(); // -1;
	const std::size_t rightOffset = 1;
	const std::size_t downOffset = modelWidth;
	const std::size_t upOffset = -downOffset;

	for (std::size_t y{start.Y}; y < start.Y + size.Rows; ++y) {
		for (std::size_t x{start.X}; x < start.X + size.Cols; ++x) {
			const auto idx = y * modelWidth + x;

			const auto currentState = _model[idx];
			_nextModel[idx] = [&]() {
				switch (currentState) {
				case State::ElectronHead:
					return State::ElectronTail;
				case State::ElectronTail:
					return State::Conductor;
				case State::Conductor: {
					const auto isHead = [&](std::size_t i) {
						return _model[i] == State::ElectronHead ? 1 : 0;
					};
					const auto headCount =
					    isHead(idx + leftOffset + upOffset) +   //
					    isHead(idx + upOffset) +                //
					    isHead(idx + rightOffset + upOffset) +  //
					    isHead(idx + leftOffset) +              //
					    isHead(idx + rightOffset) +             //
					    isHead(idx + leftOffset + downOffset) + //
					    isHead(idx + downOffset) +              //
					    isHead(idx + rightOffset + downOffset); //

					return (1 == headCount || headCount == 2)
					           ? State::ElectronHead
					           : State::Conductor;
				};
				case State::Empty:
					return currentState;
				}
			}();
		}
	}
}

MpiWireworld::MpiWireworld(const MpiEnvironment& env, const Configuration& cfg)
    : _tile(Tile::Read(cfg, env)),
      _comm([&]() -> std::unique_ptr<Communicator> {
	      switch (cfg.CommMode) {
	      case CommunicationMode::Collective:
		      return std::make_unique<CollectiveCommunicator>(
		          env, cfg.CommMode, cfg.Procs, _tile.tileSize());
	      case CommunicationMode::P2P:
		      return std::make_unique<P2PCommunicator>(
		          env, cfg.CommMode, cfg.Procs, _tile.tileSize());
	      }
	  }()) {
	_comm->Communicate(_tile.model());
}

std::ostream& operator<<(std::ostream& out, const MpiWireworld& g) {
	// for now, put only our local tile
	out << g._tile;
	return out;
}

void MpiWireworld::write() const { _tile.write(); }

void MpiWireworld::simulateStep() {
	auto& _model = _tile.model();
	auto& _nextModel = _tile.nextModel();
	const auto _tileSize = _tile.tileSize();

	//// compute the border area first, then comm
	//// and compute the core async

	/// border area
	// top
	processArea({1, 1}, {_tileSize.Cols, 1});

	// left and right
	processArea({1, 2}, {1, _tileSize.Rows - 2});
	processArea({_tileSize.Cols, 2}, {1, _tileSize.Rows - 2});

	// bottom
	processArea({1, _tileSize.Rows}, {_tileSize.Cols, 1});

	// start communication of border while processing the core
	auto req = _comm->AsyncCommunicate(_nextModel);

	/// core
	processArea({2, 2}, {_tileSize.Cols - 2, _tileSize.Rows - 2});

	req.Wait();

	std::swap(_model, _nextModel);
}