#include "BlueGene.h"

#define A_SIZE 4

#define reduceID   1
#define computeMaxID  2
#define contributeID  3

void reduce(ThreadInfo *) ;
void computeMax(ThreadInfo *) ;
void contribute(ThreadInfo *) ;

class contributeMsg: public PacketMsg 
{
public:
  int max ;
} ;

class computeMaxMsg: public PacketMsg
{} ;

class reduceMsg: public PacketMsg
{} ;

typedef struct userDataStruct
{
  int data[A_SIZE] ;
  int count ;
} userData ;

void BgInit(Main *main) 
{}

void* BgNodeInit(BgNode *bgNode) 
{

  /* ckout << "Initializing node " << bgNode->thisIndex.x << "," << bgNode->thisIndex.y << "," << bgNode->thisIndex.z << endl ; */

  bgNode->registerHandler(reduceID, reduce) ;
  bgNode->registerHandler(computeMaxID, computeMax) ;
  bgNode->registerHandler(contributeID, contribute) ;

  computeMaxMsg *msg = new computeMaxMsg;
  bgNode->addMessage(msg, computeMaxID, LARGE_WORK);

  //declare node private data
  userData *ud = new userData ;
  ud->count = 0 ;
  for(int i=0; i<A_SIZE; i++)
    ud->data[i] = 0 ;

  return (void*)ud ;
}

void BgFinish() 
{}

void computeMax(ThreadInfo *info) 
{
  int A[A_SIZE][A_SIZE];
  int i, j;
  int max = 0;

  int x,y,z;
  info->bgNode->getXYZ(x,y,z);

  // Initialize
  for (i=0;i<A_SIZE;i++) 
    for (j=0;j<A_SIZE;j++) 
      A[i][j] = info->bgNode->numBgX * info->bgNode->numBgY * info->bgNode->numBgZ - x*y*z - i*j ;
      /* A[i][j] = i*j ; */

  // Find Max
  for (i=0;i<A_SIZE;i++) 
    for (j=0;j<A_SIZE;j++) 
      if (max < A[i][j]) 
      {
  max = A[i][j];
      }

  // contribute the results for reduction 
  contributeMsg *msg = new contributeMsg;
  msg->max = max;
  info->bgNode->addMessage(msg,contributeID,SMALL_WORK);

  ckout << "computeMax in " << x << ", " << y << ", " << z << endl;
  ckout << "contributed max value " << max << endl;
}

void contribute(ThreadInfo *info)
{
  int x,y,z;
  info->bgNode->getXYZ(x,y,z);

  int count = ((userData*)(info->bgNode->nvData))->count++ ;
  ((userData*)(info->bgNode->nvData))->data[count++] = ((contributeMsg*)(info->msg))->max ;

  int reqCount ;

  if(z==info->bgNode->numBgZ-1)
  {
  reqCount = 1 ;  
  }
  else if(z>0 || (z==0 && x==info->bgNode->numBgX-1))
  {
  reqCount = 2 ;
  }
  else if(x>0 || (x==0 && y==info->bgNode->numBgY-1))
  {
  reqCount = 3 ;  
  }
  else
    reqCount = 4 ;

  if(count==reqCount)  //if data for reduction is ready
  {
    reduceMsg *msg = new reduceMsg ;
    info->bgNode->addMessage(msg, reduceID, LARGE_WORK) ;
    ckout << "contribute in " << x << ", " << y << ", " << z << endl ;
  ckout << "Values collected " << count << ", calling reduction " << endl ;
  return ;
  }
  ckout << "contribute in " << x << ", " << y << ", " << z << endl ;
  ckout << "Values collected " << count << ", reqCount " << reqCount << endl ;
}

void reduce(ThreadInfo *info) 
{
  int x,y,z;
  info->bgNode->getXYZ(x,y,z);
  ckout << "reduce in " << x << ", " << y << ", " << z << endl;

  //do reduction
  int max = 0 ;
  int count = ((userData*)(info->bgNode->nvData))->count ;
  for(int i=0; i<count; i++)
  {
  if(max<((userData*)(info->bgNode->nvData))->data[i])
    max = ((userData*)(info->bgNode->nvData))->data[i] ; 
  }

  if(x==0 && y==0 && z==0)
  {
  ckout << "Exiting: max value is " << max << endl ;
  info->bgNode->finish() ;
  return ;
  }
  //send max to destination(decide) 
  if(z>0)
    z-- ;   
  else if(x>0)
    x-- ;
  else
    y-- ;

  contributeMsg *msg = new contributeMsg;
  msg->max = max;
  info->sendPacket(x,y,z,msg,contributeID,SMALL_WORK);

  ckout << "sending max value " << max << " to " << x << ", " << y << ", " << z << endl ;
}

