#include <mpi.h>
#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include <string.h>

static int allocate_window(void * baseptr, size_t nbytes, MPI_Win *win);
static int deallocate_window(void * baseptr, MPI_Win *win);

#define COUNT 10
int main(int argc, char** argv)
{
  int myrank, nprocs, nprocs_per_node;

  MPI_Init(&argc, &argv);

  MPI_Comm_rank(MPI_COMM_WORLD, &myrank);
  MPI_Comm_size(MPI_COMM_WORLD, &nprocs);

  if (nprocs != 2)
  {
    if (myrank == 0) fprintf(stderr, "usage: the number of processes must be 2!\n");
    MPI_Abort(MPI_COMM_WORLD, 1);
    return EXIT_FAILURE;
  }

  MPI_Win w1;

  int *w1_mem;

  //allocate window 1
  allocate_window(&w1_mem, COUNT * sizeof(int), &w1);

  for (size_t idx = 0; idx < COUNT; ++idx) {
    w1_mem[idx] = idx * COUNT + 1 + myrank;
  }

  MPI_Win_lock_all(0, w1);

  MPI_Aint ldisp;
  //Displacements of Window 1
  MPI_Aint *w1_disps = malloc(nprocs * sizeof(MPI_Aint));
  assert(MPI_Get_address(w1_mem, &ldisp) == MPI_SUCCESS);
  assert(MPI_Allgather(&ldisp, 1, MPI_AINT, w1_disps, 1, MPI_AINT, MPI_COMM_WORLD) == MPI_SUCCESS);


/**************************************************************************
 * Communication with a remote rank --> Works
 **************************************************************************/

  int other = 1 - myrank;
  int recv_buf[COUNT];
  printf("rank %d reads memory from rank %d\n", myrank, other);
  MPI_Get(recv_buf, COUNT, MPI_INT, other, w1_disps[other], COUNT, MPI_INT, w1);
  MPI_Win_flush_all(w1);
  //verify values
  for (int idx = 0; idx < COUNT; ++idx) {
    assert(recv_buf[idx] == idx * COUNT + 1 + other);
  }


/**************************************************************************
 * Communication with local memory --> Hangs
 **************************************************************************/
  printf("rank %d reads memory from rank %d\n", myrank, myrank);
  MPI_Get(recv_buf, COUNT, MPI_INT, myrank, w1_disps[myrank], COUNT, MPI_INT, w1);

  //verify values
  for (int idx = 0; idx < COUNT; ++idx) {
    assert(recv_buf[idx] == idx * COUNT + 1 + myrank);
  }

  MPI_Win_unlock_all(w1);

  MPI_Barrier(MPI_COMM_WORLD);

  deallocate_window(w1_mem, &w1);

  free(w1_disps);

  printf("Test finished\n");

  MPI_Finalize();
  return EXIT_SUCCESS;
}


static int allocate_window(void * baseptr, size_t nbytes, MPI_Win *win)
{
  assert(MPI_Win_create_dynamic(MPI_INFO_NULL, MPI_COMM_WORLD, win) == MPI_SUCCESS);
  assert(MPI_Alloc_mem(nbytes, MPI_INFO_NULL, baseptr) == MPI_SUCCESS);
  return MPI_Win_attach(*win, baseptr, nbytes);
}

static int deallocate_window(void * baseptr, MPI_Win *win) {
  MPI_Win_detach(*win, baseptr);
  MPI_Win_free(win);
  MPI_Free_mem(baseptr);
}