#include <mpi.h>

class MpiComm {
	MPI_Comm _comm{MPI_COMM_NULL};
	std::size_t _size{0};
	std::size_t _rank{0};
	bool _isMaster{false};

  public:
	MPI_Comm comm() const { return _comm; }
	std::size_t size() const { return _size; }
	std::size_t rank() const { return _rank; }
	bool isMaster() const { return _isMaster; }

	friend void swap(MpiComm& first, MpiComm& second) noexcept {
		using std::swap;
		swap(first._size, second._size);
		swap(first._rank, second._rank);
		swap(first._isMaster, second._isMaster);
	}

  private: // use named ctors idiom because the ctors params are hard to understand
	MpiComm() {
		MPI_Comm_split_type(MPI_COMM_WORLD, MPI_COMM_TYPE_SHARED, 0, MPI_INFO_NULL, &_comm);
		_size = [&] {int s;  MPI_Comm_size(_comm, &s); return static_cast<std::size_t>(s); }();
		_rank = [&] { int r;  MPI_Comm_rank(_comm, &r); return static_cast<std::size_t>(r); }();
		_isMaster = _rank == 0;
	}

	MpiComm(int color) {
		MPI_Comm_split(MPI_COMM_WORLD, color, 0, &_comm);
		if (_comm != MPI_COMM_NULL) {
			_size = [&] {int s;  MPI_Comm_size(_comm, &s); return static_cast<std::size_t>(s); }();
			_rank = [&] { int r;  MPI_Comm_rank(_comm, &r); return static_cast<std::size_t>(r); }();
		}
		_isMaster = _rank == 0;
	}

  public:
	static MpiComm fromCommSplitTypeShared() { return {}; }
	static MpiComm fromCommSplitColor(int color) { return {color}; }

	~MpiComm() {
		if (_comm != MPI_COMM_NULL) { MPI_Comm_free(&_comm); };
	}

	MpiComm(const MpiComm&) = delete;
	MpiComm& operator=(const MpiComm&) = delete;
	MpiComm(MpiComm&& other) noexcept { swap(*this, other); }
	MpiComm& operator=(MpiComm&& other) noexcept {
		swap(*this, other);
		return *this;
	}

	void barrier() const { MPI_Barrier(_comm); }
};
