#include <math.h>
#include <stdbool.h>
#include <stdlib.h>
#include <stdio.h>
#include "grid.h"


grid_t *new_grid(nrow, ncol) {
    grid_t *g = malloc(sizeof(grid_t));
    g->nrow = nrow;
    g->ncol = ncol;
    g->data = calloc((nrow+2) * (ncol+2), sizeof(double));
    g->ndata = calloc((nrow+2) * (ncol+2), sizeof(double));
    return g;
}

void free_grid(grid_t *g) {
    free(g->data);
    free(g);
}

boundary_t *create_boundary(int nrow, int ncol, char mode) {
    boundary_t *bc = malloc(sizeof(boundary_t));
    int r, c;
    double dr = 1.0/(nrow+1);
    double dc = 1.0/(ncol+1);
    bc->nrow = nrow;
    bc->ncol = ncol;
    bc->northRow = calloc(ncol, sizeof(double));
    bc->eastColumn = calloc(nrow, sizeof(double));
    bc->southRow = calloc(ncol, sizeof(double));
    bc->westColumn = calloc(nrow, sizeof(double));
    switch(mode) {
    case 'h': // Horizontal
	for (c = 0; c < ncol; c++) {
	    bc->northRow[c] = dc * (c+1);
	    bc->southRow[c] = dc * (c+1);
	}
	for (r = 0; r < nrow; r++) {
	    bc->westColumn[r] = 0.0;
	    bc->eastColumn[r] = 1.0;
	}
	break;
    case 'd': // Diagonal
	for (c = 0; c < ncol; c++) {
	    bc->northRow[c] = dc * (c+1);
	    bc->southRow[c] = 1.0 - dc * (c+1);
	}
	for (r = 0; r < nrow; r++) {
	    bc->westColumn[r] = dr * (r+1);
	    bc->eastColumn[r] = 1.0 - dr * (r+1);
	}
	break;
    case 'c': // Corner
	for (c = 0; c < ncol; c++) {
	    bc->northRow[c] = 1.0 - dc * (c+1);
	    bc->southRow[c] = 0.0;
	}
	for (r = 0; r < nrow; r++) {
	    bc->westColumn[r] = 1.0 - dr * (r+1);
	    bc->eastColumn[r] = 0.0;
	}
	break;
    case 'r': // Random
	for (c = 0; c < ncol; c++) {
	    bc->northRow[c] = (double) random()/RAND_MAX;
	    bc->southRow[c] = (double) random()/RAND_MAX;
	}
	for (r = 0; r < nrow; r++) {
	    bc->westColumn[r] = (double) random()/RAND_MAX;
	    bc->eastColumn[r] = (double) random()/RAND_MAX;
	}
	break;
    }
    return bc;
}

void free_boundary(boundary_t *bc) {
    free(bc->northRow);
    free(bc->eastColumn);
    free(bc->southRow);
    free(bc->westColumn);
    free(bc);
}

void initialize_grid(grid_t *g, boundary_t *bc) {
    int r, c;
    for (r = 0; r < g->nrow; r++) {
	g->data[GINDEX(g, r,-1)] = bc->westColumn[r];
	g->ndata[GINDEX(g, r,-1)] = bc->westColumn[r];
	g->data[GINDEX(g, r,g->ncol)] = bc->eastColumn[r];
	g->ndata[GINDEX(g, r,g->ncol)] = bc->eastColumn[r];
    }
    for (c = 0; c < g->ncol; c++) {
	g->data[GINDEX(g, -1, c)] = bc->northRow[c];
	g->ndata[GINDEX(g, -1, c)] = bc->northRow[c];
	g->data[GINDEX(g, g->nrow, c)] = bc->southRow[c];
	g->ndata[GINDEX(g, g->nrow, c)] = bc->southRow[c];
    }
}

/* Next state computation for grid position r, c */
static inline double new_state(grid_t *g, int r, int c) {
    double ov = g->data[GINDEX(g, r,   c)];
    double nv = g->data[GINDEX(g, r-1, c)];
    double ev = g->data[GINDEX(g, r, c+1)];
    double sv = g->data[GINDEX(g, r+1, c)];
    double wv = g->data[GINDEX(g, r,  c-1)];
    return 0.25 * CONDUCTIVITY * (nv+ev+sv+wv) + (1-CONDUCTIVITY) * ov;
}

static double step_grid(grid_t *g) {
    double maxdiff = 0.0;
    int r, c;
#if OMP
    #pragma omp parallel for schedule(static) reduction(max:maxdiff)
#endif
    for (r = 0; r < g->nrow; r++) {
	for (c = 0; c < g->ncol; c++) {
	    double nval = new_state(g, r, c);
	    g->ndata[GINDEX(g, r,c)] = nval;
	    double diff = fabs(nval - g->data[GINDEX(g, r,   c)]);
	    maxdiff = fmax(maxdiff, diff);
	}
    }
    double *tdata = g->data;
    g->data = g->ndata;
    g->ndata = tdata;
    return maxdiff;
}

int run_grid(grid_t *g, int maxSteps, double epsilon, bool verbose) {
    double diff = 0.0;
    int s;
    for (s = 0; s < maxSteps; s++) {
	diff = step_grid(g);
	if (verbose)
	    printf("Step %d.  Max. difference = %.6f\n", s, diff);
	if (diff < epsilon) {
	    printf("Terminated after %d steps.  Max. difference = %.6f\n", s, diff);
	    return s;
	}
    }
    printf("Failed to terminate after %d steps.  Max. difference = %.6f\n", maxSteps, diff);
    return maxSteps;
}

static void show_row(grid_t *g, int r) {
    int c;
    if (r == -1 || r == g->nrow) {
	printf("      ");
	for (c = 0; c < g->ncol; c++)
	    printf("%.3f ", g->data[GINDEX(g, r,c)]);
	printf("      \n");
    } else {
	for (c = -1; c < g->ncol+1; c++)
	    printf("%.3f ", g->data[GINDEX(g, r,c)]);
	printf("\n");
    }
}

void show_grid(grid_t *g) {
    int r;
    for (r = -1; r < g->nrow+1; r++)
	show_row(g, r);
}

int solve(int nrow, int ncol, char mode, double epsilon, int maxSteps, bool verbose) {
    grid_t *g = new_grid(nrow, ncol);
    boundary_t *bc = create_boundary(nrow, ncol, mode);
    initialize_grid(g, bc);
    int steps = run_grid(g, maxSteps, epsilon, verbose);
    if (verbose)
	show_grid(g);
    free_grid(g);
    free_boundary(bc);
    return steps;
}