/*

  Explicitly solves the 2D heat conduction equation:

    dT/dt = alpha * div^2 T

  Note that this requires the HDF5 to compile ('module load hdf5' on the MSU HPCC, 
  or http://www.hdfgroup.org/HDF5/).  Compilation:

    g++ -L/usr/local/lib -lhdf5_hl -lhdf5 -h/usr/local/include conduction_2d.C -o conduction

  Run:

    ./conduction

  There are several user-modifiable parameters, defined in main().

  Written by: Brian O'Shea, oshea@msu.edu

*/

#include "hdf5.h"
#include "hdf5_hl.h"
#include <math.h>

int create_ICs(double **Temperature, 
	       int xgrid, int ygrid, 
	       double xstart, double xend, 
	       double ystart, double yend,
	       int simulation_type);

double calculate_dt(double dx, double dy, double alpha);

int calculate_Tdot(double **Temperature, double **Tdot, double alpha, 
		   int xgrid, int ygrid, double dx, double dy, int simulation_type);

int update_Temp(double **Temperature, double **Tdot, double dt, int padded_x, int padded_y);

int write_data_output(double **Temperature, int grid_x, int grid_y, int this_output);

int main(){

  /* ---------  User needs to set all of this stuff --------- */
  int grid_x=200,grid_y=200;  // grid x,y cell dimensions

  double stop_time = 5.0;  // stop time: around 5 is good for gaussian pulse, 2 for sine wave. 

  int simulation_type = 1;  // 1 = gaussian pulse, 2 = infinite sine wave

  double dt_data_output = 0.1;  // output cadence (make it small for lots of data outputs)

  /* --------- End of user-set parameters --------- */

  double x_min=-10.0, x_max=10.0;  // grid x min, max values
  double y_min=-10.0, y_max=10.0;  // grid y min, max values
  double alpha = 1.0;  // conductivity


  double **Temperature=NULL;
  double **Tdot=NULL;
  double dx, dy, dt;

  int i,j,padded_x,padded_y;

  // Grid is padded by 2 cells in each direction for boundary conditions
  padded_x = grid_x+2;
  padded_y = grid_y+2;

  /* The two arrays that we need: temperature and dT/dt.  These are
     padded by 2 compared to what the user has specified so that we can 
     take into account the boundary conditions.  So, we have to dynamically
     allocate these arrays and then set them to some default values.  */
  Temperature = new double*[padded_x];
  Tdot = new double*[padded_x];

  for(i=0; i<padded_x; i++){
    Temperature[i] = new double[padded_y];
    Tdot[i] = new double[padded_y];
  }

  for(i=0; i<padded_x; i++)
    for(j=0; j<padded_y; j++){
      Temperature[i][j] = 1.0;
      Tdot[i][j] = 0.0;
    }

  // calculate grid dx, dy
  dx = (x_max-x_min)/double(grid_x);
  dy = (y_max-y_min)/double(grid_y);

  // actually create the initial conditions
  create_ICs(Temperature, grid_x, grid_y, x_min, x_max, y_min, y_max, simulation_type);

  double this_time=0.0, next_output_time=0.0;
  int this_cycle=0;
  int this_output=0;

  // write out simulation data before calculations start so we can see the ICs
  write_data_output(Temperature, grid_x, grid_y, this_output);

  next_output_time = dt_data_output;

  /* This is the main loop.  Evolve forward in time until we reach stop_time,
     calculating timestep, dT/dt, updating the temperature, and (if we want)
     writing out the dataset.  */
  while(this_time <= stop_time){

    /* calculate timestep.  Note that this is not strictly necessary, as it does not
       change from step to step. */
    dt = calculate_dt(dx, dy, alpha);

    if(this_time + dt > next_output_time) dt = next_output_time - this_time;

    // calculate dT/dt
    calculate_Tdot(Temperature, Tdot, alpha, grid_x, grid_y, dx, dy, simulation_type);

    // update temperature
    update_Temp(Temperature, Tdot, dt, padded_x, padded_y);

    // increment time, cycle
    this_time += dt;
    this_cycle++;

    // write out data if it's time.
    if( this_time >= next_output_time){
      write_data_output(Temperature, grid_x, grid_y, ++this_output);
      next_output_time += dt_data_output;
    }

    printf("Cycle, dt, time, next output:  %d -- %lf, %lf, %lf\n", this_cycle, dt, this_time, next_output_time);

  }

  // write data out one last time.
  if( this_time >= next_output_time)
    write_data_output(Temperature, grid_x, grid_y,  ++this_output);

  return 0;
}

/* Initializes initial temperature field as either a gaussian pulse (simulation_type = 1)
   or an infinite set of sine waves (simulation_type = 2).  

Inputs:  
  **Temperature: Temperature field
  xgrid, ygrid:  cells along x,y edges (padded by 2 in each dimension)
  xstart, xend:  position of non-padded x start, end (spatial, not cells)
  ystart, yend:  as x
  simulation_type: determines whether it's a gaussian pulse (1) or sine waves (2).
*/

int create_ICs(double **Temperature, 
	       int xgrid, int ygrid, 
	       double xstart, double xend, 
	       double ystart, double yend,
	       int simulation_type){

  int i,j;

  // keep track of background temp
  double background_temperature = Temperature[0][0];
  double dx, dy, my_x, my_y, x_center, y_center;

  dx = (xend-xstart)/double(xgrid);
  dy = (yend-ystart)/double(ygrid);
  
  x_center = (xend+xstart)/2.0;
  y_center = (yend+ystart)/2.0;

  //printf("dx,dy: %e %e\n",dx, dy);
  //printf("x,y center: %e %e\n", x_center, y_center);

  if(simulation_type==1){  // set up gaussian pulse

    // make pulse fairly narrow
    double fwhm = fabs(xend-xstart)/8.0;

    // loop through cells calculating position, then setting temperature to pulse
    for(i=1; i<xgrid+1; i++)
      for(j=1; j<ygrid+1; j++){
	my_x = xstart + (double(i)-0.5)*dx;  // the -0.5 instead of +0.5 takes into account the cell of padding
	my_y = ystart + (double(j)-0.5)*dx;

	Temperature[i][j] = 10.0 *
	  exp( -1.0*(pow(my_x-x_center,2.0) + pow(my_y-y_center,2.0)) / (2.0*fwhm*fwhm) );
	
      }
    
    // set the edges of the box to the background temperature
    for(j=0;j<ygrid+2;j++)
      Temperature[0][j]=Temperature[xgrid+1][j]=background_temperature;
    for(i=0;i<xgrid+2; i++)
      Temperature[i][0]=Temperature[i][ygrid+1]=background_temperature;

  } else if(simulation_type==2){  // sine waves
    
    // calculate wave number assuming 4 waves
    double k = 2.0*M_PI*4.0/(xend-xstart);

    // loop through cells calculating position, then setting temperature to sine wave
    for(i=1; i<xgrid+1; i++)
      for(j=1; j<ygrid+1; j++){
	my_x = xstart + (double(i)-0.5)*dx;  // the -0.5 instead of +0.5 takes into account the cell of padding
	my_y = ystart + (double(j)-0.5)*dx;

	Temperature[i][j] = 5.0*(1.0+sin(k*my_x));
      }
    
    // set the edges of the box to make it periodic
    for(j=0;j<ygrid+2;j++){
      Temperature[0][j] = Temperature[xgrid][j];
      Temperature[xgrid+1][j]= Temperature[1][j];
    }
    for(i=0;i<xgrid+2; i++){
      Temperature[i][0]=Temperature[i][ygrid];
      Temperature[i][ygrid+1] = Temperature[i][1];
    }

  } else {
    printf("simulation type undefined (%d) - exiting!\n", simulation_type);
    return -1;
  }

  /* 
     // UNCOMMENT THIS IF YOU WANT AN ASCII OUTPUT OF YOUR ICs
  for(i=0; i<xgrid+2; i++){
    for(j=0; j<ygrid+2; j++){
      printf("%.2lf ",Temperature[i][j]);
    }
    printf("\n");
  }
  */

  return simulation_type;
}

/* Calculates timestep, using grid cell size and conductivity coefficient. Returns timestep. */
double calculate_dt(double dx, double dy, double alpha){
  return 0.3 * 0.5 * 0.5 * (dx*dx + dy*dy)/alpha;
}

/* Calculates dT/dt using temperature, information about grid.  Set boundary conditions appropriately.
  Temperature = temperature field.
  Tdot = dT/dt field
  alpha = conductivity
  xgrid, ygrid: size of non-padded grid (actually 2 cells longer in each direction)
  dx, dy:  cell size along x,y dimension
  simulation_time: gaussian pulse (1) or sine wave (2).  Necessary for ICs.  */
int calculate_Tdot(double **Temperature, double **Tdot, double alpha, 
		   int xgrid, int ygrid, double dx, double dy, int simulation_type){

  int i,j;

  // calculate dT/dt for non-edge cells
  for(i=1;i<xgrid+1;i++)
    for(j=1;j<ygrid+1;j++){
      Tdot[i][j] = alpha * 
	( (1.0 / (dx*dx)) * (Temperature[i+1][j]-2.0*Temperature[i][j]+Temperature[i-1][j])
	+ (1.0 / (dy*dy)) * (Temperature[i][j+1]-2.0*Temperature[i][j]+Temperature[i][j-1]) );
    }

  if(simulation_type==1){ // gaussian pulse: boundary conditions are at constant temp, so dT/dt = 0
    for(j=0;j<ygrid+2;j++)
      Tdot[0][j]=Tdot[xgrid+1][j]=0.0;
    for(i=0;i<xgrid+2; i++)
      Tdot[i][0]=Tdot[i][ygrid+1]=0.0;

  } else if(simulation_type==2){  // infinite sine pulse: BCs are periodic, so copy them from the other side

    for(j=0;j<ygrid+2;j++){
      Tdot[0][j] = Tdot[xgrid][j];
      Tdot[xgrid+1][j]= Tdot[1][j];
    }

    for(i=0;i<xgrid+2; i++){
      Tdot[i][0]=Tdot[i][ygrid];
      Tdot[i][ygrid+1] = Tdot[i][1];
    }

  } else {
    printf("simulation type (%d) undefined!\n",simulation_type);
    return -1;
  }

  return simulation_type;
}

/* Update temperature at time N+1 using temperature at time N, dT/dt at time N.

Inputs:
  Temperature: temperature field
  Tdot:  dT/dt field
  dt:  timestep
  padded_x, padded_y: number of cells along each edge (padded) */
int update_Temp(double **Temperature, double **Tdot, double dt, int padded_x, int padded_y){

  int i,j;

  for(i=0;i<padded_x;i++)
    for(j=0;j<padded_y;j++)
      Temperature[i][j] += dt*Tdot[i][j];

  return 0;
}

/* Writes out Temperature data to an HDF5 file named by the cycle number. 

Input:
   Temperature:  temperature field
   grid_x, grid_y: un-padded cell count along x, y dimensions
   this_output:  output number (used to create file name).  */
int write_data_output(double **Temperature, int grid_x, int grid_y, int this_output){

  printf("-------------\n");
  printf("Writing data. Output number %d\n",this_output);
  printf("-------------\n");

  int i,j, buffersize;

  double *output_buffer=NULL;

  output_buffer = new double[grid_x*grid_y];  // reduced in size to remove padding

  /*
    // UNCOMMENT THIS IF YOU WANT ASCII OUTPUT TO YOUR SCREEN
  for(i=0; i<grid_x+2; i++){
    for(j=0; j<grid_y+2; j++){
      printf("%.2lf ",Temperature[i][j]);
    }
    printf("\n");
  }
  */

  // copy Temperature into a contiguous 1D output puffer, removing padding as we go
  for(i=1; i<grid_x+1; i++){
    for(j=1; j<grid_y+1; j++){
      output_buffer[ (i-1) + (j-1)*grid_x] = Temperature[i][j];
    }
  }

  // write data out into an HDF5 file
  hid_t       file_id;
  hsize_t     dims[2];
  herr_t      status;

  dims[0]=grid_x;
  dims[1]=grid_y;

  char filename[32];
  snprintf(filename, sizeof(char) * 32, "output%.04i.hdf5", this_output);

  /* create a HDF5 file */
  file_id = H5Fcreate (filename, H5F_ACC_TRUNC, H5P_DEFAULT, H5P_DEFAULT);
  
  /* create and write a double-preicision type dataset named "Temperature" */
  status = H5LTmake_dataset(file_id,"/Temperature",2,dims,H5T_NATIVE_DOUBLE,output_buffer);

  /* close file */
  status = H5Fclose (file_id);

  return 0;
}