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

/*
! This program shows how you might implement an MPI_Alltoallv in
! the routine Myall2allv.  Myalltoallv is fairly efficent in that
! it does the sends/recs in orders so that there is less network
! contention.  See the class notes.
!
! Each processor send/rec a different and random amount of data
! to/from other processors.  We use MPI_Alltoall to tell how much
! data is going to be sent.
! 
*/

/* module global */
int numnodes,myid,mpi_err;
#define mpi_root 0
/* end module  */

void init_it(int  *argc, char ***argv);
void seed_random(int  id);
void random_number(float *z);
int myall2allv(int *sendbuf,int *sendcounts,int *sdispls,
               MPI_Datatype MPI_S_TYPE,
               int *recvbuf,int *recvcounts,int *rdispls,
               MPI_Datatype MPI_R_TYPE,
               MPI_Comm  comm);

void init_it(int  *argc, char ***argv) {
	mpi_err = MPI_Init(argc,argv);
    mpi_err = MPI_Comm_size( MPI_COMM_WORLD, &numnodes );
    mpi_err = MPI_Comm_rank(MPI_COMM_WORLD, &myid);
}

int main(int argc,char *argv[]){
	int *sray,*rray;
	int *sdisp,*scounts,*rdisp,*rcounts;
	int ssize,rsize,i,k,j;
	float z;

	init_it(&argc,&argv);
	scounts=(int*)malloc(sizeof(int)*numnodes);
	rcounts=(int*)malloc(sizeof(int)*numnodes);
	sdisp=(int*)malloc(sizeof(int)*numnodes);
	rdisp=(int*)malloc(sizeof(int)*numnodes);
	seed_random(myid);
	for(i=0;i<numnodes;i++){
		random_number(&z);
		scounts[i]=(int)(10.0*z)+1;
	}
	printf("myid= %d scounts=",myid);
	for(i=0;i<numnodes;i++)
		printf("%d ",scounts[i]);
	printf("\n");
	mpi_err = MPI_Alltoall(	scounts,1,MPI_INT,
						    rcounts,1,MPI_INT,
	                 	    MPI_COMM_WORLD);
/*	write(*,*)"myid= ",myid," rcounts= ",rcounts */
/* calculate displacements and the size of the arrays */
	sdisp[0]=0;
	for(i=1;i<numnodes;i++){
		sdisp[i]=scounts[i-1]+sdisp[i-1];
	}
	rdisp[0]=0;
	for(i=1;i<numnodes;i++){
		rdisp[i]=rcounts[i-1]+rdisp[i-1];
	}
	ssize=0;
	rsize=0;
	for(i=0;i<numnodes;i++){
		ssize=ssize+scounts[i];
		rsize=rsize+rcounts[i];
	}
	
/* allocate send and rec arrays */
	sray=(int*)malloc(sizeof(int)*ssize);
	rray=(int*)malloc(sizeof(int)*rsize);
	for(i=0;i<ssize;i++)
		sray[i]=myid;
/* send/rec different amounts of data to/from each processor */
	mpi_err = myall2allv(	sray,scounts,sdisp,MPI_INT,
						        rray,rcounts,rdisp,MPI_INT,
	                 	        MPI_COMM_WORLD);
	                
	printf("myid= %d rray=",myid);
	for(i=0;i<rsize;i++)
		printf("%d ",rray[i]);
	printf("\n");
    mpi_err = MPI_Finalize();
}
/*
  0:myid= 0 scounts=1 7 4
  0:myid= 0 scounts=0 1 1 1 1 1 1 2
  1:myid= 1 scounts=6 2 4
  1:myid= 1 scounts=0 0 0 0 0 0 0 1 1 2 2 2 2 2 2 2
  2:myid= 2 scounts=1 7 4
  2:myid= 2 scounts=0 0 0 0 1 1 1 1 2 2 2 2
*/

void seed_random(int  id){
	srand((unsigned int)id);
}
void random_number(float *z){
	int i;
	i=rand();
	*z=(float)i/32767;
}

int ieor(int a,int b){
 return (a ^ b);
}

int myall2allv(int *sendbuf,int *sendcounts,int *sdispls,
               MPI_Datatype MPI_S_TYPE,
               int *recvbuf,int *recvcounts,int *rdispls,
               MPI_Datatype MPI_R_TYPE,
               MPI_Comm  comm) {
  int ierr;

  int myid,numnodes;
  MPI_Status status;
  int i,n2,xchng,my_tag;

  mpi_err = MPI_Comm_size( comm, &numnodes );
  mpi_err = MPI_Comm_rank(comm, &myid);
  my_tag=4925;

/*! find the power of two >= numnodes */
  n2=1;
  while(n2 < numnodes)
  	n2=n2*2;
  n2=n2-1;
  for(i=1;i<=n2;i++) {
/*! do xor to find the processor xchng */
      xchng=ieor(i,myid);
      if(xchng <= (numnodes-1)){
          if(myid < xchng){
              ierr =  MPI_Send(&sendbuf[sdispls[xchng]],sendcounts[xchng],MPI_S_TYPE,xchng,my_tag,comm);
              ierr =  MPI_Recv(&recvbuf[rdispls[xchng]],recvcounts[xchng],MPI_R_TYPE,xchng,my_tag,comm,&status);
          }
          else {
              ierr = MPI_Recv(&recvbuf[rdispls[xchng]],recvcounts[xchng],MPI_R_TYPE,xchng,my_tag,comm,&status);
              ierr = MPI_Send(&sendbuf[sdispls[xchng]],sendcounts[xchng],MPI_S_TYPE,xchng,my_tag,comm);
          }
      }
  }
/* ! exchange data with myself */
  for(i=0;i<sendcounts[myid];i++) {
        recvbuf[rdispls[myid]+i]=sendbuf[sdispls[myid]+i];
  }
  return (ierr);
}