#include <stdio.h>
#include <math.h>
#include <omp.h>
#include <malloc.h>
#define EPSILON 0.00001
#define N 1000
#define time_steps 100

int main (int argc, char *argv[])
 {
  int i,j;
  int step;
  int threads;
  double time;
  double eps, enew;
  double time_max = 3.0;
  double alpha = 0.06;
  double dx = 1.0/N;
  double dy = 1.0/time_steps;
  double dt = time_max/time_steps;
  double dxinv = 1.0/dx;
  double dyinv = 1.0/dy;
  double dtinv = 1.0/dt;
  double divinv = 1.0/(dtinv + 2 * alpha * (dxinv * dxinv + dyinv * dyinv));
  double **t,*tS;
  double **told,*toldS;
  double minval, maxval;
  double start, end;
  char fname[40];
  FILE *out;

  if (argc < 2) threads = 1; 
  else threads = atoi(argv[1]);
  if (!omp_get_dynamic()) 
   {
    //Set the number of threads 
    omp_set_num_threads(threads);
   }
  else
   {
    printf("ERROR: set OMP_DYNAMIC=FALSE\n");
    exit(0);
   }

  start = omp_get_wtime();
  tS = (double *) malloc(N * N * sizeof(double));
  toldS = (double *) malloc(N * N * sizeof(double));
  t = (double **) malloc(N * sizeof(double *));
  told = (double **) malloc(N * sizeof(double *));
  #pragma omp parallel for
  for (i=0; i<N; i++)
   {
    t[i] = &tS[i*N];
    told[i] = &toldS[i*N]; 
   }
  // initialize interior values
  #pragma omp parallel for private(j)
  for (i=1; i<(N-1); i++)
   for (j=1; j<(N-1); j++)
    told[i][j] = 0.0;  
  // set initial boundary conditions
  #pragma omp parallel for
  for (i=0; i<N; i++)
   {
    told[i][0]   = 0.0; // left
    told[i][N-1] = 0.0; // right
   }
  #pragma omp parallel for
  for (j=0; j<N; j++)
   told[N-1][j] = 0.0; // bottom
  // for all time steps
  for (step = 1; step <= time_steps; step++)
   {
    time = step * (time_max/time_steps);
    // reset top boundary condition each timestep
     #pragma omp parallel for
    for (j=0; j<N; j++)
     told[0][j]   =  2.0 * sin(time); // top 
    do
     {
      eps = 0.0;
      #pragma omp parallel for private(j) schedule(runtime)
      for (i=1; i<(N-1); i++)
       for (j=1; j<(N-1); j++)
        t[i][j]=((told[i][j+1]+told[i][j-1])*alpha*dyinv*dyinv+(told[i+1][j]+told[i-1][j])*alpha*dxinv*dxinv+(told[i][j]*dtinv))*divinv;
      //#pragma omp parallel for private(j,enew)
      for (i=1; i<(N-1); i++)
       {
        for (j=1; j<(N-1); j++)
         {
          enew = fabs(t[i][j] - told[i][j]);
          //#pragma omp critical
          if (enew > eps) { eps = enew; }
         }
       }
      #pragma omp parallel for private(j)
      for (i=0; i<N; i++)
       for (j=0; j<N; j++)
        told[i][j] = t[i][j];
     }
    while(eps > EPSILON);

    // Dump raster date to a file
    minval = 0.0;
    maxval = 0.0;
    for (i=0; i<N; i++)
     for (j=0; j<N; j++)
      {
       if (t[i][j] < minval) { minval = t[i][j]; }
       if (t[i][j] > maxval) { maxval = t[i][j]; }
      }
    sprintf(fname,"Output\\heat%03d.raw",step); 
    out = fopen(fname,"w+b");
    for (i=0; i<N; i++)
     for (j=0; j<N; j++)
      fprintf(out,"%c",(int)(((t[i][j]-minval)*255.0)/(maxval - minval)));
    fclose(out);
    printf("Time step: %d\r",step);
   } // for all time steps  
  end = omp_get_wtime();
  printf("%d\t%.2f\n",omp_get_max_threads(),end-start);
 }