
          /*************** jacobi.C *************
	   *   Mostly coded by Robert Brunner   *
	   *   Modified for plate demo project  *
	   *   by Rui Liu, February 1999        *
	   **************************************/

#include <stdlib.h>
#include <charm++.h>
#ifdef NOWEB
#include <HeapCentLB.h>
#include <GreedyRefLB.h>
#include <WSLB.h>
#include <RecBisectBfLB.h>
#else
#include "JacobiLB.h"
#endif

#include "jacobi.h"

// #include "../Strategy/strategy.h"
// #include "../Lbdb/lb.h"

enum {original_cell_size = 66};
//enum {array_dim = 16};
enum {array_dim = 8};

CkGroupID bocnum;
CkGroupID strategyID;
CkArrayID arrayGroup;
int noBCCells;

#define VAL(a,i,j) (*((a) + (i) * cell_sz + (j)))  // i--row, j--column
// #define CELL(i,j) ((i)*array_dim + (j))
#define VAL2(a,i,j) (*((a) + (i) * (2 * cell_sz -2)+ (j)))  // i-row, --column
#define CELL(i,j) ((j)*array_dim + (i))            // i--x,column, j--y,row


class RedMsg : public CMessage_RedMsg
{
  public :
  	double val;
};


class DestinationMsg : public CMessage_DestinationMsg
{ public:
  int dest;
};

class Cell;

class Reduce : public Group
{
  private :
	int next;
	int count;
	int numExpected;
	double max;
	Cell *ids[50];
  public :
	Reduce(RedMsg *);
  	Reduce(CkMigrateMessage *m) {}
	void Deposit(double val, Cell *id, int expected);
	void Recv_bcast(RedMsg *m);
	void Recv_data(RedMsg *m);
};


class NeighborMsg : public CMessage_NeighborMsg
{
public:
  int which_neighbor;
  double data[original_cell_size];
};

class RatioMsg : public CMessage_RatioMsg
{
public:
  int refineRatio;
};

class GoMsg : public CMessage_GoMsg
{
public:
  int go;
};

class main : public Chare
{
public:
  main(CkMigrateMessage *m) {}
  main(CkArgMsg *args)
  {
    RedMsg *msg= new RedMsg;
    bocnum = CProxy_Reduce::ckNew(msg);
    arrayGroup = CProxy_Cell::ckNew(array_dim*array_dim);
    CkPrintf("\nafter creating array...\n");
    CProxy_initializer::ckNew();

//     // Create Strategy
//     StrCreateMsg *strMsg = new StrCreateMsg;
//     strMsg->funcIdx = CProxy_main::ckIdx_resume((DummyMsg1 *) 0);
//     strMsg->mainhandle = thishandle;

#ifdef NOWEB
    CreateHeapCentLB();
    //    CreateWSLB();
    //    CreateGreedyRefLB();
    //    CreateRecBisectBfLB();
#else
    CreateJacobiLB();
#endif

    noBCCells = rand() / 32767.0 * 5 + 5;
    CkPrintf("No of BC Cells: %d\n", noBCCells);
    int *bcCells = (int *) malloc(noBCCells * sizeof(int));
    CkPrintf("BC Cells:\n");
    for (int i=0;i<noBCCells; i++) {
      bcCells[i] = rand() / 32767.0 * (array_dim*array_dim-1);
//      CkPrintf("(%d)[%d, %d] ", bcCells[i],
//	       bcCells[i]%array_dim, bcCells[i]/array_dim);
    }
    CkPrintf("\n");
  }
};

class Cell : public ArrayElement1D
{
private:
  double *old_data;
  double *new_data;
  int my_x, my_y;
  int num_neighbors;
  int neighbors_reported;
  int run_until;
  int not_going;

  int refineAgent;
  int currentRefineRatio;
  int nextRefineRatio;
  // refineRatio = 1 means that no refinement,
  // refineRatio = 2 means that double both the original x and y dimensions,
  // refineRatio = 4 means that four times both the original x and y dimensions

  int cell_sz, run_done;
  int noBCCells;
  int *bcCells;

  double last_time;
//   // for load balance
//   int syncH;
//   int cltH;
//   LBDBHandle lbdb;
//   ObjHandle objHandle;

public:
  Cell(void)
  {
    //CkPrintf("Hello from PE %d index %d\n",CkMyPe(),thisIndex);
    cell_sz = original_cell_size;   /*  each cell contains cell_sz x cell_sz data points */
    run_done = RUN_UNTIL;  /*  the program will run until t = run_done timesteps*/

    refineAgent = 0;
    currentRefineRatio = 1;
    nextRefineRatio = 1;

    my_x = thisIndex % array_dim;
    my_y = thisIndex / array_dim;
    num_neighbors = 0;
    if (my_x > 0)
      num_neighbors++;
    if (my_y < array_dim-1)
      num_neighbors++;
    if (my_x < array_dim-1)
      num_neighbors++;
    if (my_y > 0)
      num_neighbors++;

    if ( (0==my_x) && (0==my_y) )
      refineAgent = 1;

    old_data = new double[cell_sz*cell_sz];
    new_data = new double [cell_sz*cell_sz];

    int i,j;
    for(i=0; i < cell_sz; i++)
      for(j=0; j < cell_sz; j++) {
	VAL(old_data,i,j) = TEMP_INITIAL;
	// VAL(new_data,i,j) = TEMP_INITIAL;
      }

    if ( ((0 == my_x) && (0 == my_y)) 
	 || ((array_dim-1 == my_x) && (array_dim-1 == my_y)) )
    {
      for (i=0;i< cell_sz;i++)
	for(j=0; j < cell_sz; j++)
	VAL(old_data,i,j) = VAL(new_data,i,j) = TEMP_SOURCE;
    }

    neighbors_reported = 0;
    run_until = 0;

#ifndef NOWEB
    CWebPlateRegisterCell();
#endif

    last_time = CmiWallTimer();
    not_going = 1;
    go_nogo();

    usesAtSync=CmiTrue;
  }

  Cell(CkMigrateMessage *msg)
  { 
	last_time = CmiWallTimer();
  }

  ~Cell() {delete [] old_data; delete [] new_data;}

  void pup(PUP::er &p)
  {
	ArrayElement1D::pup(p);//Pack superclass
	p(my_x);p(my_y);
	p(num_neighbors);
	p(neighbors_reported);
	p(run_until);p(not_going);
	p(refineAgent);
	p(currentRefineRatio);
	p(nextRefineRatio);
    p(cell_sz);
    p(run_done);

	if (p.isUnpacking()) {
		old_data = new double[cell_sz*cell_sz];
		new_data = new double [cell_sz*cell_sz];
	}

	p(old_data,cell_sz*cell_sz);
	p(new_data,cell_sz*cell_sz);

    // register this cell
#ifndef NOWEB
	if (p.isUnpacking()) CWebPlateRegisterCell();
    if (p.isPacking())  CWebPlateUnregisterCell();
#endif
  }

  void neighbor_data(NeighborMsg *msg)
  {
    int neighbor_side = (msg->which_neighbor + 2) % 4;
    int i,j;
    
//CkPrintf("neighbor_data: %d not_going:%d \n", run_until, not_going);
    if (not_going)
      go_nogo();

    //    LBDBStartTimer(lbdb, objHandle);
    
    if (neighbor_side == 0)
      for(i=1;i < original_cell_size-1; i++)
	for (int k=0;k<currentRefineRatio;k++)
	  VAL(old_data,0,currentRefineRatio*i-k) = msg->data[i-1];
    //	VAL(old_data,0,i) = msg->data[i];
    else if (neighbor_side == 1)
      for(i=1;i < original_cell_size-1; i++)
	for (int k=0;k<currentRefineRatio;k++)
	  VAL(old_data,currentRefineRatio*i-k,cell_sz-1) = msg->data[i-1];
    //	VAL(old_data,i,original_cell_size-1) = msg->data[i];
    else if (neighbor_side == 2)
      for(i=1;i < original_cell_size-1; i++)
	for (int k=0;k<currentRefineRatio;k++)
	  VAL(old_data,cell_sz-1,currentRefineRatio*i-k) = msg->data[i-1];
    //	VAL(old_data,original_cell_size-1,i) = msg->data[i];
    else 
      for(i=1;i < original_cell_size-1; i++)
	for (int k=0;k<currentRefineRatio;k++)
	  VAL(old_data,currentRefineRatio*i-k,0) = msg->data[i-1];
    //	VAL(old_data,i,0) = msg->data[i];
    neighbors_reported++;
    
    /*
    CkPrintf(
      "Index %d,%d received from neighbor %d, reported=%d expected = %d\n",
      my_x,my_y,neighbor_side,neighbors_reported,num_neighbors);
    */

    delete msg;

    if (neighbors_reported == num_neighbors) {

    int xmin, xmax, ymin, ymax;

    if (1==refineAgent) {
      int tmp[5];
#ifdef NOWEB
      if (run_until == REFINE_AT) {
	xmin = array_dim / 4;
	xmin = (array_dim*3) / 4;
	ymin = array_dim / 4;
	ymax = (array_dim * 3) / 4;
	CkPrintf("xmin,ymin,xmax,ymax=%d,%d,%d,%d.\n",xmin,ymin,xmax,ymax);
      
	for (i=xmin; i<=xmax; i++)
	  for (j=ymin; j<=ymax; j++) {
	    RatioMsg *msg = new RatioMsg;
	    msg->refineRatio = 1;
	    int send_to = CELL(i,j);
	    //CkPrintf("Index %d,%d sending to %d\n",my_x,my_y,send_to);
            CProxy_Cell cr(thisAID);
            cr[send_to].setRefineRatio(msg);
	  }	
      }
#else
      if (0 != CWebPlateCheckRefine((int *)tmp)) {
	xmin = *tmp;
	ymin = *(tmp+1);
	xmax = *(tmp+2);
	ymax = *(tmp+3);
	CkPrintf("xmin,ymin,xmax,ymax=%d,%d,%d,%d.\n",xmin,ymin,xmax,ymax);
      
	for (i=xmin; i<=xmax; i++)
	  for (j=ymin; j<=ymax; j++) {
	    RatioMsg *msg = new RatioMsg;
	    msg->refineRatio = 1;
	    int send_to = CELL(i,j);
	    //CkPrintf("Index %d,%d sending to %d\n",my_x,my_y,send_to);
            CProxy_Cell cr(thisArrayID);
            cr[send_to].setRefineRatio(msg);
	  }	
      }
#endif
    }

    /*
    if (CkMyPe() == 5) {
      if ((run_until >= 2*LDB_TIMESTEP-3)  && (run_until <=2*LDB_TIMESTEP-2)) {
	CkPrintf("[%d] step %d refining index %d\n",
		 CkMyPe(),run_until,thisIndex);
	nextRefineRatio = 2;
      }
    }
    */

      if (nextRefineRatio != currentRefineRatio)
        refine();


      int i,j;
      for(i=1; i < cell_sz-1; i++)
	for(j=1; j < cell_sz-1; j++)
	  VAL(new_data,i,j) = 0.2*(VAL(old_data,i,j+1) 
	    + VAL(old_data,i+1,j) + VAL(old_data,i,j-1) 
	    + VAL(old_data,i-1,j) + VAL(old_data,i,j));

      boundary_condition();
      //      LBDBStopTimer(lbdb, objHandle);

      //      CkPrintf("neighbor_data(): my_x, my_y: %d, %d.\n", my_x, my_y);

      double *tmp_data = old_data;
      old_data = new_data;
      new_data = tmp_data;
      neighbors_reported = 0;

      if (thisIndex==0 && run_until % PRINT_TIME == 0) {
	double now = CmiWallTimer();
	CkPrintf("TIME %d %f %f\n",run_until,now,(now-last_time)/PRINT_TIME);
	last_time = now;
      }
    // load balance
    if (0 == run_until % LDB_TIMESTEP) {
      //      AtPESync(cltH, run_until);
// CkPrintf("AtSync at %d step. \n", run_until);
      not_going = 1;
      AtSync();
      return;
    }
    else
      go_nogo();
    }  
  }

  void migrateNow(DestinationMsg *msg)
  {
    migrateMe(msg->dest);
    delete msg;
  }


  void refine()
  {
    double *old_data_tmp, *new_data_tmp;
    // later on change cell_sz according to currentRefineRatio and nextRefineRatio
    int new_cell_sz;
    new_cell_sz = cell_sz * 2 - 2;
    old_data_tmp = old_data;
    new_data_tmp = new_data;
    old_data = new double [new_cell_sz*new_cell_sz];
    new_data = new double [new_cell_sz*new_cell_sz];
    // copy
    for (int i=1;i<=cell_sz-2;i++)
      for (int j=1;j<cell_sz-2;j++)
	VAL2(old_data,2*i-1,2*j-1) = VAL2(old_data, 2*i-1, 2*j)
	  = VAL2(old_data, 2*i, 2*j-1) = VAL2(old_data, 2*i, 2*j)
	  = VAL(old_data_tmp, i, j);
    cell_sz = new_cell_sz;
    delete [] old_data_tmp;
    delete [] new_data_tmp;
 
    currentRefineRatio = nextRefineRatio;

  }

  void boundary_condition()
  {
    int i,j;

    //    if  ((0 == my_x % 2)) {
    //      for (i=0;i< cell_sz;i++)
    //	for(j=0; j < cell_sz; j++)
    //	  VAL(new_data,i,j) = TEMP_SOURCE;
    //    }

    if ((thisIndex % 23) == 0) {
      for (i=0;i< cell_sz;i++)
	for(j=0; j < cell_sz; j++)
    	  VAL(new_data,i,j) = TEMP_SOURCE;
    }
  }

  void go_nogo()
  {
    not_going = 0;
    /*
    do the following things:
    (1) check if we reach the final time step. If so then stop.
    (2) check if it's time to deposit data for display. If so then deposit.
    (3) send data to four neighbors--call neighbor_data().
    */


    run_until++;

//    CkPrintf("go_nogo(): Index %d,%d count %d\n",my_x,my_y,run_until);

    if (run_until > run_done) {
      CkExit();
      return;
    }

    //    LBDBStartTimer(lbdb, objHandle);

    if (0 == run_until % TIMESTEP_JUMP) {
      //      int data2_for_display[cell_sz][cell_sz];
      int **data2_for_display;
      int i;
data2_for_display = (int **) malloc(DISPLAY_CELL_SIZE * sizeof(int *));
      for (i = 0; i < DISPLAY_CELL_SIZE; i++)
	data2_for_display[i] = (int *) malloc(DISPLAY_CELL_SIZE * sizeof(int));

      double value, sum;
      // int ratioSquare = currentRefineRatio * currentRefineRatio;

      int shrink_ratio = (cell_sz - 2) / DISPLAY_CELL_SIZE;
      int ratioSquare = shrink_ratio * shrink_ratio;
      
      for(i=0; i < DISPLAY_CELL_SIZE; i++)
        for(int j=0; j < DISPLAY_CELL_SIZE; j++) {	  
	  sum = 0;
	  for (int k=0;k<shrink_ratio;k++)
	    for (int l=0;l<shrink_ratio;l++)
	      sum += VAL(old_data, shrink_ratio*(i+1)-k, shrink_ratio*(j+1)-l);
	  value = sum / ratioSquare;
	  //	  value = VAL(new_data,i+1,j+1);
   
	  data2_for_display[i][j] = (int)
	     ( (value-TEMP_INITIAL) / TEMP_RANGE * (COLOR_RANGE - 1) );

	}
      // convert to integers representing colors

#ifndef NOWEB
      //      CkPrintf("Before CWebPlateDataDeposit().\n");
      CWebPlateDataDeposit(run_until, my_x, my_y, DISPLAY_CELL_SIZE,
			   DISPLAY_CELL_SIZE, data2_for_display);
#endif:
      
      //      CkPrintf("After CWebPlateDataDeposit().\n");

      for (i = 0; i < DISPLAY_CELL_SIZE; i++)
	free(data2_for_display[i]);
      free(data2_for_display);

    }

    double sum;

    if (my_x > 0) {
      NeighborMsg *msg = new NeighborMsg;
      msg->which_neighbor = 0;

      /*
      for(int i=1; i < cell_sz-1; i++)
	msg->data[i-1] = VAL(old_data,1,i);
	*/
      
      for(int i=1; i < original_cell_size-1; i++) {
	sum = 0;
	for (int k=0; k<currentRefineRatio; k++)
	  sum += VAL(old_data,1,currentRefineRatio*i-k);
	msg->data[i-1] = sum / currentRefineRatio;
      }
      
      int send_to = CELL(my_x-1,my_y);
      //CkPrintf("Index %d,%d sending to %d\n",my_x,my_y,send_to);
      CProxy_Cell cr(thisArrayID);
      cr[send_to].neighbor_data(msg);
    }

    if (my_y < array_dim-1) {
      NeighborMsg *msg = new NeighborMsg;
      msg->which_neighbor = 1;
      
      /*
      for(int i=1; i < cell_sz-1; i++)
	msg->data[i-1] = VAL(old_data,i,cell_sz-2);
	*/
      
      for(int i=1; i < original_cell_size-1; i++) {
        sum = 0;
	for (int k=0; k<currentRefineRatio; k++)
	  sum += VAL(old_data,currentRefineRatio*i-k,cell_sz-2);
	msg->data[i-1] = sum / currentRefineRatio;
      }
      
      int send_to = CELL(my_x,my_y+1);
      //CkPrintf("Index %d,%d sending to %d\n",my_x,my_y,send_to);
      CProxy_Cell cr(thisArrayID);
      cr[send_to].neighbor_data(msg);
    }

    if (my_x < array_dim-1) {
      NeighborMsg *msg = new NeighborMsg;
      msg->which_neighbor = 2;
      
      /*
      for(int i=1; i < cell_sz-1; i++)
	msg->data[i-1] = VAL(old_data,cell_sz-2,i);
	*/
      
      for(int i=1; i < original_cell_size-1; i++) {
	sum = 0;
	for (int k=0; k<currentRefineRatio; k++)
	  sum += VAL(old_data,cell_sz-2,currentRefineRatio*i-k);
	msg->data[i-1] = sum / currentRefineRatio;
      }
      
      int send_to = CELL(my_x+1,my_y);
      //CkPrintf("Index %d,%d sending to %d\n",my_x,my_y,send_to);
      CProxy_Cell cr(thisArrayID);
      cr[send_to].neighbor_data(msg);
    }

    if (my_y > 0) {
      NeighborMsg *msg = new NeighborMsg;
      msg->which_neighbor = 3;

      /*
      for(int i=1; i < cell_sz-1; i++)
	msg->data[i-1] = VAL(old_data,i,1);
	*/	
      
      for(int i=1; i < original_cell_size-1; i++) {
	sum = 0;
     	for (int k=0; k<currentRefineRatio; k++)
	  sum += VAL(old_data,currentRefineRatio*i-k,1);
	msg->data[i-1] = sum / currentRefineRatio;
      }
      
      int send_to = CELL(my_x,my_y-1);
      //CkPrintf("Index %d,%d sending to %d\n",my_x,my_y,send_to);
      CProxy_Cell cr(thisArrayID);
      cr[send_to].neighbor_data(msg);
    }
    //    LBDBStopTimer(lbdb, objHandle);
  }

  void setRefineRatio(RatioMsg *m)
  {
    //    refineRatio = m->refineRatio;
    nextRefineRatio = currentRefineRatio * 2;
    CkPrintf("setRefineRatio: [%d, %d]: nextRR, curRR = %d, %d.\n",
	     my_x, my_y, nextRefineRatio, currentRefineRatio);
    delete m;
  }

  void ResumeFromSync(void) {
    CProxy_Cell p(thisArrayID);
    p[thisIndex].resume();
  }

  void resume(void)
  {
    if (not_going)
      go_nogo();
  }
};


Reduce :: Reduce(RedMsg *)
{
  count=0;
  next=0;
  numExpected=0;
  max=0.0;
}

void Reduce ::  Deposit(double val, Cell *id, int expected)
{
  //CkPrintf("Deposit called in %d numExpected =%d \n", CkMyPe(), expected);
  numExpected=expected + CmiNumSpanTreeChildren(CkMyPe());
  ids[next++]=id;
  if (next >=50) {
	//CkPrintf("ids OVERFLOW!!!!!!!!!!!!!!\n"); 
  }
  if (val > max) max=val;
  count++;
  //CkPrintf("Deposit in %d count = %d num=%d\n", CkMyPe(), count, numExpected);
  if (count == numExpected) {
        count=0;
	RedMsg *msg=(RedMsg *) new RedMsg;
	msg->val=max;
	max=0.0;
	if (CkMyPe()==0) {
                CProxy_Reduce cr(thisgroup); cr.Recv_bcast(msg);
		return;
	}
	//CkPrintf("%d Sending data to parent\n", CkMyPe());
        CProxy_Reduce cr(thisgroup); 
        cr.Recv_data(msg, CmiSpanTreeParent(CkMyPe()));
  }
}

void Reduce :: Recv_data(RedMsg *m)
{
  //CkPrintf("%d received data from child\n", CkMyPe());
  if (m->val > max) max=m->val;
  count++;
  //CkPrintf("recv_data in %d count = %d num=%d\n", CkMyPe(), count, numExpected);
  if (count == numExpected) {
        count=0;
	RedMsg *msg= new RedMsg;
	msg->val=max;
	max=0.0;

	if (CkMyPe()==0) {
                CProxy_Reduce cr(thisgroup); cr.Recv_bcast(msg);
		return;
	}
	//CkPrintf("%d Sending data to parent\n", CkMyPe());
        CProxy_Reduce cr(thisgroup);
        cr.Recv_data(msg, CmiSpanTreeParent(CkMyPe()));
  }
}
	  
void Reduce ::  Recv_bcast(RedMsg *m)
{
  CkPrintf("%d received bcast from root!!!!! Wall time= %f Timer= %f\n", CkMyPe(), CmiWallTimer(), CmiTimer());
  for (int i=0;i<next;i++) {
	ids[i]->go_nogo();
  }
  max=0.0;
  next=0;
}

#include "jacobi.def.h"


// extern "C" void OMMigrate(OMID omId, void *userData, ObjID objID, int dest)
// {
//   DestinationMsg *msg = new DestinationMsg;
//   msg->dest = dest;
//   (CProxy_Cell(arrayGroup)[objID[0]]).migrateNow(msg);
// }
