#include <math.h>
#include <stdbool.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>

#include <mpi.h>

#include "grid.h"

/* Should receives be asynchronous */
#ifndef ASYNCH_RECV
#define ASYNCH_RECV 0
#endif


/* MPI communication routines */

// Master sends data.  Others receive
static void broadcast_receive_data(double *data, int count);

// Mechanisms to exchange data between adjacent rows
static void start_send_data(double *data, int count, int process_id, MPI_Request *request);

#if ASYNCH_RECV
/* Asynchronous */
static void start_receive_data(double *data, int count, int process_id, MPI_Request *request);
#else
/* Synchronous */
static void receive_data(double *data, int count, int process_id);
#endif

static void finish_data(MPI_Request *request);


// Find maximum value across processes
static double global_maximum(double local_max, int process_id, int process_count);

grid_t *new_grid(int nrow, int ncol, int process_id, int process_count) {
    int row_count = nrow / process_count;
    int row_start = row_count * process_id;
    int extra = nrow % process_count;
    if (process_id < extra) {
	/* Distribute extra processing among first rows */
	row_count += 1;
	row_start += process_id;
    } else {
	row_start += extra;
    }
    grid_t *g = malloc(sizeof(grid_t));
    g->nrow = nrow;
    g->ncol = ncol;
    g->process_id = process_id;
    g->process_count = process_count;
    g->row_start = row_start;
    g->row_count = row_count;
    if (row_count == 0) {
	g->data = NULL;
	g->ndata = NULL;
    } else {
	g->data = calloc((row_count+2) * (ncol+2), sizeof(double));
	g->ndata = calloc((row_count+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->bcount = nrow*2 + ncol*2;
    // Trick: Allocate all 4 as contiguous storage
    bc->bdata = calloc(bc->bcount, sizeof(double));
    bc->northRow =   bc->bdata;
    bc->eastColumn = bc->bdata+nrow;
    bc->southRow =   bc->bdata+nrow+ncol;
    bc->westColumn = bc->bdata+nrow+ncol+nrow;
    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;
    default: // Do nothing
	break;
    }
    return bc;
}

void free_boundary(boundary_t *bc) {
    free(bc->bdata);
    free(bc);
}

void initialize_grid(grid_t *g, boundary_t *bc) {
    int r, c;
    int process_id = g->process_id;
    if (process_id == 0) {
	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];
	}
    }
    if (process_id == g->process_count - 1) {
	for (c = 0; c < g->ncol; c++) {
	    g->data[GINDEX(g, g->row_count, c)] = bc->southRow[c];
	    g->ndata[GINDEX(g, g->row_count, c)] = bc->southRow[c];
	}
    }

    for (r = 0; r < g->row_count; r++) {
	for (c = 0; c < g->ncol; c++) {
	    int gr = g->row_start + r;
	    g->data[GINDEX(g, r,-1)] = bc->westColumn[gr];
	    g->ndata[GINDEX(g, r,-1)] = bc->westColumn[gr];
	    g->data[GINDEX(g, r,g->ncol)] = bc->eastColumn[gr];
	    g->ndata[GINDEX(g, r,g->ncol)] = bc->eastColumn[gr];
	}
    }
}

/* Exchange rows with neighbors to north and south */
static void exchange_rows(grid_t *g) {
    if (g->row_count == 0)
	return;
    int process_id = g->process_id;
    bool north_neighbor = process_id > 0;
    bool south_neighbor = process_id < g->process_count - 1;
    MPI_Request north_request;
    MPI_Request south_request;
#if ASYNCH_RECV
    MPI_Request north_recv_request;
    MPI_Request south_recv_request;
#endif

    if (north_neighbor) {
	double *north_boundary = &g->data[GINDEX(g, 0, 0)];
	start_send_data(north_boundary, g->ncol, process_id - 1, &north_request);
    }
    if (south_neighbor) {
	double *south_boundary = &g->data[GINDEX(g, g->row_count-1, 0)];
	start_send_data(south_boundary, g->ncol, process_id + 1, &south_request);
    }
    if (north_neighbor) {
	double *north_ghost = &g->data[GINDEX(g, -1, 0)];
#if ASYNCH_RECV
	start_receive_data(north_ghost, g->ncol, process_id - 1, &north_recv_request);
#else
	receive_data(north_ghost, g->ncol, process_id - 1);
#endif
    }
    if (south_neighbor) {
	double *south_ghost = &g->data[GINDEX(g, g->row_count, 0)];
#if ASYNCH_RECV
	start_receive_data(south_ghost, g->ncol, process_id + 1, &south_recv_request);
#else
	receive_data(south_ghost, g->ncol, process_id + 1);
#endif
    }
    if (north_neighbor) {
	finish_data(&north_request);
#if ASYNCH_RECV
	finish_data(&north_recv_request);
#endif
    }
    if (south_neighbor) {
	finish_data(&south_request);
#if ASYNCH_RECV
	finish_data(&south_recv_request);
#endif
    }
}

/* 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;
}

/* Perform one step of grid solving */
static double step_grid(grid_t *g) {
    int row_count = g->row_count;
    int process_id = g->process_id;

    double local_maxdiff = 0.0;
    int r, c;
    for (r = 0; r < row_count; 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)]);
	    local_maxdiff = fmax(local_maxdiff, diff);
	}
    }

    /* Swap old and new data */
    double *tdata = g->data;
    g->data = g->ndata;
    g->ndata = tdata;

    exchange_rows(g);

    double maxdiff = global_maximum(local_maxdiff, process_id, g->process_count);
    return maxdiff;
}

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

int solve(int nrow, int ncol, char mode, double epsilon, int maxSteps, bool verbose, int process_id, int process_count) {

    grid_t *g = new_grid(nrow, ncol, process_id, process_count);
    char local_mode = process_id == 0 ? mode : 'x';
    boundary_t *bc = create_boundary(nrow, ncol, local_mode);
    broadcast_receive_data(bc->bdata, bc->bcount);
    initialize_grid(g, bc);
    free_boundary(bc);
    int steps = run_grid(g, maxSteps, epsilon, verbose);
    free_grid(g);
    return steps;
}


/************** MPI Routines ******************/

static void broadcast_receive_data(double *data, int count) {
    /* Master process will send.  Others will receive */
    MPI_Bcast(data, count, MPI_DOUBLE, 0, MPI_COMM_WORLD);
}

// Exchange data between adjacent rows
static void start_send_data(double *data, int count, int process_id, MPI_Request *request) {
    MPI_Isend(data, count, MPI_DOUBLE, process_id, 0, MPI_COMM_WORLD, request);
}

#if ASYNCH_RECV
static void start_receive_data(double *data, int count, int process_id, MPI_Request *request) {
    MPI_Irecv(data, count, MPI_DOUBLE, process_id, 0, MPI_COMM_WORLD, request);
}
#endif

static void finish_data(MPI_Request *request) {
    MPI_Wait(request, MPI_STATUS_IGNORE);
}

#if !ASYNCH_RECV
static void receive_data(double *data, int count, int process_id) {
    MPI_Recv(data, count, MPI_DOUBLE, process_id, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
}
#endif

#define USEREDUCE 0
// Find maximum value across processes
#if USEREDUCE
static double global_maximum(double local_max, int process_id, int process_count) {
    double result = 0.0;
    MPI_Allreduce(&local_max, &result, 1, MPI_DOUBLE, MPI_MAX, MPI_COMM_WORLD);
    return result;
}
#else
static double global_maximum(double local_max, int process_id, int process_count) {
    double maxval = local_max;
    if (process_id == 0) {
	int pid;
	double otherval = 0.0;
	for (pid = 1; pid < process_count; pid++) {
	    MPI_Recv(&otherval, 1, MPI_DOUBLE, pid, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
	    maxval = fmax(maxval, otherval);
	}
	broadcast_receive_data(&maxval, 1);
    } else {
	MPI_Send(&maxval, 1, MPI_DOUBLE, 0, 0, MPI_COMM_WORLD);
	broadcast_receive_data(&maxval, 1);
    }
    return maxval;
}
#endif