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

/*****************************************************************************

 * Function: bcast_scatter_LR_allgather

 * Return: int

 * Inputs:
    buff: send input buffer
    count: number of elements to send
    data_type: data type of elements being sent
    root: source of data
    comm: communicator

 * Descrp: broadcasts using a scatter followed by LR allgather.

 * Auther: MPICH / modified by Ahmad Faraj

 ****************************************************************************/
int
bcast_scatter_LR_allgather(void * buff, int count, MPI_Datatype data_type,
			   int root, MPI_Comm comm)
{
  MPI_Aint extent;
  MPI_Status status;
  int i, j, k, src, dst, rank, num_procs;
  int mask, relative_rank, curr_size, recv_size, send_size, nbytes;
  int scatter_size, left, right, next_src, * recv_counts, * disps;
  int tag = 1, success = 0, failure = 1;

  MPI_Comm_rank(comm, &rank);
  MPI_Comm_size(comm, &num_procs);
  MPI_Type_extent(data_type, &extent);

  
  nbytes = extent * count;
  scatter_size = (nbytes + num_procs - 1) / num_procs; // ceiling division 
  curr_size = (rank == root) ? nbytes : 0; // root starts with all the data
  relative_rank = (rank >= root) ? rank - root : rank - root + num_procs;

  mask = 0x1;
  while (mask < num_procs) 
    {
      if (relative_rank & mask) 
	{
	  src = rank - mask; 
	  if (src < 0) src += num_procs;
	  recv_size = nbytes - relative_rank * scatter_size;
	  //  recv_size is larger than what might actually be sent by the
	  //  sender. We don't need compute the exact value because MPI
	  //  allows you to post a larger recv.
	  if (recv_size <= 0) 
	    curr_size = 0; // this process doesn't receive any data
	                   // because of uneven division 
	  else 
	    {
	      MPIC_Recv(buff + relative_rank * scatter_size, recv_size,
			MPI_BYTE, src, tag, comm, &status);
	      MPI_Get_count(&status, MPI_BYTE, &curr_size);
	    }
	  break;
	}
      mask <<= 1;
    }
  
  // This process is responsible for all processes that have bits
  // set from the LSB upto (but not including) mask.  Because of
  // the "not including", we start by shifting mask back down
  // one.

  mask >>= 1;
  while (mask > 0) 
    {
      if (relative_rank + mask < num_procs) 
	{
	  send_size = curr_size - scatter_size * mask; 
	  // mask is also the size of this process's subtree 
	  
	  if (send_size > 0) 
	    {
	      dst = rank + mask;
	      if (dst >= num_procs) dst -= num_procs;
	      MPIC_Send (buff + scatter_size * (relative_rank + mask),
			 send_size, MPI_BYTE, dst, tag, comm);

	      curr_size -= send_size;
	    }
	}
      mask >>= 1;
    }

  // done scatter now do allgather
  recv_counts = (int *) malloc(sizeof(int) * num_procs);
  if (!recv_counts)
    {
      printf("bcast-scatter-LR-allgather:95: cannot allocate memory\n");
      MPI_Finalize();
      exit(failure);
    }

  disps = (int *) malloc(sizeof(int) * num_procs);  
  if (!disps)
    {
      printf("bcast-scatter-LR-allgather:103: cannot allocate memory\n");
      MPI_Finalize();
      exit(failure);
    }
  
  for (i = 0; i < num_procs; i++)
    {
      recv_counts[i] = nbytes - i * scatter_size;
      if (recv_counts[i] > scatter_size)
        recv_counts[i] = scatter_size;
      if (recv_counts[i] < 0)
        recv_counts[i] = 0;
    }
  
  disps[0] = 0;
  for (i = 1; i < num_procs; i++)
    disps[i] = disps[i - 1] + recv_counts[i - 1];
  
  left  = (num_procs + rank - 1) % num_procs;
  right = (rank + 1) % num_procs;
          
  src = rank;
  next_src = left;

  for (i = 1; i < num_procs; i++)
    {
      MPIC_Sendrecv(buff + disps[(src - root + num_procs) % num_procs],
		    recv_counts[(src - root + num_procs) % num_procs],
		    MPI_BYTE, right, tag,
		    buff + disps[(next_src - root + num_procs) % num_procs],
		    recv_counts[(next_src - root + num_procs) % num_procs],
		    MPI_BYTE, left, tag, comm, &status);
      src = next_src;
      next_src = (num_procs + next_src - 1) % num_procs;
    }
  
  
  free(recv_counts);
  free(disps);

  return success;
}