#pragma once
#include <cstddef>
#include <initializer_list>
#include <mpi.h>
#include <vector>

struct SubarrayDimensionDefinition { // helper container
	std::size_t size{};
	std::size_t subSize{};
	std::size_t start{};
};

class SubarrayDefinition { // helper container for MPI Datatype creation
	std::vector<int> _sizes;
	std::vector<int> _subSizes;
	std::vector<int> _starts;

  public:
	int dims() const { return _sizes.size(); }
	int* sizes() { return _sizes.data(); }
	int* subSizes() { return _subSizes.data(); }
	int* starts() { return _starts.data(); }

	SubarrayDefinition(std::initializer_list<SubarrayDimensionDefinition>
	                       saDimDefs) {
		for (auto& dd : saDimDefs) {
			_sizes.push_back(static_cast<int>(dd.size));
			_subSizes.push_back(static_cast<int>(dd.subSize));
			_starts.push_back(static_cast<int>(dd.start));
		}
	}
};

class MpiSubarray { // wrapper for creating and destroying the type
	MPI_Datatype _type{MPI_DATATYPE_NULL};

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

	MpiSubarray(SubarrayDefinition sd) {
		MPI_Type_create_subarray(sd.dims(),     // ndims
		                         sd.sizes(),    // array_of_sizes
		                         sd.subSizes(), // array_of_subsizes
		                         sd.starts(),   // array_of_starts
		                         MPI_ORDER_C,   // order
		                         MPI_CHAR,      // oldtype
		                         &_type         // newtype
		                         );
		MPI_Type_commit(&_type);
	}
	~MpiSubarray() {
		if (_type != MPI_DATATYPE_NULL) { MPI_Type_free(&_type); }
	}
	MpiSubarray(MpiSubarray&) = delete;
	MpiSubarray& operator=(MpiSubarray&) = delete;
	MpiSubarray(MpiSubarray&& other) noexcept { swap(*this, other); }
	MpiSubarray& operator=(MpiSubarray&& other) noexcept {
		swap(*this, other);
		return *this;
	}
};