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


namespace ep_lib
{


  int MPI_Win_create(void *base, MPI_Aint size, int disp_unit, MPI_Info info, MPI_Comm comm, 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->is_ep = comm.is_ep;

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

    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_create from server, rank = %d\n", rank);
      for(int i=0; i<num_ep_max; i++)
      {
        ::MPI_Win_create(base, to_mpi_aint(size.mpi_aint), disp_unit, to_mpi_info(info.mpi_info), to_mpi_comm(comm.mpi_comm), 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_max; i++)
      {
        if(rank_loc == i)
        {
          printf("Calling MPI_Win_create from client, rank = %d\n", rank);
          ::MPI_Win_create(base, to_mpi_aint(size.mpi_aint), disp_unit, to_mpi_info(info.mpi_info), to_mpi_comm(comm.mpi_comm), static_cast< ::MPI_Win*>(win->client_win)); 
          win->comm = comm;
        }
        MPI_Barrier_local(comm);
      }

      // #pragma omp critical (win_create)
      // {
      //   ::MPI_Win_create(base, to_mpi_aint(size.mpi_aint), disp_unit, to_mpi_info(info.mpi_info), to_mpi_comm(comm.mpi_comm), static_cast< ::MPI_Win*>(win->client_win)); 
      //   win->comm = comm;
      // }
    }
    return 0;
  }



  int MPI_Win_free(MPI_Win *win)
  {
    int rank, rank_loc, num_ep;
    

    MPI_Comm_rank(win->comm, &rank);
    rank_loc = win->comm.ep_comm_ptr->size_rank_info[1].first;
    num_ep = win->comm.ep_comm_ptr->size_rank_info[1].second;
    
    
    int num_ep_max;
    MPI_Allreduce(&num_ep, &num_ep_max, 1, MPI_INT, MPI_MAX, win->comm);

    //printf("rank_loc = %d, thread_num = %d, num_ep_max = %d\n", rank_loc, omp_get_thread_num(), num_ep_max);

    if(num_ep == 1)
    {
      printf("Calling MPI_Win_free from server, rank = %d\n", rank);
      for(int i=0; i<num_ep_max; i++)
      {
        ::MPI_Win_free(static_cast< ::MPI_Win*>(win->server_win[i]));
      }
    }
    else
    {
      for(int i=0; i<num_ep_max; i++)
      {
        if(rank_loc == i)
        {
          printf("Calling MPI_Win_free from client, rank = %d\n", rank);
          ::MPI_Win_free(static_cast< ::MPI_Win*>(win->client_win)); 
        }
        MPI_Barrier_local(win->comm);
      }
    }
    return 0;
  }



  int MPI_Win_fence(int assert, MPI_Win win)
  {
    int rank, rank_loc, num_ep;
    
    MPI_Comm_rank(win.comm, &rank);
    rank_loc = win.comm.ep_comm_ptr->size_rank_info[1].first;
    num_ep = win.comm.ep_comm_ptr->size_rank_info[1].second;
    
    int num_ep_max;
    MPI_Allreduce(&num_ep, &num_ep_max, 1, MPI_INT, MPI_MAX, win.comm);

    if(num_ep == 1)
    {
      printf("Calling MPI_Win_fence from server, rank = %d\n", rank);
      for(int i=0; i<num_ep_max; i++)
      {
        ::MPI_Win_fence(assert, *static_cast< ::MPI_Win*>(win.server_win[i]));
      }
    }
    else
    {
      for(int i=0; i<num_ep_max; i++)
      {
        if(rank_loc == i)
        {
          printf("Calling MPI_Win_fence from client, rank = %d\n", rank);
          ::MPI_Win_fence(assert, *static_cast< ::MPI_Win*>(win.client_win)); 
        }
        MPI_Barrier_local(win.comm);
      }
    }

    return 0;  
  }

}