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

int bcast_knary_segment_size;
int knary = 2;

int bcast_knary(void *buf,int count, MPI_Datatype datatype, int root, MPI_Comm comm)
{
  int tag = 5000;
  MPI_Status status;
  MPI_Request req[10];
  int rank;
  int i,j;
  int total_node;
  int target;

  MPI_Comm_rank(MPI_COMM_WORLD,&rank);
  MPI_Comm_size(MPI_COMM_WORLD,&total_node);

  int **send_to;
  int *recv_from;
  int *sequence;
  int bcast_knary_phase;

  //MALLOC
  recv_from = (int *) malloc (total_node * sizeof (int));
  sequence = (int *) malloc (total_node * sizeof (int));
  send_to = (int **) malloc (knary * sizeof (int *));
  for (i=0;i<knary;i++)
    {
      send_to[i] = (int *) malloc (total_node * sizeof (int));
    }

  //initial everything to -1                                                                                                    
  for (i=0;i<total_node;i++)
    {
      recv_from[i] = -1;
      sequence[i] = -1;
      for (j=0;j<knary;j++)
        {
          send_to[j][i] = -1;
        }
    }

  // go through node 0 to node num_machine - 1 and compute send_to[phase][machine] and recv_from[machine]                       
  for (i=0;i<total_node;i++)
    {
      //node i send to (knary_number * i)+ k, where 0<k<=knary_number                                                           
      //for example, for binary, node i sends to node 2*i+1 and node 2*i+2                                                      
      for (j=0;j<knary;j++)
        {
          target = (knary * i) + j + 1;

          if (target < total_node)
            {
              send_to[j][i] = target;
              recv_from[target] = i;
            }
          else
            {
              send_to[j][i] = -1;
            }
        }
    }

  // compute sequence[] which is a level that node resides.                                                                     
  //node 0 is in level 0, there are (knary_number) nodes in level 1                                                             
  //(knary_number * knary_number) nodes in level 2 and so forth                                                                 
  sequence[0] = 0;
  int curr_node = 1;
  int curr_size = 1;
  int level = 1;
  while(curr_node < total_node)
    {
      curr_size *= knary;
      for (i=0;i<curr_size;i++)
        {
          sequence[curr_node] = level;
          curr_node++;
          if (curr_node >= total_node)
            {
              break;
            }
        }
      level++;
    }
  // number of phase of this knary tree is equal to a level of the last node                                                    
  bcast_knary_phase = sequence[total_node-1];

  // non-pipeline
  if (count <= bcast_knary_segment_size)
    {
      for (i=0;i<bcast_knary_phase;i++)
	{
	  // loop through all send
	  for (j=0;j<knary;j++)
	    {
	      if ((send_to[j][rank] != -1) && (sequence[rank] == i))
		MPI_Send(buf,count,datatype,send_to[j][rank],tag,MPI_COMM_WORLD);
	    }

	  // receive
	  if ((recv_from[rank] != -1) && (sequence[rank] == (i+1)))
	    MPI_Recv(buf,count,datatype,recv_from[rank],tag,MPI_COMM_WORLD,&status);
	}
      for (i=0;i<knary;i++)
	free(send_to[i]);
      free(send_to);
      free(recv_from);
      free(sequence);
      return 0;
    }
  
  // pipelining

  int pipe_length, last_count, phase;

  pipe_length = (count-1)/bcast_knary_segment_size + 1;
  last_count = count%bcast_knary_segment_size;
  phase = pipe_length+bcast_knary_phase-1;

  // last message has the same size as all the rest
  if (last_count == 0)
    {
      for (i=0;i<phase;i++)
	{
	  // send
	  for (j=0;j<knary;j++)
	    {
	      if ((send_to[j][rank] != -1) && (i >= sequence[rank]) && (i < (pipe_length + sequence[rank])))
		MPI_Send(buf+(bcast_knary_segment_size * (i-sequence[rank])),bcast_knary_segment_size,datatype,send_to[j][rank],tag,MPI_COMM_WORLD);
	    }
	  
	  // receive
	  if ((recv_from[rank] != -1) && (i >= (sequence[rank]-1)) && (i < pipe_length + sequence[rank] -1))
	    MPI_Recv(buf+(bcast_knary_segment_size*(i-sequence[rank]+1)),bcast_knary_segment_size,datatype,recv_from[rank],tag,MPI_COMM_WORLD,&status);
	}
    }
  
  // last message has smaller size than the rest
  else
    {
      count = bcast_knary_segment_size;
      for (i=0;i<phase;i++)
	{
	  // send
	  for (j=0;j<knary;j++)
	    {
	      if ((send_to[j][rank] != -1) && (i >= sequence[rank]) && (i < (pipe_length + sequence[rank])))
		{
		  if (i == (pipe_length + sequence[rank] -1))
		    count = last_count;

		  MPI_Send(buf+(bcast_knary_segment_size * (i-sequence[rank])),count,datatype,send_to[j][rank],tag,MPI_COMM_WORLD);
		}
	    }

	  // receive
	  if ((recv_from[rank] != -1) && (i >= (sequence[rank]-1)) && (i < pipe_length + sequence[rank] -1))
	    {
              if (i == (pipe_length + sequence[rank] -2))
		count = last_count;
	      MPI_Recv(buf+(bcast_knary_segment_size*(i-sequence[rank]+1)),count,datatype,recv_from[rank],tag,MPI_COMM_WORLD,&status);
	    }
	}
    }
  
  for (i=0;i<knary;i++)
    free(send_to[i]);
  free(send_to);
  free(recv_from);
  free(sequence);
  return 0;
}