#include "jacobi2d.decl.h"
#include <vector>
#include <assert.h>
#include <sys/time.h>
#define LDB_ITER 20
// See README for documentation

/*readonly*/ CProxy_Main mainProxy;
/*readonly*/ int block_height;
/*readonly*/ int block_width;
/*readonly*/ int array_height;
/*readonly*/ int array_width;
             int enter_once  = 1;

// specify the number of worker chares in each dimension
/*readonly*/ int num_chare_rows;
/*readonly*/ int num_chare_cols;

// We want to wrap entries around, and because mod operator % sometimes misbehaves on negative values, 
// I just wrote these simple wrappers that will make the mod work as expected. -1 maps to the highest value.
#define wrap_x(a)  (((a)+num_chare_cols)%num_chare_cols)
#define wrap_y(a)  (((a)+num_chare_rows)%num_chare_rows)

CkArrayID a;

//#define total_iterations 200
float progStartTime;

class Main : public CBase_Main
{
public:
    float *totalTime;
    int *totalObjs;
    int recieve_count;
    CProxy_Jacobi array;
    int num_chares;
    int iterations;
    int total_iterations;
    double perIterStartTime;
    double totalIterTime;
    double lbdOverhead;

    Main(CkArgMsg* m) {
	totalTime = new float[CkNumPes()];
	totalObjs = new int[CkNumPes()];
	for(int i=0;i<CkNumPes();i++) 
	{
		totalTime[i] = 0.0;
		totalObjs[i] = 0;
	}
        if (m->argc < 3) {
          CkPrintf("%s [array_size] [block_size]\n", m->argv[0]);
          CkAbort("Abort");
        }
        totalIterTime = 0.0;
        lbdOverhead = 0.0;

        // set iteration counter to zero
        iterations=0;

        // store the main proxy
        mainProxy = thisProxy;

        array_height = atoi(m->argv[1]);
	      array_width = atoi(m->argv[2]);
        block_height = atoi(m->argv[3]);
	block_width = atoi(m->argv[4]);
        if (array_width < block_width || array_width % block_width != 0)
          CkAbort("array_size % block_size != 0!");

        num_chare_rows = array_height / block_height;
        num_chare_cols = array_width / block_width;
        // print info
        //CkPrintf("Running Jacobi on %d processors with (%d,%d) elements\n", CkNumPes(), num_chare_rows, num_chare_cols);

	total_iterations = 200;
	if (m->argc > 5) {
	  total_iterations = atoi(m->argv[5]);
	}

        // Create new array of worker chares
        array = CProxy_Jacobi::ckNew(num_chare_cols, num_chare_rows);

        // save the total number of worker chares we have in this simulation
        num_chares = num_chare_rows*num_chare_cols;


        //Start the computation
        perIterStartTime = CkWallTimer();
        progStartTime = CkWallTimer();
        recieve_count = 0;
        array.begin_iteration();


    }

    // Each worker reports back to here when it completes an iteration
    void report(int row, int col,int core,float mytime) {
      recieve_count++;

      totalTime[core] += mytime;
      totalObjs[core]++;

      double perIterTime = CkWallTimer() - perIterStartTime;

      if (num_chares == recieve_count) {
        if (++iterations == total_iterations) {
          CkPrintf("Completed %d iterations; last iteration time: %.6lf\n", iterations, perIterTime);
          totalIterTime += perIterTime;
          //find maximum totalTIme
          double totalProgramTime = (CmiWallTimer()-progStartTime);
          //double timePerIter = totalProgramTime/iterations;

          // calculate the max idle time
          double maximum=0.0;
          double minimum = totalTime[0];
          int penum_min, penum_max;
          for(int i=0;i<CkNumPes();i++){
             if(totalTime[i] >= maximum){
               maximum = totalTime[i];
               penum_max = i;
             }
             if(totalTime[i] <= minimum){
               minimum = totalTime[i];
               penum_min = i;
             }
          }
          double maxIdleTime =  maximum - minimum;
          double percentageIdleTime = (maxIdleTime/totalProgramTime)*100.00;

          // calculate the average idle time
          double totalidletime=0.0;
          for(int i=0;i<CkNumPes();i++){
            totalidletime += (maximum-totalTime[i]);
          }
          double averageidletime  =  totalidletime/ CkNumPes();
          double percentageAvIdleTime = (averageidletime / totalProgramTime)*100;



          CkPrintf("PE\tTotal Time\tobjs\n");
          CkPrintf("-----------------------------------------------------\n");
          for(int i=0;i<CkNumPes();i++)
          {
            CkPrintf("%d\t%lf\t%d\n",i,totalTime[i], totalObjs[i]);
          }
          CkPrintf("\n\n");
          CkPrintf("Total Program Time:%.15lf \n",  totalProgramTime);
          CkPrintf("Total Iteration Time:%.15lf \n",  totalIterTime);
          CkPrintf("lbd Overhead Iteration Time:%.15lf \n",  lbdOverhead);
          CkPrintf("Max idle time:%.15lf \n", maxIdleTime);
          CkPrintf("Average idle time:%.15lf\n", averageidletime);
          CkPrintf("Percentage max idle time:%.15lf \n", percentageIdleTime);
          CkPrintf("Percentage average idle time:%.15lf\n", percentageAvIdleTime);
          CkPrintf("max pe (%d): min pe (%d)\n\n\n", penum_max, penum_min); 

          CkPrintf("%.15lf:%.15lf:%.15lf:%.15lf:%.15lf:%.15lf\n", 
                  totalProgramTime,totalIterTime, maxIdleTime,averageidletime, percentageIdleTime,percentageAvIdleTime   );
          CkExit();        
      } else {
        CkPrintf("starting new iteration; iteration %d time: %.6lf\n", iterations, perIterTime);
        if(iterations == LDB_ITER + 1) {
          lbdOverhead = perIterTime;
        } else {
          totalIterTime += perIterTime;
        }
        recieve_count=0;
        perIterStartTime = CkWallTimer();
        if(iterations != LDB_ITER) {
          array.begin_iteration();
        } else {
          CkPrintf("Before atsync\n");
          for(int i=0;i<CkNumPes();i++) {
            CkPrintf("%d\t%lf\t%d\n",i,totalTime[i], totalObjs[i]);
          }
        }
     }
  }
}
};

class Jacobi: public CBase_Jacobi {
public:

    //Jacobi_SDAG_CODE
    int messages_due;
    int itercnt;
    double **temperature;
    bool hasSent;
    //Timer times[20][20];

    // Constructor, initialize values
    Jacobi() {
        int i,j;
        usesAtSync = true;
          // allocate two dimensional array
        temperature = new double*[block_height+2];
        for (i=0; i<block_height+2; i++)
          temperature[i] = new double[block_width+2];
        messages_due = 4;
        for(i=0;i<block_height+2;++i){
          for(j=0;j<block_width+2;++j){
            temperature[i][j] = 0.0;
          }
        }
        itercnt=0;
        hasSent=false;
        BC();
    }
    
    
    void ResumeFromSync(){
      //ckout << "iter=" << itercnt << ": ResumeFromSync called on " << CkMyPe() << endl;
      begin_iteration();
    }

    // Enforce some boundary conditions
    void BC(){
        // Heat left and top edges of each chare's block
	    for(int i=1;i<block_height+1;++i)
        temperature[i][1] = 255.0;
      for(int j=1;j<block_width+1;++j)
        temperature[1][j] = 255.0;
    }

    void pup(PUP::er &p){
        CBase_Jacobi::pup(p);
        //__sdag_pup(p);

        p|messages_due;
        p|itercnt;
        p|hasSent;
        p|usesAtSync;
        if(p.isUnpacking()){
          temperature = new double*[block_height+2];
          for (int i=0; i<block_height+2; i++)
            temperature[i] = new double[block_width+2];
        }
        for (int i=0; i<block_height+2; i++) {
          PUParray(p, temperature[i], block_width+2);
        }
    }

    // a necessary function which we ignore now
    // if we were to use load balancing and migration
    // this function might become useful
    Jacobi(CkMigrateMessage* m) {
    }

    ~Jacobi() { 
      for (int i=0; i<block_height; i++)
        delete [] temperature[i];
      delete [] temperature; 
    }

    // Perform one iteration of work
    // The first step is to send the local state to the neighbors
    void begin_iteration(void) {
      itercnt++;
      //ckout << "iter=" << itercnt << ":(" << thisIndex.x << "," << thisIndex.y << ") is on " << CkMyPe() << endl;
      // Copy left column and right column into temporary arrays
      double *left_edge = new double[block_height];
      double *right_edge = new double[block_height];

      for(int i=0;i<block_height;++i){
        left_edge[i] = temperature[i+1][1];
        right_edge[i] = temperature[i+1][block_width];
      }

      // Send my left edge
      thisProxy(wrap_x(thisIndex.x-1), thisIndex.y).ghostsFromRight(block_height, left_edge);
      // Send my right edge
      thisProxy(wrap_x(thisIndex.x+1), thisIndex.y).ghostsFromLeft(block_height, right_edge);
      // Send my top edge
      thisProxy(thisIndex.x, wrap_y(thisIndex.y-1)).ghostsFromBottom(block_width, &temperature[1][1]);
      // Send my bottom edge
      thisProxy(thisIndex.x, wrap_y(thisIndex.y+1)).ghostsFromTop(block_width, &temperature[block_height][1]);

      hasSent=true;
      check_and_compute();
      delete [] right_edge;
      delete [] left_edge;
    }

    void ghostsFromRight(int width, double ghost_values[]) {
      for(int i=0;i<width;++i){
        temperature[i+1][block_width+1] = ghost_values[i];
      }
      messages_due--;
      check_and_compute();
    }

    void ghostsFromLeft(int width, double ghost_values[]) {
      for(int i=0;i<width;++i){
          temperature[i+1][0] = ghost_values[i];
      }
      messages_due--;
      check_and_compute();
    }

    void ghostsFromBottom(int width, double ghost_values[]) {
      for(int i=0;i<width;++i){
          temperature[block_height+1][i+1] = ghost_values[i];
      }
      messages_due--;
      check_and_compute();
    }

    void ghostsFromTop(int width, double ghost_values[]) {
      for(int i=0;i<width;++i){
          temperature[0][i+1] = ghost_values[i];
      }
      messages_due--;
      check_and_compute();
    }

    void check_and_compute() {
	    float timeToCompute = CmiWallTimer();
       if (messages_due == 0 && hasSent) {
          messages_due = 4;
          hasSent=false;
          compute();
          int currPe = CkMyPe();
          //ckout << "iter=" << itercnt << ": (" << thisIndex.x << "," << thisIndex.y << ") reporting pe " << CkMyPe() << endl;
          mainProxy.report(thisIndex.x, thisIndex.y, CkMyPe(), CmiWallTimer()-timeToCompute);
          if (itercnt == LDB_ITER) {
            //ckout << "iter=" << itercnt << ": AtSync reported from " << CkMyPe() << endl;
            AtSync();
          }
        }
    }

    // Check to see if we have received all neighbor values yet
    // If all neighbor values have been received, we update our values and proceed
    void compute() {
      // We must create a new array for these values because we don't want to update any of the
      // the values in temperature[][] array until using them first. Other schemes could be used
      // to accomplish this same problem. We just put the new values in a temporary array
      // and write them to temperature[][] after all of the new values are computed.
      double new_temperature[block_height+2][block_width+2];

      for(int i=1;i<block_height+1;++i){
        for(int j=1;j<block_width+1;++j){
          // update my value based on the surrounding values
          new_temperature[i][j] = (temperature[i-1][j]+temperature[i+1][j]+temperature[i][j-1]+temperature[i][j+1]+temperature[i][j]) / 5.0;
        }
      }

      for(int i=0;i<block_height+2;++i)
        for(int j=0;j<block_width+2;++j)
          temperature[i][j] = new_temperature[i][j];

      // Enforce the boundary conditions again
      BC();
    }




};

#include "jacobi2d.def.h"
