#include "ep_lib.hpp"
#include <mpi.h>
#include "ep_declaration.hpp"
#include "ep_mpi.hpp"

namespace ep_lib
{



  int MPI_Alloc_mem(MPI_Aint size, MPI_Info info, void *baseptr)
  {
    ::MPI_Alloc_mem(to_mpi_aint(size), to_mpi_info(info), baseptr);
    //::MPI_Alloc_mem(size.mpi_aint, MPI_INFO_NULL_STD, baseptr);
    return 0;
   }
   
  int MPI_Alloc_mem(unsigned long size, MPI_Info info, void *baseptr)
  {
    ::MPI_Alloc_mem(size, *(static_cast< ::MPI_Info*>(info->mpi_info)), baseptr);
    //::MPI_Alloc_mem(size, MPI_INFO_NULL_STD, baseptr);
    return 0;
  }

  int MPI_Win_allocate (MPI_Aint size, int disp_unit, MPI_Info info, MPI_Comm comm, void *baseptr, MPI_Win *win)
  {
    int rank, rank_loc, num_ep;
    MPI_Comm_rank(comm, &rank);
    rank_loc = comm->ep_comm_ptr->size_rank_info[1].first;
    num_ep = comm->ep_comm_ptr->size_rank_info[1].second;
    
    *win = new ep_win;
    
    (*win)->is_ep = comm->is_ep;

    int num_ep_max;
    MPI_Allreduce(&num_ep, &num_ep_max, 1, MPI_INT, MPI_MAX, comm);

    assert(num_ep_max > 1);
    
    info = new ep_info;

    if(num_ep == 1)  // server
    {
      for(int i=0; i<num_ep_max; i++)
      {
        (*win)->server_win[i] = new ::MPI_Win;
      }

      printf("Calling MPI_Win_allocate from server, rank = %d\n", rank);
      for(int i=0; i<num_ep_max; i++)
      {
        ::MPI_Win_allocate(to_mpi_aint(size), disp_unit, to_mpi_info(info), to_mpi_comm(comm->mpi_comm), baseptr, static_cast< ::MPI_Win*>((*win)->server_win[i]));
        (*win)->comm = comm;
      }
    }
    else  // client
    {
      (*win)->client_win = new ::MPI_Win;
      
      for(int i=0; i<num_ep; i++)
      {
        if(rank_loc == i)
        {
          printf("Calling MPI_Win_allocate from client, rank = %d\n", rank);
          ::MPI_Win_allocate(to_mpi_aint(size), disp_unit, to_mpi_info(info), to_mpi_comm(comm->mpi_comm), baseptr, static_cast< ::MPI_Win*>((*win)->client_win)); 
          (*win)->comm = comm;
        }
        MPI_Barrier_local(comm);
      }

      if(num_ep_max > num_ep)
      {
        #pragma omp master
        {
          for(int i=0; i<num_ep_max-num_ep; i++)
          {
            (*win)->null_win[i] = new ::MPI_Win;
            ::MPI_Win_allocate(0, disp_unit, to_mpi_info(info), to_mpi_comm(comm->mpi_comm), baseptr, static_cast< ::MPI_Win*>((*win)->null_win[i])); 
          }  
        }  
      }

    }
    return 0;
  }

}