/* Heat flow model in C + HDF5
 * 
 * The left edge of a 2D grid starts with an initial temperature. As time
 * passes, the heat flows through the grid.
 *
 * This is an "I am the average of my neighbors" model:
 *  temp = (t1 + t2 + ... + tN) / N
 * 
 * Author: Aaron Weeden, Shodor, 2014
 */

/* LIBS */
#include <hdf5.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

/* GLOBALS */
float **Grid;   // Grid of temperature data
float **Next_Grid; // Next time step of temperature data
int N_Rows = 5; // Number of rows in the 2D temperature grid
int N_Cols = 5; // Number of columns in the 2D temperature grid
int N_Time_Steps = 5; // Number of time steps in the simulation
float Init_Temp = 5.0; // Initial temperature for the left edge of the grid

/* FUNCTIONS */

/* get arguments from command line and set global vars with them */
void parseArgs(int argc, char **argv) {
    int c;

    while((c = getopt(argc, argv, "r:c:t:i:")) != -1) {
        switch(c) {
            case 'r':
                N_Rows = atoi(optarg);
                break;
            case 'c':
                N_Cols = atoi(optarg);
                break;
            case 't':
                N_Time_Steps = atoi(optarg);
                break;
            case 'i':
                Init_Temp = atof(optarg);
                break;
            case '?':
            default:
                fprintf(stderr, "Usage: %s [OPTIONS]\n", argv[0]);
                fprintf(stderr, "OPTIONS:\n");
                fprintf(stderr, "\t-r int   : number of rows\n");
                fprintf(stderr, "\t-c int   : number of columns\n");
                fprintf(stderr, "\t-t int   : number of time steps\n");
                fprintf(stderr, "\t-f float : initial temperature\n");
                exit(EXIT_FAILURE);
        }
    }
}

/* check if the simulation can continue with valid parameters, exit if not */
void validateInput() {
    bool err = false;

    if (N_Rows <= 0) {
        fprintf(stderr, "ERROR: Number of rows (-r %d) must be positive.\n",
                N_Rows);
        err = true;
    }
    if (N_Cols <= 0) {
        fprintf(stderr, "ERROR: Number of columns (-c %d) must be positive.\n",
                N_Cols);
        err = true;
    }
    if (N_Time_Steps < 0) {
        fprintf(stderr, "ERROR: Number of time steps (-t %d) must be ",
                N_Time_Steps);
        fprintf(stderr, "non-negative.\n");
        err = true;
    }
    if (err)
        exit(EXIT_FAILURE);
}

/* allocate memory for grids */
void allocMem() {
    int row_mem_size = N_Cols * sizeof(float);
    int grid_mem_size = N_Rows * row_mem_size;
    int row;

    if ((Grid = (float**)malloc(grid_mem_size)) == NULL) {
        fprintf(stderr, "ERROR allocating memory for 2D grid.\n");
        exit(EXIT_FAILURE);
    }
    for (row=0; row<N_Rows; row++) {
        if ((Grid[row] = (float*)malloc(row_mem_size)) == NULL) {
            fprintf(stderr, "ERROR allocating memory for 2D grid, row %d.\n",
                    row);
            exit(EXIT_FAILURE);
        }
    }

    if ((Next_Grid = (float**)malloc(grid_mem_size)) == NULL) {
        fprintf(stderr, "ERROR allocating memory for next grid.\n");
        exit(EXIT_FAILURE);
    }
    for (row=0; row<N_Rows; row++) {
        if ((Next_Grid[row] = (float*)malloc(row_mem_size)) == NULL) {
            fprintf(stderr, "ERROR allocating memory for next grid, row %d.\n",
                    row);
            exit(EXIT_FAILURE);
        }
    }
}

/* print the temperature at each cell of the grid */
void displayGrid(int time_step) {
    int row, col;

    printf("Time step %d:\n", time_step);
    for (row=0; row<N_Rows; row++) {
        for (col=0; col<N_Cols; col++) {
            printf("%.2f ", Grid[row][col]);
        }
        printf("\n");
    }
    printf("\n");
}

void writeHDF5(int time_step) {
    char filename[80], dataset_name[80];
    hid_t file_id, dataset_id, dataspace_id;
    hsize_t dims[2] = {N_Rows, N_Cols};
    float data[N_Rows][N_Cols];
    int row, col;

    sprintf(filename, "heat%d.h5", time_step);
    sprintf(dataset_name, "/heat%d", time_step);

    file_id = H5Fcreate(filename, H5F_ACC_TRUNC, H5P_DEFAULT, H5P_DEFAULT);

    dataspace_id = H5Screate_simple(2, dims, NULL);

    dataset_id = H5Dcreate2(file_id, dataset_name, H5T_IEEE_F32BE, dataspace_id,
            H5P_DEFAULT, H5P_DEFAULT, H5P_DEFAULT);

    for (row=0; row<N_Rows; row++) {
        for (col=0; col<N_Cols; col++) {
            data[row][col] = Grid[row][col];
        }
    }
    H5Dwrite(dataset_id, H5T_NATIVE_FLOAT, H5S_ALL, H5S_ALL, H5P_DEFAULT, data);

    H5Dclose(dataset_id);
    H5Sclose(dataspace_id);
    H5Fclose(file_id);
}

/* 0 out the grid, except for the left edge */
void initGrid() {
    int row, col;

    for (row=0; row<N_Rows; row++)
        for (col=0; col<N_Cols; col++)
            if (col == 0)
                Grid[row][col] = Init_Temp;
            else
                Grid[row][col] = 0.0;
}

/* process 'I am the average of my neighbors' */
void setNextGrid() {
    int row, col, n_neighbors;;

    for (row=0; row<N_Rows; row++) {
        for (col=0; col<N_Cols; col++) {
            n_neighbors=4;
            Next_Grid[row][col] = 0.0;
            if (row > 0)
                Next_Grid[row][col] += Grid[row-1][col];
            else
                n_neighbors--;
            if (row < N_Rows-1)
                Next_Grid[row][col] += Grid[row+1][col];
            else
                n_neighbors--;
            if (col > 0)
                Next_Grid[row][col] += Grid[row][col-1];
            else
                n_neighbors--;
            if (col < N_Cols-1)
                Next_Grid[row][col] += Grid[row][col+1];
            else
                n_neighbors--;

            Next_Grid[row][col] /= n_neighbors;
        }
    }
}

/* copy the next grid into the current grid */
void advanceGrid() {
    int row, col;

    for (row=0; row<N_Rows; row++) {
        for (col=0; col<N_Cols; col++) {
            Grid[row][col] = Next_Grid[row][col];
        }
    }
}

/* free memory allocated for grids */
void freeMem() {
    int row;
    for (row=N_Rows-1; row>=0; row--)
        free(Next_Grid[row]);
    free(Next_Grid);
    for (row=N_Rows-1; row>=0; row--)
        free(Grid[row]);
    free(Grid);
}

int main(int argc, char **argv) {
    int i;
    char filename[80];

    /* set up */
    parseArgs(argc, argv);
    validateInput();
    allocMem();
    initGrid();

    /* simulate */
    displayGrid(0); 
    writeHDF5(0);
    for (i=0; i<N_Time_Steps; i++) {
        setNextGrid();
        advanceGrid();
        displayGrid(i+1);
        writeHDF5(i+1);
    }

    /* tear down */
    freeMem();
    return 0;
}