# cython: language_level=3
# cython: warn.undeclared=True
# cython: warn.unreachable=True
# cython: warn.maybe_uninitialized=True
# cython: warn.unused=True
# cython: warn.unused_arg=True
# cython: warn.unused_result=True
# cython: warn.multiple_declarators=True
# distutils: language = c++

import cython
from libc.stdint cimport uint64_t
from libcpp cimport bool


# C++ stuff that should be visible from Cython
cdef extern from "interface_mpi.h" namespace "tbplas::python":
    void mpienv_init_cxx()
    void mpienv_finalize_cxx()

    cdef cppclass BaseMPIEnv:
        BaseMPIEnv() except +
        int getRank()
        int getSize()
        bool isMaster()
        bool isSPMode()
        bool isMPIEnabled()
        void barrier()
        void free()
        void BcastRawDouble(double* source, uint64_t num_data, int root)
        void BcastRawComplex(double complex* source, uint64_t num_data, int root)


# Python functions and classes for importing
def mpienv_init():
    mpienv_init_cxx()


def mpienv_finalize():
    mpienv_finalize_cxx()


cdef class CPPBaseMPIEnv:
    cdef BaseMPIEnv* c_mpi_env

    def __cinit__(self):
        self.c_mpi_env = new BaseMPIEnv()

    def __dealloc__(self):
        del self.c_mpi_env

    def get_rank(self):
        return self.c_mpi_env.getRank()

    def get_size(self):
        return self.c_mpi_env.getSize()

    def is_master(self):
        return self.c_mpi_env.isMaster()

    def is_sp_mode(self):
        return self.c_mpi_env.isSPMode()

    def is_mpi_enabled(self):
        return self.c_mpi_env.isMPIEnabled()

    def barrier(self):
        self.c_mpi_env.barrier()

    def free(self):
        self.c_mpi_env.free()

    def bcast_double(self, double [::1] data):
        self.c_mpi_env.BcastRawDouble(&data[0], data.size, 0)

    def bcast_complex(self, double complex [::1] data):
        self.c_mpi_env.BcastRawComplex(&data[0], data.size, 0)
