#include <stdio.h>
#include <omp.h>
#include <stdlib.h>
#include <string.h>
#include <mpi.h>
#include "dimension.h"
#define FILENAME "foutput_cpu.dat"
// Change here the number of steps, the cell geometry, etc
#define STEPITER 1000
#define delx 0.5
#define dely 0.25
// end change here.
/***/
       extern void itstep( int mx, int my, void *f, void *newf, double rdx2, double rdy2, double beta );
       int main( int argc, char **argv )
{
       int niter, i, j, n, mx1, mx2, my1, my_number, n_of_nodes, totalmx, partmx, leftmx, mx, my, k;
       int prev, next, rank, size;
       FILE *fp;
       double howlong;
       double rdx2, rdy2, beta;
/***/
       MPI_Init( &argc, &argv );
       MPI_Comm_rank( MPI_COMM_WORLD, &my_number );
       rank = my_number;
       MPI_Comm_size( MPI_COMM_WORLD, &n_of_nodes );
       size = n_of_nodes;
       prev = (rank > 0)      ? rank-1 : MPI_PROC_NULL;
       next = (rank < size-1) ? rank+1 : MPI_PROC_NULL;
       if ( argc != 4 )
        {
	 if ( !my_number )
	  {
	   fprintf( stderr, "Usage: %s <nrows> <ncolumns> <niter>\n", argv[0] );
          }
         MPI_Finalize();
	 return (-1);
	}
       totalmx = mx = (int)atol( argv[1] );
       my = (int)atol( argv[2] );
       if ( my < 1 )
        {
	 if ( !my_number )
	  {
	   fprintf( stderr, "Number of columns (%d) should be positive\n", my );
	  }
	 MPI_Finalize(); 
	 return (-1); 
	}
       niter = (int)atol( argv[3] );
       if ( my < 1 )
        {
	 if ( !niter )
	  {
	   fprintf( stderr, "Number of iterations (%d) should be positive\n", niter );
	  }
	 MPI_Finalize(); 
	 return (-1); 
	}
/* Compute the number of rows per node: */
       mx = (totalmx+n_of_nodes-1)/n_of_nodes;
/* This is the number of rows for all but the last: */
       partmx = mx;
/* This is the number of rows for the last: */
/* It cannot be greater than partmx, but it can be non-positive: */
       leftmx = totalmx - partmx*(n_of_nodes-1);
       if ( leftmx < 1 )
        {
	 if ( !my_number )
	  {
	   fprintf( stderr, "Cannot distribute rows, too many processors\n" );
	  } 
         MPI_Finalize();
	 return (-1); 
	}
       if ( my_number == (n_of_nodes-1) ) mx = leftmx;
/* End rows distribution. */
       partmx += 2;	
       mx += 2;
       my += 2;
/* Here we know the array sizes, so make the arrays themselves: */
       {
        DIM2( double,    f, my );
        DIM2( double, newf, my );
/***/
        f =    (typeof(f))malloc( partmx*sizeof(*f) );
        newf = (typeof(f))malloc( partmx*sizeof(*f) );
        if ( (!f) || (!newf) )
         {
          fprintf( stderr, "Cannot allocate, exiting\n" );
          exit( -1 );
         }
        rdx2 = 1./delx/delx;
        rdy2 = 1./dely/dely;
        beta = 1.0/(2.0*(rdx2+rdy2)); 
        if ( !my_number )
         {
	  printf( "Solving heat conduction task on %d by %d grid by %d processors\n",
                  totalmx, my-2, n_of_nodes );	 
	  fflush( stdout );
	 }
        for ( i = 0; i < mx; i++ )
         {
          for ( j = 0; j < my; j++ )
           {
            if ( ((i==0)&&(my_number==0)) || (j==0) || ((i==(mx-1))&&(my_number==(n_of_nodes-1))) || (j==(my-1)) )
             newf[i][j] = f[i][j] = 1.0;
            else
             newf[i][j] = f[i][j] = 0.0;      
           }
         }	
        mx1 = mx - 1;
        my1 = my - 1; 
        MPI_Barrier( MPI_COMM_WORLD );
/* Iteration loop: */
        howlong = MPI_Wtime();
        for ( n = 0; n < niter; n++ )
         {
	  if ( !my_number )
	   {
	    if ( !(n%STEPITER) )
             printf( "Iteration %d\n", n );
	   }
/* Step of calculation starts here: */ 
	  itstep( mx, my, f, newf, rdx2, rdy2, beta );
/* Do all the transfers: */
	  MPI_Sendrecv( &f[1][0],    mx, MPI_DOUBLE, prev, 0, &f[0][0],    mx, MPI_DOUBLE, prev, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE );
          MPI_Sendrecv( &f[mx-2][0], mx, MPI_DOUBLE, next, 0, &f[mx-1][0], mx, MPI_DOUBLE, next, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE );	   
	 }
#pragma omp barrier
#pragma omp master
        if ( !my_number )
         {
          printf( "Elapsed time: %f sec\n", (float)(MPI_Wtime()-howlong) );	
	  fp = fopen( FILENAME, "w" );
	  fclose( fp );
	 }
        for ( j = 0; j < n_of_nodes; j++ )	
         {
	  MPI_Barrier( MPI_COMM_WORLD );
	  if ( j == my_number )
	   {
	    fp = fopen( FILENAME, "a" );
	    for ( i = 1; i < (mx-1); i++ )
	     { 
	      fwrite( &(newf[i][1]), my-2, sizeof(newf[0][0]), fp );
//	      for ( k = 0; k < my; k++ ) fprintf( fp, "%f ", newf[i][k] );
//	      fprintf( fp, "\n" );
	     } 
	    fclose( fp );
	   }
	 }
#pragma omp barrier	 
       }	 
       MPI_Finalize();
       return 0;
}