#ifndef DOCTEST_MPI_SUB_COMM_H
#define DOCTEST_MPI_SUB_COMM_H

#include "mpi.h"
#include "doctest/doctest.h"
#include <cassert>
#include <string>

namespace doctest {

inline int mpi_world_nb_procs() {
    int n;
    MPI_Comm_size(MPI_COMM_WORLD, &n);
    return n;
}

struct mpi_sub_comm {
    int nb_procs;
    int rank;
    MPI_Comm comm;

    mpi_sub_comm(const mpi_sub_comm &) = delete;
    mpi_sub_comm &operator=(const mpi_sub_comm &) = delete;

    mpi_sub_comm(int nb_prcs) noexcept
        : nb_procs(nb_prcs), rank(-1), comm(MPI_COMM_NULL) {
        int comm_world_rank;
        MPI_Comm_rank(MPI_COMM_WORLD, &comm_world_rank);
        if (nb_procs > mpi_world_nb_procs()) {
            if (comm_world_rank == 0) {
                MESSAGE(
                    "Unable to run test: need ",
                    std::to_string(nb_procs),
                    " procs",
                    " but program launched with only ",
                    std::to_string(doctest::mpi_world_nb_procs()),
                    "."
                );
                CHECK(nb_procs <= mpi_world_nb_procs());
            }
        } else {
            int color = MPI_UNDEFINED;
            if (comm_world_rank < nb_procs) {
                color = 0;
            }
            MPI_Comm_split(MPI_COMM_WORLD, color, comm_world_rank, &comm);

            if (comm != MPI_COMM_NULL) {
                MPI_Comm_rank(comm, &rank);
                assert(rank == comm_world_rank);
            }
        }
    }

    void destroy_comm() {
        if (comm != MPI_COMM_NULL) {
            MPI_Comm_free(&comm);
        }
    }

    mpi_sub_comm(mpi_sub_comm &&x)
        : nb_procs(x.nb_procs), rank(x.rank), comm(x.comm) {
        x.comm = MPI_COMM_NULL;
    }
    mpi_sub_comm &operator=(mpi_sub_comm &&x) {
        destroy_comm();
        nb_procs = x.nb_procs;
        rank = x.rank;
        comm = x.comm;
        x.comm = MPI_COMM_NULL;
        return *this;
    }

    ~mpi_sub_comm() {
        destroy_comm();
    }
};

} // namespace doctest

#endif // DOCTEST_SUB_COMM_H
