///
/// Functions and classes for handling parallel computation environment
///
/// @author: Yunhai Li (yhli)
///

#ifndef TBPLAS_BASE_PARALLEL_H
#define TBPLAS_BASE_PARALLEL_H

#ifdef WITH_MPI
#include <mpi.h>
#endif
#include <complex>
#include <eigen3/Eigen/Dense>
#include <tuple>

// Fake datatypes when MPI is disabled
#ifndef WITH_MPI
typedef int MPI_Comm;
#endif

namespace tbplas::base {

// Functions for initializing and finalizing MPI environments
// Should be called in the main function
void MPIENV_INIT(int argc, char* argv[]);
void MPIENV_FINALIZE();

// Wrappers over MPI calls, dedicated for user-defined MPI environment classes
// ----- Generated with macro expansion -----
/* begin func_claim */
void AllreduceScalar_(const MPI_Comm& comm, const int& source, int& dest);
void InplaceAllreduceScalar_(const MPI_Comm& comm, int& source);
void BcastScalar_(const MPI_Comm& comm, int& source, int root);
void SendScalar_(const MPI_Comm& comm, const int& source, int destID, int tag);
void RecvScalar_(const MPI_Comm& comm, int& dest, int sourceID, int tag);
/* end func_claim */
/* begin func_claim */
void AllreduceScalar_(const MPI_Comm& comm, const double& source, double& dest);
void InplaceAllreduceScalar_(const MPI_Comm& comm, double& source);
void BcastScalar_(const MPI_Comm& comm, double& source, int root);
void SendScalar_(const MPI_Comm& comm, const double& source, int destID, int tag);
void RecvScalar_(const MPI_Comm& comm, double& dest, int sourceID, int tag);
/* end func_claim */
/* begin func_claim */
void AllreduceScalar_(const MPI_Comm& comm, const std::complex<double>& source, std::complex<double>& dest);
void InplaceAllreduceScalar_(const MPI_Comm& comm, std::complex<double>& source);
void BcastScalar_(const MPI_Comm& comm, std::complex<double>& source, int root);
void SendScalar_(const MPI_Comm& comm, const std::complex<double>& source, int destID, int tag);
void RecvScalar_(const MPI_Comm& comm, std::complex<double>& dest, int sourceID, int tag);
/* end func_claim */
/* begin func_claim */
void AllreduceMatrixXi_(const MPI_Comm& comm, const Eigen::MatrixXi& source, Eigen::MatrixXi& dest);
void InplaceAllreduceMatrixXi_(const MPI_Comm& comm, Eigen::MatrixXi& source);
void BcastMatrixXi_(const MPI_Comm& comm, Eigen::MatrixXi& source, int root);
void SendMatrixXi_(const MPI_Comm& comm, const Eigen::MatrixXi& source, int destID, int tag);
void RecvMatrixXi_(const MPI_Comm& comm, Eigen::MatrixXi& dest, int sourceID, int tag);
/* end func_claim */
/* begin func_claim */
void AllreduceMatrixXd_(const MPI_Comm& comm, const Eigen::MatrixXd& source, Eigen::MatrixXd& dest);
void InplaceAllreduceMatrixXd_(const MPI_Comm& comm, Eigen::MatrixXd& source);
void BcastMatrixXd_(const MPI_Comm& comm, Eigen::MatrixXd& source, int root);
void SendMatrixXd_(const MPI_Comm& comm, const Eigen::MatrixXd& source, int destID, int tag);
void RecvMatrixXd_(const MPI_Comm& comm, Eigen::MatrixXd& dest, int sourceID, int tag);
/* end func_claim */
/* begin func_claim */
void AllreduceMatrixXcd_(const MPI_Comm& comm, const Eigen::MatrixXcd& source, Eigen::MatrixXcd& dest);
void InplaceAllreduceMatrixXcd_(const MPI_Comm& comm, Eigen::MatrixXcd& source);
void BcastMatrixXcd_(const MPI_Comm& comm, Eigen::MatrixXcd& source, int root);
void SendMatrixXcd_(const MPI_Comm& comm, const Eigen::MatrixXcd& source, int destID, int tag);
void RecvMatrixXcd_(const MPI_Comm& comm, Eigen::MatrixXcd& dest, int sourceID, int tag);
/* end func_claim */
/* begin func_claim */
void AllreduceVectorXi_(const MPI_Comm& comm, const Eigen::VectorXi& source, Eigen::VectorXi& dest);
void InplaceAllreduceVectorXi_(const MPI_Comm& comm, Eigen::VectorXi& source);
void BcastVectorXi_(const MPI_Comm& comm, Eigen::VectorXi& source, int root);
void SendVectorXi_(const MPI_Comm& comm, const Eigen::VectorXi& source, int destID, int tag);
void RecvVectorXi_(const MPI_Comm& comm, Eigen::VectorXi& dest, int sourceID, int tag);
/* end func_claim */
/* begin func_claim */
void AllreduceVectorXd_(const MPI_Comm& comm, const Eigen::VectorXd& source, Eigen::VectorXd& dest);
void InplaceAllreduceVectorXd_(const MPI_Comm& comm, Eigen::VectorXd& source);
void BcastVectorXd_(const MPI_Comm& comm, Eigen::VectorXd& source, int root);
void SendVectorXd_(const MPI_Comm& comm, const Eigen::VectorXd& source, int destID, int tag);
void RecvVectorXd_(const MPI_Comm& comm, Eigen::VectorXd& dest, int sourceID, int tag);
/* end func_claim */
/* begin func_claim */
void AllreduceVectorXcd_(const MPI_Comm& comm, const Eigen::VectorXcd& source, Eigen::VectorXcd& dest);
void InplaceAllreduceVectorXcd_(const MPI_Comm& comm, Eigen::VectorXcd& source);
void BcastVectorXcd_(const MPI_Comm& comm, Eigen::VectorXcd& source, int root);
void SendVectorXcd_(const MPI_Comm& comm, const Eigen::VectorXcd& source, int destID, int tag);
void RecvVectorXcd_(const MPI_Comm& comm, Eigen::VectorXcd& dest, int sourceID, int tag);
/* end func_claim */

/**
 * @brief Range info for distributing jobs.
 * 
 */
class RangeInfo {
public:
    int k_min = 0;
    int k_max = 0;
};

/// @brief Base class for MPI environments
class BaseMPIEnv {
public:
    BaseMPIEnv();
    explicit BaseMPIEnv(const MPI_Comm& comm);
    [[nodiscard]] int getRank() const;
    [[nodiscard]] int getSize() const;
    [[nodiscard]] const MPI_Comm& getComm() const;
    [[nodiscard]] bool isMaster() const;
    [[nodiscard]] bool isSPMode() const;
    [[nodiscard]] bool isMPIEnabled() const;
    void barrier() const;
    void free();
    [[nodiscard]] RangeInfo distRange(int n_max, int num_group) const;

    // ----- Generated with macro expansion -----
    /* begin class_func_claim */
    void AllreduceScalar(const int& source, int& dest) const;
    void InplaceAllreduceScalar(int& source) const;
    void BcastScalar(int& source, int root = 0) const;
    void SendScalar(const int& source, int destID, int tag) const;
    void RecvScalar(int& dest, int sourceID, int tag) const;
    /* end class_func_claim */
    /* begin class_func_claim */
    void AllreduceScalar(const double& source, double& dest) const;
    void InplaceAllreduceScalar(double& source) const;
    void BcastScalar(double& source, int root = 0) const;
    void SendScalar(const double& source, int destID, int tag) const;
    void RecvScalar(double& dest, int sourceID, int tag) const;
    /* end class_func_claim */
    /* begin class_func_claim */
    void AllreduceScalar(const std::complex<double>& source, std::complex<double>& dest) const;
    void InplaceAllreduceScalar(std::complex<double>& source) const;
    void BcastScalar(std::complex<double>& source, int root = 0) const;
    void SendScalar(const std::complex<double>& source, int destID, int tag) const;
    void RecvScalar(std::complex<double>& dest, int sourceID, int tag) const;
    /* end class_func_claim */
    /* begin class_func_claim */
    void AllreduceMatrixXi(const Eigen::MatrixXi& source, Eigen::MatrixXi& dest) const;
    void InplaceAllreduceMatrixXi(Eigen::MatrixXi& source) const;
    void BcastMatrixXi(Eigen::MatrixXi& source, int root = 0) const;
    void SendMatrixXi(const Eigen::MatrixXi& source, int destID, int tag) const;
    void RecvMatrixXi(Eigen::MatrixXi& dest, int sourceID, int tag) const;
    /* end class_func_claim */
    /* begin class_func_claim */
    void AllreduceMatrixXd(const Eigen::MatrixXd& source, Eigen::MatrixXd& dest) const;
    void InplaceAllreduceMatrixXd(Eigen::MatrixXd& source) const;
    void BcastMatrixXd(Eigen::MatrixXd& source, int root = 0) const;
    void SendMatrixXd(const Eigen::MatrixXd& source, int destID, int tag) const;
    void RecvMatrixXd(Eigen::MatrixXd& dest, int sourceID, int tag) const;
    /* end class_func_claim */
    /* begin class_func_claim */
    void AllreduceMatrixXcd(const Eigen::MatrixXcd& source, Eigen::MatrixXcd& dest) const;
    void InplaceAllreduceMatrixXcd(Eigen::MatrixXcd& source) const;
    void BcastMatrixXcd(Eigen::MatrixXcd& source, int root = 0) const;
    void SendMatrixXcd(const Eigen::MatrixXcd& source, int destID, int tag) const;
    void RecvMatrixXcd(Eigen::MatrixXcd& dest, int sourceID, int tag) const;
    /* end class_func_claim */
    /* begin class_func_claim */
    void AllreduceVectorXi(const Eigen::VectorXi& source, Eigen::VectorXi& dest) const;
    void InplaceAllreduceVectorXi(Eigen::VectorXi& source) const;
    void BcastVectorXi(Eigen::VectorXi& source, int root = 0) const;
    void SendVectorXi(const Eigen::VectorXi& source, int destID, int tag) const;
    void RecvVectorXi(Eigen::VectorXi& dest, int sourceID, int tag) const;
    /* end class_func_claim */
    /* begin class_func_claim */
    void AllreduceVectorXd(const Eigen::VectorXd& source, Eigen::VectorXd& dest) const;
    void InplaceAllreduceVectorXd(Eigen::VectorXd& source) const;
    void BcastVectorXd(Eigen::VectorXd& source, int root = 0) const;
    void SendVectorXd(const Eigen::VectorXd& source, int destID, int tag) const;
    void RecvVectorXd(Eigen::VectorXd& dest, int sourceID, int tag) const;
    /* end class_func_claim */
    /* begin class_func_claim */
    void AllreduceVectorXcd(const Eigen::VectorXcd& source, Eigen::VectorXcd& dest) const;
    void InplaceAllreduceVectorXcd(Eigen::VectorXcd& source) const;
    void BcastVectorXcd(Eigen::VectorXcd& source, int root = 0) const;
    void SendVectorXcd(const Eigen::VectorXcd& source, int destID, int tag) const;
    void RecvVectorXcd(Eigen::VectorXcd& dest, int sourceID, int tag) const;
    /* end class_func_claim */

    /**
     * @brief Temporary patch for Python interface.
     * 
     */
    void BcastRawDouble(double* source, size_t num_data, int root = 0) const;
    void BcastRawComplex(std::complex<double>* source, size_t num_data, int root = 0) const;

protected:
    void updateRankSize();

protected:
    MPI_Comm m_Comm; // communicator
    int m_Rank = 0; // rank of this process in the communicator
    int m_Size = 1; // total number of processes in the communicator
};

} // namespace tbplas::base

#endif
