#pragma once

#include "Configuration.hpp"
#include "Particle.hpp"
#include <cstddef>
#include <initializer_list>
#include <mpi.h>

namespace {
struct DataMemberDefinition { // helper container
	int blocklength{0};
	MPI_Aint displacement{0};
	MPI_Datatype datatype{0};
};

class StructDefinition { // helper container for MPI Datatype creation
	std::vector<int> _blocklengths;
	std::vector<MPI_Datatype> _datatypes;
	std::vector<MPI_Aint> _displacements;

  public:
	int size() const { return static_cast<int>(_blocklengths.size()); }
	int* blocklengths() { return _blocklengths.data(); }
	MPI_Datatype* datatypes() { return _datatypes.data(); }
	MPI_Aint* displacements() { return _displacements.data(); }

	StructDefinition(std::initializer_list<DataMemberDefinition> dataMemberDefinitions) {
		for (auto& md : dataMemberDefinitions) {
			_blocklengths.push_back(md.blocklength);
			_datatypes.push_back(md.datatype);
			_displacements.push_back(md.displacement);
		}
	}
};

template <typename T> class MpiType { // wrapper for creating and destroying the type
	MPI_Datatype _type{MPI_DATATYPE_NULL};

  public:
	friend void swap(MpiType& first, MpiType& second) noexcept {
		using std::swap;
		swap(first._type, second._type);
	}
	MPI_Datatype type() const { return _type; }

	MpiType(std::initializer_list<DataMemberDefinition> dataMemberDefinitions) {
		StructDefinition sd{dataMemberDefinitions};
		MPI_Datatype tempType;
		MPI_Type_create_struct(sd.size(), sd.blocklengths(), sd.displacements(), sd.datatypes(), &tempType);
		MPI_Type_create_resized(tempType, 0, sizeof(T), &_type);
		MPI_Type_commit(&_type);
		MPI_Type_free(&tempType);
	}
	~MpiType() {
		if (_type != MPI_DATATYPE_NULL) { MPI_Type_free(&_type); }
	}
	MpiType(MpiType&) = delete;
	MpiType& operator=(MpiType&) = delete;
	MpiType(MpiType&& other) noexcept { swap(*this, other); }
	MpiType& operator=(MpiType&& other) noexcept {
		swap(*this, other);
		return *this;
	}
};
}

// =========== BEGIN TYPES ==============

MpiType<Vec3> GenVec3Type() {
	return {
	    {1, offsetof(Vec3, X), MPI_DOUBLE}, //
	    {1, offsetof(Vec3, Y), MPI_DOUBLE}, //
	    {1, offsetof(Vec3, Z), MPI_DOUBLE}  //
	};
}

MpiType<Particle> GenParticleType() {
	auto vec3Type = GenVec3Type();
	return {
	    {1, offsetof(Particle, Mass), MPI_DOUBLE},          //
	    {1, offsetof(Particle, Location), vec3Type.type()}, //
	    {1, offsetof(Particle, Velocity), vec3Type.type()}  //
	};
}

MpiType<Particle> GenSparseLocatedParticleType() {
	auto vec3Type = GenVec3Type();
	return {
	    {1, offsetof(Particle, Location), vec3Type.type()} //
	};
}

MpiType<Configuration> GenConfigurationType() {
	return {
	    {1, offsetof(Configuration, NoIterations), MPI_UINT64_T}, //
	    {1, offsetof(Configuration, NoParticles), MPI_UINT64_T}   //
	};
}
