#include "mpi.h"
#include <stdio.h>

typedef struct _SHMEM{
  MPI_Win shmem_win;
  MPI_Win lock_win;
  MPI_Datatype dtype;
  int dtype_size;
  int dim;
} *SHMEM;

void MPE_Mutex_create();
void MPE_Mutex_lock();
void MPE_Mutex_unlock();
void MPE_Mutex_free();

static int MPE_COUNTER_KEYVAL;
static double *shmem_mem_ptr;
static MPI_Win dbl_lock;

int shmem_mem_alloc_(MPI_Fint *Fcomm, MPI_Fint *Fdtype, MPI_Fint *Fmem )
{

  int size, chunk2, sizeoftype;
  MPI_Aint local_size;
  MPI_Datatype dtype;
  MPI_Comm comm;
  extern double *shmem_mem_ptr;

  dtype = MPI_Type_f2c(*Fdtype);
  comm=MPI_Comm_f2c(*Fcomm);

  MPI_Type_size(dtype, &sizeoftype);
  local_size = *Fmem * sizeoftype;

  /*  printf(" shmem_mem allocating %d bytes \n",local_size);*/
  MPI_Alloc_mem(local_size, MPI_INFO_NULL, &shmem_mem_ptr);
  /*  printf(" allocated  memory %d \n",shmem_mem_ptr);*/

  return (MPI_Fint) shmem_mem_ptr ;
}

int shmem_mem_free_()
{
  extern double *shmem_mem_ptr;
  MPI_Free_mem(shmem_mem_ptr);
  return 0;
}

int shmem_create_(MPI_Fint *Fcomm, MPI_Fint *Fdim, MPI_Fint *Fdtype, SHMEM *shmem)
{
  SHMEM new_shmem;
  int size, chunk2, sizeoftype, sizeofint;
  int dim,dim2,i,iam;
  MPI_Aint local_size;
  MPI_Datatype dtype;
  MPI_Fint dtype2;
  MPI_Comm comm;
  int *shmem_win_ptr;

  dim=*Fdim;
  dtype = MPI_Type_f2c(*Fdtype);
  comm=MPI_Comm_f2c(*Fcomm);

  new_shmem = (SHMEM)malloc (sizeof(struct _SHMEM));
  if (!new_shmem) return 1;

  MPI_Comm_size(comm, &size);
  MPI_Type_size(dtype, &sizeoftype);
  local_size = dim * sizeoftype;

  MPI_Alloc_mem(local_size, MPI_INFO_NULL, &shmem_win_ptr);

  MPI_Type_size(MPI_INT, &sizeofint);
  dim2=dim*(sizeoftype/sizeofint);
  for (i=0;i<dim2;i++)shmem_win_ptr[i] = (int) 0 ;

  MPI_Win_create(shmem_win_ptr, local_size, sizeoftype,
		 MPI_INFO_NULL, comm, &new_shmem->shmem_win);
  MPE_Mutex_create(comm, &new_shmem->lock_win);

  new_shmem->dtype=dtype;
  new_shmem->dtype_size=sizeoftype;
  new_shmem->dim=dim;
  *shmem=new_shmem;

  MPI_Comm_rank(MPI_COMM_WORLD, &iam);
  /*  printf("shmem_creat iam %d lock %d \n",iam,new_shmem->lock_win);*/
  return 0;
}
int shmem_free_(int *fshmem)
{
  int flag;
  void *shmem_win_ptr;
  SHMEM shmem;

  shmem = (SHMEM) *fshmem;
  MPI_Win_get_attr(shmem->shmem_win, MPI_WIN_BASE, &shmem_win_ptr,&flag);
  MPI_Win_free(&shmem->shmem_win);
  MPE_Mutex_free(shmem->lock_win);
  MPI_Free_mem(shmem_win_ptr);
  free(shmem);
  return 0;
}

void dlb_create_(MPI_Fint *Fcomm)
{
  MPI_Comm comm;

  comm=MPI_Comm_f2c(*Fcomm);
  MPE_Mutex_create(comm, &dbl_lock);

  return;
}

void dlb_nxtval_(MPI_Fint *value)
{
  int ierr;

  ierr=MPE_Counter_inc(dbl_lock, 0, value, 1);

}

void dlb_free_(MPI_Fint *Fcomm)
{

  MPE_Mutex_free(dbl_lock);

  return;
}

void MPE_Mutex_free( MPI_Win counter_win )
{
  extern int MPE_COUNTER_KEYVAL;
  MPI_Aint *myval_p;
  int flag; 
  void *counter_mem;

  MPI_Win_get_attr(counter_win, MPE_COUNTER_KEYVAL, &myval_p, &flag);
  MPI_Win_get_attr(counter_win, MPI_WIN_BASE, &counter_mem, &flag);

  MPI_Win_free(&counter_win);
  MPI_Free_mem(counter_mem);
  free (myval_p);
  return;
}

int shmem_dacc_(int *fshmem, int *fadr, int *flen, int *fnde, double *buf)
{
  int adr,len,nde;
  int iam,localsize,i;
  double *tmp_buf;
  MPI_Aint disp;
  SHMEM shmem;

  shmem = (SHMEM) *fshmem;
  adr=*fadr;
  len=*flen;
  nde=*fnde;

  MPI_Comm_rank(MPI_COMM_WORLD, &iam);
  /*  printf("shmem_acc1 iam %d len %d nde %d adr %d \n",iam,len,nde,adr);
      printf("shmem_acc2 iam %d lock %d \n",iam,shmem->lock_win);*/

  MPE_Mutex_lock(shmem->lock_win, nde);
  MPI_Win_lock(MPI_LOCK_SHARED, nde, MPI_MODE_NOCHECK, shmem-> shmem_win);
  /*  localsize = shmem->dtype_size *  len;
  tmp_buf = (double *)malloc (localsize);
  MPI_Get(tmp_buf,len,shmem->dtype,nde,adr,len,shmem->dtype,shmem->shmem_win);
  for (i=0;i<len;i++) tmp_buf[i]=buf[i]+tmp_buf[i];
  MPI_Put(tmp_buf,len,shmem->dtype,nde,adr,len,shmem->dtype,shmem->shmem_win);
  free(tmp_buf);*/
 
  MPI_Accumulate(buf, len, shmem->dtype,nde, adr, len, shmem->dtype,
                     MPI_SUM, shmem->shmem_win);

  MPI_Win_unlock(nde, shmem->shmem_win);
  MPE_Mutex_unlock(shmem->lock_win, nde);

  return 0;
}
int shmem_put_(int *fshmem, int *fadr, int *flen, int *fnde, void *buf)
{
  int adr,len,nde;
  int iam;
  MPI_Aint disp;
  SHMEM shmem;

  shmem = (SHMEM) *fshmem;
  adr=*fadr;
  len=*flen;
  nde=*fnde;

  MPI_Comm_rank(MPI_COMM_WORLD, &iam);
  /*  printf("shmem_put iam %d len %d nde %d adr %d \n",iam,len,nde,adr);*/

  MPE_Mutex_lock(shmem->lock_win, nde);
  MPI_Win_lock(MPI_LOCK_SHARED, nde, MPI_MODE_NOCHECK, shmem-> shmem_win);
  MPI_Put(buf,len,shmem->dtype,nde,adr,len,shmem->dtype,shmem->shmem_win);
  MPI_Win_unlock(nde, shmem->shmem_win);
  MPE_Mutex_unlock(shmem->lock_win, nde);

  return 0;
}

int shmem_get_(int *fshmem, int *fadr, int *flen, int *fnde, void *buf)
{
  int adr,len,nde;
  int iam;
  MPI_Aint disp;
  SHMEM shmem;

  shmem = (SHMEM) *fshmem;
  adr=*fadr;
  len=*flen;
  nde=*fnde;

  MPI_Comm_rank(MPI_COMM_WORLD, &iam);
  /*  printf("shmem_get iam %d len %d nde %d adr %d \n",iam,len,nde,adr);*/

  MPE_Mutex_lock(shmem->lock_win, nde);
  MPI_Win_lock(MPI_LOCK_SHARED, nde, MPI_MODE_NOCHECK, shmem-> shmem_win);
  MPI_Get(buf,len,shmem->dtype,nde,adr,len,shmem->dtype,shmem->shmem_win);
  MPI_Win_unlock(nde, shmem->shmem_win);
  MPE_Mutex_unlock(shmem->lock_win, nde);

  return 0;
}


void MPE_Mutex_create(MPI_Comm old_comm, MPI_Win *counter_win)
{
  extern int MPE_COUNTER_KEYVAL;
  int size, rank, *counter_mem, i, *myval_p;

  MPI_Comm_rank(old_comm, &rank);
  MPI_Comm_size(old_comm, &size);

  MPI_Alloc_mem(size*sizeof(int), MPI_INFO_NULL, &counter_mem);
  for (i=0; i<size; i++) counter_mem[i] = 0;
  MPI_Win_create(counter_mem, size*sizeof(int), sizeof(int), 
                 MPI_INFO_NULL, old_comm, counter_win);

  if (MPE_COUNTER_KEYVAL == MPI_KEYVAL_INVALID){
    MPI_Win_create_keyval(MPI_WIN_NULL_COPY_FN,
                          MPI_WIN_NULL_DELETE_FN,
			  &MPE_COUNTER_KEYVAL, NULL);
  }
  myval_p = (int *)malloc(size*sizeof(int));
  for (i=0; i<size; i++) myval_p[i] = 0;
  MPI_Win_set_attr( *counter_win, MPE_COUNTER_KEYVAL, myval_p);
}

int MPE_Counter_inc(MPI_Win counter_win, int nde, int *value, int increment)
{
  extern int MPE_COUNTER_KEYVAL;
  MPI_Group group;
  int rank, size, myval, flag, i, *val;
  MPI_Aint *myval_p;

  MPI_Win_get_group(counter_win, &group);
  MPI_Group_rank(group, &rank);
  MPI_Group_size(group, &size);
  MPI_Group_free(&group);
  MPI_Win_get_attr(counter_win, MPE_COUNTER_KEYVAL, &myval_p, &flag);
  myval = myval_p[nde];

  val = (int *)malloc(size*sizeof(int));
  MPI_Win_lock(MPI_LOCK_EXCLUSIVE,nde,0,counter_win);
  for (i=0; i<size; i++){
    if (i==rank)
      MPI_Accumulate(&increment, 1, MPI_INT, nde, i, 1, MPI_INT, 
                     MPI_SUM, counter_win);
    else
      MPI_Get(&val[i], 1, MPI_INT, nde, i, 1, MPI_INT, counter_win);
  }
  MPI_Win_unlock(nde,counter_win);

  myval_p[nde] = myval_p[nde]+increment;

  *value = myval;
  for (i=0; i<size; i++){
    if (i != rank) *value = *value + val[i];
  }
  free(val);
  return 0;
}

void MPE_Mutex_lock(MPI_Win win, int nde)
{
  int value, ierr;
  ierr=MPE_Counter_inc(win, nde, &value, 1);
  while (value != 0){
    ierr=MPE_Counter_inc(win, nde, &value, -1);
    ierr=MPE_Counter_inc(win, nde, &value, 1);
  }
}
void MPE_Mutex_unlock(MPI_Win win, int nde)
{
  int value, ierr;
  ierr=MPE_Counter_inc(win, nde, &value, -1);
}