#include <mpi.h>

template <typename T> class MpiSharedMemWin {
	MPI_Win _memoryWindow{MPI_WIN_NULL};
	std::size_t _localCount{0};
	T* _basePtr{nullptr};

  public:
	MPI_Win memoryWindow() const { return _memoryWindow; }
	T* basePtr() const { return _basePtr; }

	friend void swap(MpiSharedMemWin& first, MpiSharedMemWin& second) noexcept {
		using std::swap;
		swap(first._memoryWindow, second._memoryWindow);
		swap(first._localCount, second._localCount);
		swap(first._basePtr, second._basePtr);
	}

	MpiSharedMemWin(std::size_t localWindowCount, MPI_Comm comm) : _localCount(localWindowCount) {
		const auto dispUnit = sizeof(T);
		const auto bytes = dispUnit * _localCount;
		T* tmpLocalBasePtr;
		MPI_Win_allocate_shared(static_cast<MPI_Aint>(bytes), dispUnit, MPI_INFO_NULL, comm, &tmpLocalBasePtr, &_memoryWindow);

		// Get pointer to start of shared memory region
		int rank{0};
		MPI_Aint windowSize;
		int tmpDispUnit;
		MPI_Win_shared_query(_memoryWindow, rank, &windowSize, &tmpDispUnit, &_basePtr);
	}
	~MpiSharedMemWin() {
		if (_memoryWindow != MPI_WIN_NULL) { MPI_Win_free(&_memoryWindow); }
	}
	MpiSharedMemWin(MpiSharedMemWin&) = delete;
	MpiSharedMemWin& operator=(MpiSharedMemWin&) = delete;
	MpiSharedMemWin(MpiSharedMemWin&& other) noexcept { swap(*this, other); }
	MpiSharedMemWin& operator=(MpiSharedMemWin&& other) noexcept {
		swap(*this, other);
		return *this;
	}

	T& operator[](std::size_t idx) { return _basePtr[idx]; }
	const T& operator[](std::size_t idx) const { return _basePtr[idx]; }

	using iterator = T*;
	using const_iterator = const T*;
	iterator begin() { return _basePtr; }
	const_iterator begin() const { return _basePtr; }
	iterator end() { return _basePtr + _localCount; }
	const_iterator end() const { return _basePtr + _localCount; }
};
