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

    }
    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; 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);
      }


      if(num_ep_max > num_ep)
      {
        #pragma omp master
        {
          for(int i=0; i<num_ep_max-num_ep; i++)
          {
            ::MPI_Win_free(static_cast< ::MPI_Win*>((*win)->null_win[i])); 
          }  
        }  
      }

    }
    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; 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);
      }


      if(num_ep_max > num_ep)
      {
        #pragma omp master
        {
          for(int i=0; i<num_ep_max-num_ep; i++)
          {
            ::MPI_Win_fence(assert, *static_cast< ::MPI_Win*>(win->null_win[i])); 
          }  
        }  
      }
    }

    return 0;  
  }

}