#include "slice.h"

// main represents code shared by all processes:
int main(int argc, char** argv)
{
  int rank, size, i, j;
  int themaster = 0;
  MPI_Init(&argc,&argv);
  MPI_Comm_rank(MPI_COMM_WORLD,&rank);
  MPI_Comm_size(MPI_COMM_WORLD,&size);
  if (size != (P*P)+1) {printf("wrong np\n"); exit(1);}

  // seperate worker from master process:  (must be done by all procs)
  MPI_Group worldgroup, workersgroup;
  MPI_Comm_group(MPI_COMM_WORLD,&worldgroup);
  MPI_Group_excl(worldgroup,1,&themaster,&workersgroup);
  MPI_Comm_create(MPI_COMM_WORLD,workersgroup,&WORKERS);

  if (rank==0)
      master();
  else  
      worker();

  MPI_Finalize();
}

// Called by root to record relative ranks of each process in WORKERS.
// That is, if a process has rank 2 in MPI_COMM_WORLD and 
void registerworkers()
{
  int i, y;
  MPI_Status stat;
    // needs to know how to contact each proc in WORKERS
  for(i=0;i<P*P;i++)
    {
      MPI_Recv(&y,1,MPI_INT,MPI_ANY_SOURCE,1,MPI_COMM_WORLD,&stat);
      locate[y] = stat.MPI_SOURCE;
      // printf("%d in WORKERS has is %d in COMM_WORLD\n",y,locate[y]);
      // fflush(stdout);
    }
} // registerworkers

// called by Worker to register itself with coordinator:
void registercoord(int y)  // y is your rank in WORKERS
{
   MPI_Send(&y,1,MPI_INT,0,1,MPI_COMM_WORLD);  // tag 1 used
}


/* 3D visualization routines:  - active only when GRAPHICS is 1*/
void setup()
{ 
  int M2;
  M2 = N+2;
  tiltR = toRadians(220);   // 140 for rear view
  ZR = 0.5; 
  XDIM=M2*3;  ZDIM=M2;  YDIM = 2*M2;  // x,z will be x,y in simulation
  xshift = (int)(ZDIM * cos(tiltR));
  yshift = (int)(ZDIM * sin(tiltR));

  MPE_Open_graphics(&graph,MPI_COMM_WORLD,(char*)0,0,0,
		    XDIM+(int)abs(xshift)+10,YDIM+(int)abs(yshift)+10, 
		    MPE_GRAPH_INDEPENDENT);
} //setup

void visualize(NROW *BB)
{
  int i, j, k, kr, kd, mrank, U; 
  double z, zr, zd;
  MPE_Color mc;
  // erase background
  MPE_Fill_rectangle(graph,0,0,XDIM+(int)abs(xshift)+10,
		    YDIM+(int)abs(yshift)+10,MPE_WHITE);
  drawbox(graph,0,0,0,XDIM-1,YDIM-1,ZDIM-1,MPE_GREEN);
  U = 4;  // gap between lines in grid
  // draw as a grid
  for(i=0;i<N+2;i+=U)   // rows  - z coord on screen
    for(j=0;j<N+2;j+=U) // cols  - x coord on screen
      {
	mrank = locate[i/(M-1)];
        mc = mcolor[mrank%9];
	z = -1*BB[i][j]*(YDIM/2.0);
	k = (int)(z + (YDIM/2.0)+0.5);
	if (k<0) k=0; if (k>=YDIM) k=YDIM-1;
        if (j+U<N+2) 
	  {
	    zr=-1*BB[i][j+U]*(YDIM/2.0);
	    kr = (int)(zr + (YDIM/2.0)+0.5);
	    if (kr<0) kr=0; if (kr>=YDIM) kr=YDIM-1;
	    drawline(graph,3*j,k,i,3*(j+U),kr,i,mc);
	  }
	if (i+U<N+2)
	  {
	    zd=-1*BB[i+U][j]*(YDIM/2.0);
	    kd = (int)(zd + (YDIM/2.0)+0.5);
	    if (kd<0) kd=0; if (kd>=YDIM) kd=YDIM-1;
	    drawline(graph,3*j,k,i,3*j,kd,i+U,mc);
	  }
      } // for for
  MPE_Update(graph);
  if (mousesynch) MPE_Get_mouse_press( graph, &i, &j, &k );
} //visualize



// Master function
void master()
{
  int i, j, x, y, r, c;
  int resume = 1;
  int vote = 0;  // 0 means should not resume
  static MPI_Status stat[P*P];
  static MPI_Request req[P][P];

  registerworkers(); // sets locate[] matrix

  // create Data matrix
  Data = (NROW *)malloc(sizeof(double)*(N+2)*(N+2));
  Q = (NROW *)malloc(sizeof(double)*(N+2)*M); // for comm
  
 // GENERATE Data matrix
  // interior
  for(i=1;i<N+1;i++) // rows (y)
    for(j=1;j<N+1;j++) // cols (x)
    { 
      Data[i][j] = sin(PI/(1+(random() % 4)));
    }
  // set boundary values:
  for(j=0;j<N+2;j++) 
    { 
      Data[0][j] = sin((PI*j)/(N+1));
      Data[N+1][j] = sin((PI*j)/(N+1))*exp(-1*PI);
    }
  for(i=0;i<N+2;i++)
    { Data[i][0] = 0;  Data[i][N+1] = 0; }

  // setup graphics
  if (GRAPHICS)  
    { setup();   visualize(Data); 
    }

  // send slices to workers
  for(i=0;i<P*P;i++)
    MPI_Send(SLICE(Data,i),(N+2)*M,MPI_DOUBLE,locate[i],0,MPI_COMM_WORLD);
  ///////////////////////

  while (resume)
    {
        MPI_Allreduce(&vote,&resume,1,MPI_INT,MPI_LOR,MPI_COMM_WORLD);
	//printf("master:vote is %d\n",resume);
	if (GRAPHICS)  // if using graphics, need to receive from workers
	  {
	    for(i=0;i<P*P;i++)
	      { 
    	        MPI_Recv(Q,(N+2)*M,MPI_DOUBLE,
			 locate[i],3,MPI_COMM_WORLD,stat); 
		// copy back to Data matrix, without affecting boundary 
		for(r=0;r<M-2;r++) 
		  for(c=0;c<N;c++)
		    { Data[r+(i*(N/(P*P))+1)][c+1] = Q[r+1][c+1]; }
	      } // for i,j

	    visualize(Data);  
	  } // if show graphics
    } // while resume

  if (!GRAPHICS) // receive final results
    {
	    for(i=0;i<P*P;i++)
	      { 
    	        MPI_Recv(Q,(N+2)*M,MPI_DOUBLE,
			 locate[i],3,MPI_COMM_WORLD,stat); 
		// copy back to Data matrix, without affecting boundary 
		for(r=0;r<M-2;r++) 
		  for(c=0;c<N;c++)
		    { Data[r+(i*(N/(P*P))+1)][c+1] = Q[r+1][c+1]; }
	      } // for i,j      
    }
  printf("Convergence!\n");
  
  ///////////////////////
} //master


// Worker function
void worker()
{

  // find my rank within workers process
  MPI_Comm_rank(WORKERS,&mrank);
  // register rank with master
  registercoord(mrank);
  // create work matrices:
  
  // set ranks for north/south neighbors:
  nrank = mrank-1;   srank = mrank+1;

  Q = (NROW *)malloc(sizeof(double)*M*(N+2));
  R = (NROW *)malloc(sizeof(double)*M*(N+2));      
  // receive submatrix from master
  MPI_Recv(Q,M*(N+2),MPI_DOUBLE,0,0,MPI_COMM_WORLD,stat);

  // ...

} //worker