// SPDX-FileComment: This file is part of TNL - Template Numerical Library (https://tnl-project.org/)
// SPDX-License-Identifier: MIT

#pragma once

#include <stdexcept>

#ifdef HAVE_MPI
   #include <mpi.h>
#else
   #include "DummyDefs.h"
   #include <cstring>  // std::memcpy
   #include <TNL/Exceptions/MPISupportMissing.h>
#endif

#include <TNL/Assert.h>
#include "getDataType.h"
#include "Profiling.h"

namespace TNL::MPI {

// wrappers for basic MPI functions

[[nodiscard]] inline bool
Initialized()
{
#ifdef HAVE_MPI
   int flag;
   MPI_Initialized( &flag );
   return flag != 0;
#else
   return true;
#endif
}

[[nodiscard]] inline bool
Finalized()
{
#ifdef HAVE_MPI
   int flag;
   MPI_Finalized( &flag );
   return flag != 0;
#else
   return false;
#endif
}

// wrappers for MPI helper functions

/**
 * \brief Wrapper for \ref MPI_Dims_create.
 *
 * \param nnodes - number of nodes in the grid
 * \param ndims - number of dimensions of the Cartesian grid
 * \param dims - distribution of processes into the \e dim-dimensional
 *               Cartesian grid (array of length \e ndims)
 *
 * Negative input values of \e dims[i] are erroneous. An error will occur if
 * \e nnodes is not a multiple of the product of all non-zero values \e dims[i].
 *
 * See the MPI documentation for more information.
 */
inline void
Compute_dims( int nnodes, int ndims, int* dims )
{
#ifdef HAVE_MPI
   int prod = 1;
   for( int i = 0; i < ndims; i++ ) {
      if( dims[ i ] < 0 )
         throw std::invalid_argument( "Negative value passed to MPI::Compute_dims in the dims array argument." );
      if( dims[ i ] > 0 )
         prod *= dims[ i ];
   }

   if( nnodes % prod != 0 )
      throw std::logic_error( "The program tries to call MPI_Dims_create with wrong dimensions."
                              "The product of the non-zero values dims[i] is "
                              + std::to_string( prod )
                              + " and the "
                                "number of processes ("
                              + std::to_string( nnodes ) + ") is not a multiple of the product." );

   MPI_Dims_create( nnodes, ndims, dims );
#else
   for( int i = 0; i < ndims; i++ )
      dims[ i ] = 1;
#endif
}

// wrappers for MPI communication functions

inline void
Barrier( MPI_Comm communicator = MPI_COMM_WORLD )
{
   if( communicator == MPI_COMM_NULL )
      throw std::invalid_argument( "Barrier cannot be called with MPI_COMM_NULL" );
#ifdef HAVE_MPI
   if( ! Initialized() || Finalized() )
      throw std::logic_error( "MPI is not initialized" );
   MPI_Barrier( communicator );
#endif
}

inline void
Waitall( MPI_Request* reqs, int length )
{
#ifdef HAVE_MPI
   if( ! Initialized() || Finalized() )
      throw std::logic_error( "MPI is not initialized" );
   MPI_Waitall( length, reqs, MPI_STATUSES_IGNORE );
#endif
}

template< typename T >
void
Send( const T* data, int count, int dest, int tag, MPI_Comm communicator = MPI_COMM_WORLD )
{
   if( communicator == MPI_COMM_NULL )
      throw std::invalid_argument( "Send cannot be called with MPI_COMM_NULL" );
#ifdef HAVE_MPI
   if( ! Initialized() || Finalized() )
      throw std::logic_error( "MPI is not initialized" );
   MPI_Send( reinterpret_cast< const void* >( data ), count, getDataType< T >(), dest, tag, communicator );
#endif
}

template< typename T >
void
Recv( T* data, int count, int src, int tag, MPI_Comm communicator = MPI_COMM_WORLD )
{
   if( communicator == MPI_COMM_NULL )
      throw std::invalid_argument( "Recv cannot be called with MPI_COMM_NULL" );
#ifdef HAVE_MPI
   if( ! Initialized() || Finalized() )
      throw std::logic_error( "MPI is not initialized" );
   MPI_Recv( reinterpret_cast< void* >( data ), count, getDataType< T >(), src, tag, communicator, MPI_STATUS_IGNORE );
#endif
}

template< typename T >
void
Sendrecv( const T* sendData,
          int sendCount,
          int destination,
          int sendTag,
          T* receiveData,
          int receiveCount,
          int source,
          int receiveTag,
          MPI_Comm communicator = MPI_COMM_WORLD )
{
   if( communicator == MPI_COMM_NULL )
      throw std::invalid_argument( "Sendrecv cannot be called with MPI_COMM_NULL" );
#ifdef HAVE_MPI
   if( ! Initialized() || Finalized() )
      throw std::logic_error( "MPI is not initialized" );
   MPI_Sendrecv( reinterpret_cast< const void* >( sendData ),
                 sendCount,
                 getDataType< T >(),
                 destination,
                 sendTag,
                 reinterpret_cast< void* >( receiveData ),
                 receiveCount,
                 getDataType< T >(),
                 source,
                 receiveTag,
                 communicator,
                 MPI_STATUS_IGNORE );
#else
   throw Exceptions::MPISupportMissing();
#endif
}

template< typename T >
[[nodiscard]] MPI_Request
Isend( const T* data, int count, int dest, int tag, MPI_Comm communicator = MPI_COMM_WORLD )
{
   if( communicator == MPI_COMM_NULL )
      throw std::invalid_argument( "Isend cannot be called with MPI_COMM_NULL" );
#ifdef HAVE_MPI
   if( ! Initialized() || Finalized() )
      throw std::logic_error( "MPI is not initialized" );
   MPI_Request req;
   MPI_Isend( reinterpret_cast< const void* >( data ), count, getDataType< T >(), dest, tag, communicator, &req );
   return req;
#else
   return MPI_REQUEST_NULL;
#endif
}

template< typename T >
[[nodiscard]] MPI_Request
Irecv( T* data, int count, int src, int tag, MPI_Comm communicator = MPI_COMM_WORLD )
{
   if( communicator == MPI_COMM_NULL )
      throw std::invalid_argument( "Irecv cannot be called with MPI_COMM_NULL" );
#ifdef HAVE_MPI
   if( ! Initialized() || Finalized() )
      throw std::logic_error( "MPI is not initialized" );
   MPI_Request req;
   MPI_Irecv( reinterpret_cast< void* >( data ), count, getDataType< T >(), src, tag, communicator, &req );
   return req;
#else
   return MPI_REQUEST_NULL;
#endif
}

template< typename T >
void
Allreduce( const T* data, T* reduced_data, int count, const MPI_Op& op, MPI_Comm communicator = MPI_COMM_WORLD )
{
   if( communicator == MPI_COMM_NULL )
      throw std::invalid_argument( "Allreduce cannot be called with MPI_COMM_NULL" );
#ifdef HAVE_MPI
   if( ! Initialized() || Finalized() )
      throw std::logic_error( "MPI is not initialized" );
   getTimerAllreduce().start();
   MPI_Allreduce( reinterpret_cast< const void* >( data ),
                  reinterpret_cast< void* >( reduced_data ),
                  count,
                  getDataType< T >(),
                  op,
                  communicator );
   getTimerAllreduce().stop();
#else
   std::memcpy( reinterpret_cast< void* >( reduced_data ), reinterpret_cast< const void* >( data ), count * sizeof( T ) );
#endif
}

// in-place variant of Allreduce
template< typename T >
void
Allreduce( T* data, int count, const MPI_Op& op, MPI_Comm communicator = MPI_COMM_WORLD )
{
   if( communicator == MPI_COMM_NULL )
      throw std::invalid_argument( "Allreduce cannot be called with MPI_COMM_NULL" );
#ifdef HAVE_MPI
   if( ! Initialized() || Finalized() )
      throw std::logic_error( "MPI is not initialized" );
   getTimerAllreduce().start();
   MPI_Allreduce( MPI_IN_PLACE, reinterpret_cast< void* >( data ), count, getDataType< T >(), op, communicator );
   getTimerAllreduce().stop();
#endif
}

template< typename T >
void
Reduce( const T* data, T* reduced_data, int count, const MPI_Op& op, int root, MPI_Comm communicator = MPI_COMM_WORLD )
{
   if( communicator == MPI_COMM_NULL )
      throw std::invalid_argument( "Reduce cannot be called with MPI_COMM_NULL" );
#ifdef HAVE_MPI
   if( ! Initialized() || Finalized() )
      throw std::logic_error( "MPI is not initialized" );
   MPI_Reduce( reinterpret_cast< const void* >( data ),
               reinterpret_cast< void* >( reduced_data ),
               count,
               getDataType< T >(),
               op,
               root,
               communicator );
#else
   std::memcpy( reinterpret_cast< void* >( reduced_data ), reinterpret_cast< const void* >( data ), count * sizeof( T ) );
#endif
}

template< typename T >
void
Bcast( T* data, int count, int root, MPI_Comm communicator = MPI_COMM_WORLD )
{
   if( communicator == MPI_COMM_NULL )
      throw std::invalid_argument( "Bcast cannot be called with MPI_COMM_NULL" );
#ifdef HAVE_MPI
   if( ! Initialized() || Finalized() )
      throw std::logic_error( "MPI is not initialized" );
   MPI_Bcast( reinterpret_cast< void* >( data ), count, getDataType< T >(), root, communicator );
#endif
}

template< typename T >
void
Alltoall( const T* sendData, int sendCount, T* receiveData, int receiveCount, MPI_Comm communicator = MPI_COMM_WORLD )
{
   if( communicator == MPI_COMM_NULL )
      throw std::invalid_argument( "Alltoall cannot be called with MPI_COMM_NULL" );
#ifdef HAVE_MPI
   if( ! Initialized() || Finalized() )
      throw std::logic_error( "MPI is not initialized" );
   MPI_Alltoall( reinterpret_cast< const void* >( sendData ),
                 sendCount,
                 getDataType< T >(),
                 reinterpret_cast< void* >( receiveData ),
                 receiveCount,
                 getDataType< T >(),
                 communicator );
#else
   if( sendCount != receiveCount )
      throw std::invalid_argument( "Alltoall: sendCount must be equal to receiveCount when running without MPI" );
   std::memcpy(
      reinterpret_cast< void* >( receiveData ), reinterpret_cast< const void* >( sendData ), sendCount * sizeof( T ) );
#endif
}

}  // namespace TNL::MPI
