// ------------------------------------------------------------------------- //
// --- This is the raw MPI stereo deal-style program                     --- //
// ------------------------------------------------------------------------- //
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <mpi.h>
#include <sys/times.h>
#include "stereo.h"

#define DEBUG  0     // set to 1 for debugging mode, 0 for normal mode
#define WINX   13    // window rows
#define WINY   13    // window columns

// For eSkel
#define STAGES  3   // pipeline stages

// For display
#define CASEDEAL 1
#define CASE3ST  2
#define CASESEQ  3
#define CASEMPI  4

#define BYEBYE 1         //end of work tag
#define SPECIALBYEBYE 2  //end of work tag to the deal worker which will have to forward BYEBYEto the final stage

double  dsecsG=0, esecsG=0;
double st3secs[ITERS];  // times between completions at stage 3, 
                        // st3secs[0] must be ignored, leaving
                        // ITERS-1 interesting readings

double truethrput;


// ------------------------------------------------------------------------- //
// gendiffimg forms the difference image
void gendiffimg(int ref[(ROWS+WINY)*COLS], int m1[(ROWS+WINY)*COLS], int m2[(ROWS+WINY)*COLS],
		float diffimg[(ROWS+WINY)*COLS], int curdisp) {
  int  i,j;

  for (i=0;i<ROWS+WINY;i++)
    for (j=0;j<COLS-2*curdisp;j++)
      diffimg[i*COLS+j] = (ref[i*COLS+j] - m1[i*COLS+j+curdisp])*(ref[i*COLS+j] - m1[i*COLS+j+curdisp]) +
	              (ref[i*COLS+j] - m2[i*COLS+j+2*curdisp])*(ref[i*COLS+j] - m2[i*COLS+j+2*curdisp]);
  
}

// ------------------------------------------------------------------------- //
// updatedispimg sums the difference image over a moving window and updates
// the current best error and best disparity images.
void updatedispimg(float diffimg[(ROWS+WINY)*COLS], float curbesterr[ROWS*COLS], 
		   int curbestdisp[ROWS*COLS], int curdisp) {
  int   i,j;
  float sum,csum[COLS];

  for (j=0;j<COLS;j++) {
    csum[j] = 0.0;
    for (i=0;i<WINX;i++)
      csum[j] = csum[j] + diffimg[i*COLS+j];
  }

  for (i=0;i<ROWS;i++) {
    sum = 0.0;
    for (j=0;j<7;j++)
      sum = sum + csum[j];
    
    if (sum < curbesterr[i*COLS+0]) {
      curbesterr[i*COLS+0]  = sum;
      curbestdisp[i*COLS+0] = curdisp;
    }

    for (j=1;j<7;j++) {
      sum = sum + csum[j+6];
      if (sum < curbesterr[i*COLS+j]) {
	curbesterr[i*COLS+j]  = sum;
	curbestdisp[i*COLS+j] = curdisp;
      }
    }

    for (j=7;j<COLS-6;j++) {
      sum = sum - csum[j-7];
      sum = sum + csum[j+6];
      if (sum < curbesterr[i*COLS+j]) {
	curbesterr[i*COLS+j]  = sum;
	curbestdisp[i*COLS+j] = curdisp;
      }
    }

    for (j=COLS-6;j<COLS;j++) {
      sum = sum - csum[j-7];
      if (sum < curbesterr[i*COLS+j]) {
	curbesterr[i*COLS+j]  = sum;
	curbestdisp[i*COLS+j] = curdisp;
      }
    }
    
    for (j=0;j<COLS;j++)
      csum[j] = csum[j] - diffimg[i*COLS+j] + diffimg[(i+WINY)*COLS+j];

  }
  
}

// ------------------------------------------------------------------------- //
// error displays an error message raised in testdata
void error(int i, int j, int tardis, int disp) {
  
  printf("*** ERROR in testdata at position (%d,%d). Expected disp %d but got %d *** \n",
	 i,j,tardis,disp);
  
}

// ------------------------------------------------------------------------- //
// testdata is testing output (matched with getdata)
void testdata(int disp[ROWS*COLS]) {

  int i,j,ipos,jpos;
  int tardis,isize,jsize;

  tardis = 12;
  isize=37;
  jsize=37;

  ipos = (ROWS-isize)/2 - 1 ;
  jpos = (COLS-jsize)/2 - tardis - 1;


  for (i=ipos+10;i<=ipos+isize-11;i++)
    for (j=jpos+10;j<=jpos+jsize-11;j++)
      if (disp[i*COLS+j]!=tardis) 
	error(i,j,tardis,disp[i*COLS+j]);
  
  for (i=0;i<=ipos-10;i++)
    for (j=0;j<=jpos+jsize+2*tardis+9;j++)
      if (disp[i*COLS+j]!=0) 
	error(i,j,0,disp[i*COLS+j]);
  
  for (i=0;i<=ipos+isize+9;i++)
    for (j=jpos+jsize+2*tardis+9;j<COLS;j++)
      if (disp[i*COLS+j]!=0) 
	error(i,j,0,disp[i*COLS+j]);
  
  for (i=ipos+isize+9;i<ROWS;i++)
    for (j=jpos-10;j<COLS;j++)
      if (disp[i*COLS+j]!=0) 
	error(i,j,0,disp[i*COLS+j]);

  for (i=ipos-10;i<ROWS;i++)
    for (j=0;j<=jpos-10;j++)
      if (disp[i*COLS+j]!=0) 
	error(i,j,0,disp[i*COLS+j]);

}

// ------------------------------------------------------------------------- //

void printperf(int runtype, int np, double dsecs, double esecs, int iter, double truethrput, double rawtime) {
  double secs;
  double dpercent, epercent;

  if (DEBUG)
    printf("Entering printperf, %f %f %d \n",dsecs,esecs,iter);

  dsecs = dsecs/iter;
  esecs = esecs/iter;
  secs  = dsecs+esecs;

  //printf("%f %f %f \n",dsecs,esecs,secs);
  
  dpercent = dsecs*100/secs;
  epercent = esecs*100/secs;

  //printf("%f %f \n",dpercent,epercent);

  printf("***********************************************\n");
  printf("* Multiple-Baseline Stereo (%3d x %3d)        *\n", ROWS, COLS);
  printf("***********************************************\n");
  printf("* Time for one image set:                     *\n");
  printf("* diff     : %9.6f (%2.0f percent)           *\n",dsecs,dpercent);
  printf("* error    : %9.6f (%2.0f percent)           *\n",esecs,epercent);
  printf("* total(s) : %9.6f                        *\n"    ,secs);
  printf("***********************************************\n\n");
  //printf("* Throughput :   %6.2f frames/sec            *\n",1.0/(dsecs+esecs));
  //printf("***********************************************\n");

  switch (runtype) {
    case CASEDEAL:
      printf("Deal%d,",np-2);
      break;
    case CASE3ST:
      printf("3ST,");
      break;
    case CASESEQ:
      printf("Seq,");
      break;
    case CASEMPI:
      printf("MPI,");
      break;
  }
  printf(" %d procs,", np);
  printf(" throughput was %4.2f frames/sec, run-time on P%d was %4.2f secs, %d images.\n",truethrput, np-1, rawtime, ITERS);
}


// ------------------------------------------------------------------------- //

// The first stage: data generation
void stage1 (int np) { 
  int    count; 
  int data[3*(ROWS+WINY)*COLS];
  int  dealw;
  int  nextw;

  // The three arrays to transmit
  int *ref, *m1, *m2;
  
  int size;
  int i,j,ipos,jpos;
  int tardis,isize,jsize;

  ref = &data[0];
  m1  = &data[(ROWS+WINY)*COLS];
  m2  = &data[2*(ROWS+WINY)*COLS];
    
  dealw = np-2;  // number of workders in the next stage
  nextw = 0;     // first worker to send to, by pid within that imaginary group
  if (DEBUG) printf("Entering stage1 ...  \n");
  
  // We transmit data of type MPI_int_array, of length 3 (ie three arrays together)
  
  MPI_Type_size(MPI_int_array,&size);
  
  // Filling the arrays
  tardis = 12;
  isize=37;
  jsize=37;
  
  for (count=0;count<ITERS;count++) {
    if (DEBUG) printf("Generate a data and send it to the next stage - %d \n",count+1);
    
    for (i=0;i<6;i++)
      for (j=0;j<COLS;j++) {
	ref[i*COLS+j]=0;
	m1[i*COLS+j] =0;
	m2[i*COLS+j] =0;
      }
    
    for (i=6;i<ROWS+6;i++)
      for (j=0;j<COLS;j++) {
	ref[i*COLS+j]=i+j;
	m1[i*COLS+j] =i+j;
	m2[i*COLS+j] =i+j;
      }
    
    for (i=ROWS+6;i<ROWS+13;i++)
      for (j=0;j<COLS;j++) {
	ref[i*COLS+j]=0;
	m1[i*COLS+j] =0;
	m2[i*COLS+j] =0;
      }
    
    ipos = (ROWS-isize)/2 + 6 - 1;
    jpos = (COLS-jsize)/2 - tardis - 1;
    
    for (i=ipos;i<ipos+isize;i++)
      for (j=jpos;j<jpos+jsize;j++) {
	ref[i*COLS+j]         = i-ipos + j-jpos;
	m1[i*COLS+j+tardis]   = i-ipos + j-jpos;
	m2[i*COLS+j+2*tardis] = i-ipos + j-jpos;
      }
    
    MPI_Send(data, 3, MPI_int_array, nextw+1, 0, MPI_COMM_WORLD);  //+1 includes offset within MPI_COMM_WORLD
    nextw = (nextw+1)%dealw;
  }  

  MPI_Send(data, 0, MPI_int_array, nextw+1, SPECIALBYEBYE, MPI_COMM_WORLD);  //+1 includes offset within MPI_COMM_WORLD
  nextw = (nextw+1)%dealw;

  for (count=1;count<np-2;count++) {
    MPI_Send(data, 0, MPI_int_array, nextw+1, BYEBYE, MPI_COMM_WORLD);  //+1 includes offset within MPI_COMM_WORLD
    nextw = (nextw+1)%dealw;
  }  
}


// treatment of one set of images
// worker of the Deal nested in stage 2

void stage2 (int np, int myrank) {
  int i,j; 
  int data[5*(ROWS+WINY)*COLS];
  
  int *ref, *m1, *m2;
  float curbesterr[ROWS*COLS];
  int   curbestdisp[ROWS*COLS];
  float diffimg[(ROWS+WINY)*COLS]; // intermediate image
  int dataout[(ROWS+WINY)*COLS];

  int curdisp,size;
  int myiters, imagecnt;
  double secs;
  MPI_Status status;
  

  MPI_Type_size(MPI_int_array,&size);

  MPI_Recv(data,5,MPI_int_array, 0, MPI_ANY_TAG, MPI_COMM_WORLD, &status);
  while ( status.MPI_TAG == 0) {

    ref = &data[0];
    m1  = &data[(ROWS+WINY)*COLS];
    m2  = &data[2*(ROWS+WINY)*COLS];
       
    for (i=0;i<ROWS;i++)
      for (j=0;j<COLS;j++) {
          curbesterr[i*COLS+j]  = 9999999.0;
  	  curbestdisp[i*COLS+j] = 0;
    }   
    
    for (curdisp=0;curdisp<MAXDISP;curdisp++) {
      secs = MPI_Wtime();
      gendiffimg(ref,m1,m2,diffimg,curdisp);
      dsecsG = dsecsG + MPI_Wtime() - secs;
      secs = MPI_Wtime();
      updatedispimg(diffimg,curbesterr,curbestdisp,curdisp);
      esecsG = esecsG + MPI_Wtime() - secs;
    }
    
    MPI_Send(curbestdisp, ROWS*COLS, MPI_INT, np-1, 0, MPI_COMM_WORLD);
    MPI_Recv(data,5,MPI_int_array, 0, MPI_ANY_TAG, MPI_COMM_WORLD, &status);
  }
  if (status.MPI_TAG == SPECIALBYEBYE)  {
     MPI_Send(curbestdisp, 0, MPI_INT, np-1, BYEBYE, MPI_COMM_WORLD);
  }
}

void stage3 (int np) {
  int i,j, imagecnt=0; 

  double oldsecs = 0, newsecs;
  MPI_Status status;
  
  int curbestdisp[(ROWS+WINY)*COLS];
  int dealw, nextw;
    
  dealw = np-2;  // number of workders in the next stage
  nextw = 0;     // first worker to send to, by pid within that imaginary group


  MPI_Recv(curbestdisp, 1, MPI_int_array, nextw+1, MPI_ANY_TAG, MPI_COMM_WORLD, &status);  //+1 includes offset within MPI_COMM_WORLD
  nextw = (nextw+1)%dealw;
  while (status.MPI_TAG != BYEBYE) {
    testdata(curbestdisp);
    newsecs = MPI_Wtime();
    st3secs[imagecnt++] = newsecs - oldsecs;
    oldsecs = newsecs;
    MPI_Recv(curbestdisp, 1, MPI_int_array, nextw+1, MPI_ANY_TAG, MPI_COMM_WORLD, &status);  //+1 includes offset within MPI_COMM_WORLD
    nextw = (nextw+1)%dealw;
  }

  // Compute and send true throughput
  truethrput = 0;
  for (i=1; i<ITERS; i++) truethrput += st3secs[i];
  truethrput = 1/(truethrput/(ITERS-1));
} 

// ------------------------------------------------------------------------- //
int main (int argc, char *argv[])
{
  int     i, count, curdisp;
  float   start;

  // pipeline stuff
  int mystagenum;  
  int size;
  
  void *buf;
  MPI_Status junk;
  int myrank;
  double secs, runtime;
  int nprocs;

  // Initializing MPI and eSkel
  MPI_Init(&argc, &argv);
  MPI_Comm_size(MPI_COMM_WORLD, &nprocs);
  MPI_Comm_rank(MPI_COMM_WORLD, &myrank);

  // Starting the application
  printf("Entering the Raw MPI Multiple-Baseline Stereo application, Process %d\n", myrank);

  secs = MPI_Wtime();

  dsecsG = 0;
  esecsG = 0;

  // Customized types
  MPI_Type_contiguous((ROWS+WINY)*COLS, MPI_INT, &MPI_int_array);
  MPI_Type_commit(&MPI_int_array);

  if (myrank==0)             mystagenum = 0;
  else if (myrank==nprocs-1) mystagenum = 2;
  else                       mystagenum = 1;
  
  switch (mystagenum) {
    case 0: stage1(nprocs); break;
    case 1: stage2(nprocs, myrank); break;
    case 2: stage3(nprocs); break;
  }

  runtime = MPI_Wtime()-secs;
  
  // Process 0 in the global communicator collect results and print them
  if (myrank==0) {
    buf = (double *) malloc(sizeof(double)*2);
    for (i=1; i<nprocs-1; i++) {
      MPI_Recv(buf,2,MPI_DOUBLE,i,MPI_ANY_TAG,MPI_COMM_WORLD,&junk); 
      dsecsG += ((double *)buf)[0];
      esecsG += ((double *)buf)[1];
    }
    MPI_Recv(&runtime,1,MPI_DOUBLE,nprocs-1,MPI_ANY_TAG,MPI_COMM_WORLD,&junk); 
    MPI_Recv(&truethrput,1,MPI_DOUBLE,nprocs-1,MPI_ANY_TAG,MPI_COMM_WORLD,&junk); 
    printperf(CASEMPI,nprocs,dsecsG,esecsG,ITERS,truethrput,runtime);
  }
  
  if (mystagenum==1) {
    buf = (double *) malloc(sizeof(double)*2);
    ((double *)buf)[0]=dsecsG;
    ((double *)buf)[1]=esecsG;
    MPI_Send(buf,2,MPI_DOUBLE,0,0,MPI_COMM_WORLD);
  }

  if (mystagenum==2) {
    MPI_Send(&runtime,1,MPI_DOUBLE,0,0,MPI_COMM_WORLD);
    MPI_Send(&truethrput,1,MPI_DOUBLE,0,0,MPI_COMM_WORLD);
  }

  // Finalize MPI
  MPI_Finalize();
  return 0;
}