#ifndef HMLP_MPI_HPP

#define HMLP_MPI_HPP

#include <external/mpi_prototypes.h>

#include <vector>

#include <stdlib.h>

#include <stdint.h>

#include <limits.h>

#include <type_traits>

#include <cstdio>

#include <cstdint>

#include <cassert>

#include <vector>

#if SIZE_MAX == UCHAR_MAX

#define HMLP_MPI_SIZE_T MPI_UNSIGNED_CHAR

#elif SIZE_MAX == USHRT_MAX

#define HMLP_MPI_SIZE_T MPI_UNSIGNED_SHORT

#elif SIZE_MAX == UINT_MAX

#define HMLP_MPI_SIZE_T MPI_UNSIGNED

#elif SIZE_MAX == ULONG_MAX

#define HMLP_MPI_SIZE_T MPI_UNSIGNED_LONG

#elif SIZE_MAX == ULLONG_MAX

#define HMLP_MPI_SIZE_T MPI_UNSIGNED_LONG_LONG

#else

#error "what is happening here?"

#endif

/** -lmemkind */

#ifdef HMLP_MIC_AVX512

#include <hbwmalloc.h>

#include <hbw_allocator.h>

#endif // ifdef HMLP}_MIC_AVX512

using namespace std;

namespace hmlp

{

/** MPI compatable interface */

namespace mpi

{

typedef MPI_Status Status;

typedef MPI_Request Request;

typedef MPI_Comm Comm;

typedef MPI_Datatype Datatype;

typedef MPI_Op Op;

/** MPI 1.0 functionality */

int Init( int *argc, char ***argv );

int Initialized( int *flag );

int Finalize( void );

int Finalized( int *flag );

int Send( const void *buf, int count, Datatype datatype, int dest, int tag, Comm comm );

int Isend( const void *buf, int count, Datatype datatype, int dest, int tag, Comm comm, Request *request );

int Recv( void *buf, int count, Datatype datatype, int source, int tag, Comm comm, Status *status );

int Irecv( void *buf, int count, Datatype datatype, int source, int tag, Comm comm, Request *request );

int Sendrecv( void *sendbuf, int sendcount, Datatype sendtype, int dest, int sendtag, void *recvbuf, int recvcount, Datatype recvtype, int source, int recvtag, Comm comm, Status *status );

int Get_count( Status *status, Datatype datatype, int *count );

int Comm_size( Comm comm, int *size );

int Comm_rank( Comm comm, int *rank );

int Comm_dup( Comm comm, Comm *newcomm );

int Comm_split( Comm comm, int color, int key, Comm *newcomm );

int Type_contiguous( int count, Datatype oldtype, Datatype *newtype );

int Type_commit( Datatype *datatype );

int Test( Request *request, int *flag, Status *status );

int Barrier( Comm comm );

int Ibarrier( Comm comm, Request *request );

int Bcast( void *buffer, int count, Datatype datatype, int root, Comm comm );

int Reduce( void *sendbuf, void *recvbuf, int count, Datatype datatype, Op op, int root, Comm comm );

int Gather( const void *sendbuf, int sendcount, Datatype sendtype, void *recvbuf, int recvcount, Datatype recvtype, int root, Comm comm );

int Gatherv( void *sendbuf, int sendcount, Datatype sendtype, void *recvbuf, const int *recvcounts, const int *displs, Datatype recvtype, int root, Comm comm );

int Scan( void *sendbuf, void *recvbuf, int count, Datatype datatype, Op op, Comm comm );

int Allreduce( void* sendbuf, void* recvbuf, int count, Datatype datatype, Op op, Comm comm );

int Allgather( void *sendbuf, int sendcount, Datatype sendtype, void *recvbuf, int recvcount, Datatype recvtype, Comm comm );

int Allgatherv( void *sendbuf, int sendcount, Datatype sendtype, void *recvbuf, int *recvcounts, int *displs, Datatype recvtype, Comm comm );

int Alltoall( void *sendbuf, int sendcount, Datatype sendtype, void *recvbuf, int recvcount, Datatype recvtype, Comm comm );

int Alltoallv( void *sendbuf, int *sendcounts, int *sdispls, Datatype sendtype, void *recvbuf, int *recvcounts, int *rdispls, Datatype recvtype, Comm comm );

/** MPI 2.0 and 3.0 functionality */

int Init_thread( int *argc, char ***argv, int required, int *provided );

int Probe( int source, int tag, Comm comm, Status *status );

int Iprobe( int source, int tag, Comm comm, int *flag, Status *status );

/** HMLP MPI extension */

void PrintProgress( const char *s, mpi::Comm comm );

class MPIObject

{

  public:

    MPIObject() {};

    {

    mpi::Comm GetComm() { return comm; };

      mpi::Comm GetPrivateComm() { return private_comm; };

      int GetCommSize() { return comm_size; };

      int GetCommRank() { return comm_rank; };

      int Comm_size()

      {

         int size;

         mpi::Comm_size( comm, &size );

         return size;

      };

      int Comm_rank()

      {

         int rank;

         mpi::Comm_rank( comm, &rank );

         return rank;

      };

    int PrivateBarrier() { return mpi::Barrier( private_comm ); };

  private:

    /** This is the assigned communicator. */

    mpi::Comm comm = MPI_COMM_WORLD;

    /** This communicator is duplicated from comm */

    mpi::Comm private_comm;

    /** (Default) there is only one MPI process. */

    int comm_size = 1;

    /** (Default) the only MPI process has rank 0. */

    int comm_rank = 0;

}; /** end class MPIObject */

/** Type insensitive interface */

template<typename T>

struct NumberIntPair { T val; int key; };

template<typename T>

{

  Datatype datatype;

  if ( std::is_same<T, int>::value )

    datatype = MPI_INT;

  else if ( std::is_same<T, float>::value )

    datatype = MPI_FLOAT;

  else if ( std::is_same<T, double>::value )

    datatype = MPI_DOUBLE;

  else if ( std::is_same<T, size_t>::value )

    datatype = HMLP_MPI_SIZE_T;

  else if ( std::is_same<T, NumberIntPair<float> >::value )

    datatype = MPI_FLOAT_INT;

  else if ( std::is_same<T, NumberIntPair<double> >::value )

  {

    datatype = MPI_DOUBLE_INT;

  }

  else

  {

    //printf( "request datatype is not supported in the simplified interface\n" );

    //exit( 1 );

  }

}; /** end GetMPIDatatype() */

template<typename TSEND>

int Send( TSEND *buf, int count,

    int dest, int tag, Comm comm )

{

  Datatype datatype = GetMPIDatatype<TSEND>();

}; /** end Send() */

template<typename TSEND>

int Isend( TSEND *buf, int count,

		int dest, int tag, Comm comm, Request *request )

{

  Datatype datatype = GetMPIDatatype<TSEND>();

}; /** end Isend() */

template<typename TRECV>

int Recv( TRECV *buf, int count,

    int source, int tag, Comm comm, Status *status )

{

  Datatype datatype = GetMPIDatatype<TRECV>();

}; /** end Recv() */

template<typename T>

int Bcast( T *buffer, int count, int root, Comm comm )

{

  Datatype datatype = GetMPIDatatype<T>();

}; /** end Bcast() */

template<typename TSEND, typename TRECV>

int Sendrecv(

    TSEND *sendbuf, int sendcount, int dest,   int sendtag,

    TRECV *recvbuf, int recvcount, int source, int recvtag,

    Comm comm, Status *status )

{

  Datatype sendtype = GetMPIDatatype<TSEND>();

  Datatype recvtype = GetMPIDatatype<TRECV>();

  return Sendrecv(

      sendbuf, sendcount, sendtype,   dest, sendtag,

      recvbuf, recvcount, recvtype, source, recvtag,

}; /** end Sendrecv() */

template<typename T>

int Reduce( T *sendbuf, T *recvbuf, int count,

    Op op, int root, Comm comm )

{

  Datatype datatype = GetMPIDatatype<T>();

}; /** end Reduce() */

template<typename TSEND, typename TRECV>

int Gather( const TSEND *sendbuf, int sendcount,

    TRECV *recvbuf, int recvcount,

    int root, Comm comm )

{

  Datatype sendtype = GetMPIDatatype<TSEND>();

  Datatype recvtype = GetMPIDatatype<TRECV>();

  return Gather( sendbuf, sendcount, sendtype,

};

template<typename TSEND, typename TRECV>

int Gatherv(

    TSEND *sendbuf, int sendcount,

    TRECV *recvbuf, const int *recvcounts, const int *displs,

    int root, Comm comm )

{

  Datatype sendtype = GetMPIDatatype<TSEND>();

  Datatype recvtype = GetMPIDatatype<TRECV>();

  return Gatherv( sendbuf, sendcount, sendtype,

}; /** end Gatherv() */

template<typename T>

int Allreduce( T* sendbuf, T* recvbuf, int count, Op op, Comm comm )

{

  Datatype datatype = GetMPIDatatype<T>();

}; /** end Allreduce() */

template<typename TSEND, typename TRECV>

int Allgather(

    TSEND *sendbuf, int sendcount,

    TRECV *recvbuf, int recvcount, Comm comm )

{

  Datatype sendtype = GetMPIDatatype<TSEND>();

  Datatype recvtype = GetMPIDatatype<TRECV>();

  return Allgather(

      sendbuf, sendcount, sendtype,

}; /** end Allgather() */

template<typename TSEND, typename TRECV>

    TSEND *sendbuf, int sendcount,

    TRECV *recvbuf, int *recvcounts, int *displs, Comm comm )

{

  return Allgatherv(

      sendbuf, sendcount, sendtype,

}; /** end Allgatherv() */

template<typename TSEND, typename TRECV>

int Alltoall(

    TSEND *sendbuf, int sendcount,

    TRECV *recvbuf, int recvcount, Comm comm )

{

  Datatype sendtype = GetMPIDatatype<TSEND>();

  Datatype recvtype = GetMPIDatatype<TRECV>();

  return Alltoall(

      sendbuf, sendcount, sendtype,

}; /** end Alltoall() */

template<typename TSEND, typename TRECV>

    TSEND *sendbuf, int *sendcounts, int *sdispls,

    TRECV *recvbuf, int *recvcounts, int *rdispls, Comm comm )

{

  return Alltoallv(

      sendbuf, sendcounts, sdispls, sendtype,

}; /** end Alltoallv() */

/**

 *  @brief This is a short hand for sending a vector, which

 *         involves two MPI_Send() calls.

 */

#ifdef HMLP_MIC_AVX512

/** use hbw::allocator for Intel Xeon Phi */

template<class T, class Allocator = hbw::allocator<T> >

#else

/** use default stl allocator */

template<class T, class Allocator = std::allocator<T> >

#endif

    std::vector<T, Allocator> &bufvector, int dest, int tag, Comm comm )

{

  Datatype datatype = GetMPIDatatype<T>();

  /** send the count size first */

  //printf( "beg send count %lu to %d\n", count, dest );

  Send( &count, 1, dest, tag, comm );

  //printf( "end send count %lu to %d\n", count, dest );

  /** now send the vector itself */

}; /** end SendVector() */

/**

 *  @brief This is a short hand for receving a vector, which

 *         involves two MPI_Recv() calls.

 */

#ifdef HMLP_MIC_AVX512

/** use hbw::allocator for Intel Xeon Phi */

template<class T, class Allocator = hbw::allocator<T> >

#else

/** use default stl allocator */

template<class T, class Allocator = std::allocator<T> >

#endif

    std::vector<T, Allocator> &bufvector, int source, int tag, Comm comm, Status *status )

{

  Datatype datatype = GetMPIDatatype<T>();

  /** recv the count size first */

  //printf( "beg recv count %lu from %d\n", count, source );

  Recv( &count, 1, source, tag, comm, status );

  //printf( "end recv count %lu from %d\n", count, source );

  /** resize receiving vector */

  /** now recv the vector itself */

}; /** end RecvVector() */

#ifdef HMLP_MIC_AVX512

/** use hbw::allocator for Intel Xeon Phi */

template<class T, class Allocator = hbw::allocator<T> >

#else

/** use default stl allocator */

template<class T, class Allocator = std::allocator<T> >

#endif

    vector<T, Allocator> &sendvector,   int dest, int sendtag,

    vector<T, Allocator> &recvvector, int source, int recvtag,

    Comm comm, Status *status )

{

  Datatype datatype = GetMPIDatatype<T>();

  /** exechange size */

  Sendrecv(

      &send_size, 1,   dest, sendtag,

      &recv_size, 1, source, recvtag,

      comm, status );

  /** resize receiving vector */

  /** exechange buffer */

      sendvector.data(), send_size,   dest, sendtag,

      recvvector.data(), recv_size, source, recvtag,

}; /** end ExchangeVector() */

template<typename T>

{

  Datatype datatype = GetMPIDatatype<T>();

  /** Use typeless Gather() for counting. */

  Gather( &send_size, 1, recv_sizes.data(), 1, root, comm );

  /** Accumulate received displs. */

  /** Resize recvvector to fit. */

  /** Gather vectors. */

}; /** end GatherVector() */

#ifdef HMLP_MIC_AVX512

/** use hbw::allocator for Intel Xeon Phi */

template<class T, class Allocator = hbw::allocator<T> >

#else

/** use default stl allocator */

template<class T, class Allocator = std::allocator<T> >

#endif

    const vector<vector<T, Allocator>>& sendvector,

    vector<vector<T, Allocator>>& recvvector, Comm comm )

{

  vector<T> sendbuf;

  vector<T> recvbuf;

  //printf( "rank %d ", rank );

  {

        sendvector[ p ].begin(),

        sendvector[ p ].end() );

    //printf( "%d ", sendcounts[ j ] );

  }

  //printf( "\n" );

  //printf( "before Alltoall (mpi size = %d)\n", size ); fflush( stdout );

  //Barrier( comm );

  /** exchange sendcount */

  Alltoall( sendcounts.data(), 1, recvcounts.data(), 1, comm );

  //printf( "after Alltoall\n" ); fflush( stdout );

  //Barrier( comm );

  //printf( "rank %d ", rank );

  size_t total_recvcount = 0;

  {

    //printf( "%d ", recvcounts[ j ] );

  }

  //printf( "\n" );

  /** resize receving buffer */

  //printf( "before Alltoall2 recvbuf.size() %lu\n", recvbuf.size() ); fflush( stdout );

  //Barrier( comm );

      sendbuf.data(), sendcounts.data(), sdispls.data(),

      recvbuf.data(), recvcounts.data(), rdispls.data(), comm );

  //printf( "after Alltoall2\n" ); fflush( stdout );

  //Barrier( comm );

  {

        recvbuf.begin() + rdispls[ p ],

  }

}; /** end AlltoallVector() */

/**

 *  This interface supports hmlp::Data. [ m, n ] will be set

 *  after the routine.

 *

 */

//template<typename T>

//int AlltoallVector(

//    std::vector<hmlp::Data<T>> &sendvector,

//    std::vector<hmlp::Data<T>> &recvvector, Comm comm )

//{

//  int size = 0;

//  int rank = 0;

//  Comm_size( comm, &size );

//  Comm_rank( comm, &rank );

//

//  assert( sendvector.size() == size );

//  assert( recvvector.size() == size );

//

//  std::vector<T> sendbuf;

//  std::vector<T> recvbuf;

//  std::vector<int> sendcounts( size, 0 );

//  std::vector<int> recvcounts( size, 0 );

//  std::vector<int> sdispls( size + 1, 0 );

//  std::vector<int> rdispls( size + 1, 0 );

//

//

//  for ( size_t p = 0; p < size; p ++ )

//  {

//    sendcounts[ p ] = sendvector[ p ].size();

//    sdispls[ p + 1 ] = sdispls[ p ] + sendcounts[ p ];

//    sendbuf.insert( sendbuf.end(),

//        sendvector[ p ].begin(),

//        sendvector[ p ].end() );

//  }

//

//  /** exchange sendcount */

//  Alltoall( sendcounts.data(), 1, recvcounts.data(), 1, comm );

//

//  /** compute total receiving count */

//  size_t total_recvcount = 0;

//  for ( size_t p = 0; p < size; p ++ )

//  {

//    rdispls[ p + 1 ] = rdispls[ p ] + recvcounts[ p ];

//    total_recvcount += recvcounts[ p ];

//  }

//

//  /** resize receving buffer */

//  recvbuf.resize( total_recvcount );

//

//

//

//

//

//

//}; /** end AlltoallVector() */

}; /** end namespace mpi */

}; /** end namespace hmlp */

#endif /** define HMLP_MPI_HPP */