/*
 * Red-Black Successive Over-Relaxation.
 * Sequential implementation for MPI assignment
 * Programming Large-Scale Parallel Systems 2022.
 * 
 * Original by Rob van Nieuwpoort, changes by Henk Dreuning.
 *
 */

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <string.h>
#include <sys/time.h>

#define TOLERANCE 0.00002        /* Termination criterion. */
#define MAGIC 0.8                /* Magic factor. */

int is_even(int x) {
    return !(x & 1);
}

/* 'plus' around G[row][col] */
double stencil(double** G, int row, int col) {
    return (G[row-1][col] + G[row+1][col] +
        G[row][col-1] + G[row][col+1] ) / 4.0;
}

void allocate_grid(double ***Gptr, int N) {
    /* Allocate array of pointers (1 per row). */
    double **G = (double **)malloc(N * sizeof(double *));
    if (G == NULL) {
        fprintf(stderr, "Could not allocate memory for the grid. Exiting...\n");
        exit(1);
    }

    /* Allocate rows. */
    for (int i = 0; i < N; i++) {
        G[i] = (double *)malloc(N * sizeof(double));
        if ( G[i] == NULL ) {
            fprintf(stderr, "Could not allocate memory for the grid. Exiting...\n");
            exit(1); 
        }
    }

    *Gptr = G;
}

/* Initialize the grid. */
void initialize_grid(double **G, int N) {
    for (int i = 0; i < N; i++){
        for (int j = 0; j < N; j++){
            if (i == 0) G[i][j] = 4.56;
            else if (i == N-1) G[i][j] = 9.85;
            else if (j == 0) G[i][j] = 7.32;
            else if (j == N-1) G[i][j] = 6.88;
            else G[i][j] = 0.0;
        }
    }
}

void print_grid(double **G, int N) {
    for (int i = 1; i < N-1; i++) {
        for (int j = 1; j < N-1; j++) {
            printf("%10.3f ", G[i][j]);
        }
        printf("\n");
    }
}


int main (int argc, char *argv[]) {
    int N;                              /* Grid size. */
    double new_val, r, omega;           /* Variables related to the relaxation factor. */
    double stopdiff, maxdiff, diff;     /* Difference between old and new values and
                                         * termination criterion. */
    double **G;                         /* The grid. */
    int iteration;                      /* Counters. */
    struct timeval start;               /* Timing variables. */
    struct timeval end;
    double elapsed_time;
    int verbose = 0;

    N = 10000;

    for (int i = 1; i < argc; i++) {
        if(!strcmp(argv[i], "-verbose"))
            verbose = 1;
        else
            N = atoi(argv[i]);
    }

    fprintf(stderr, "Running RB-SOR with grid size %d x %d\n", N, N);

    /* Add a row/column on the top, bottom, left and right of the grid.
     * The computation will only be performed on the inner N*N elements. */
    N += 2;

    /* Determine the relaxation factor and termination criterion. */
    r        = 0.5 * ( cos( M_PI / N ) + cos( M_PI / N ) );
    omega    = 2.0 / ( 1 + sqrt( 1 - r * r ) );
    stopdiff = TOLERANCE / ( 2.0 - omega );
    omega   *= MAGIC;

    /* Allocate and initialize the grid. */
    allocate_grid(&G, N);
    initialize_grid(G, N);

    /* Start timing. */
    if(gettimeofday(&start, 0) != 0) {
        fprintf(stderr, "Could not start timer. Exiting...\n");
        exit(1);
    }

    /* Perform RB-SOR. */
    iteration = 0;
    do {
        maxdiff = 0.0;
        for (int phase = 0; phase < 2; phase++){
            for (int i = 1; i < N-1; i++){
                for (int j = 1 + (is_even(i) ^ phase); j < N-1; j += 2){
                    new_val = stencil(G,i,j);
                    diff = fabs(new_val - G[i][j]);
                    if (diff > maxdiff)
                        maxdiff = diff;
                    G[i][j] = G[i][j] + omega * (new_val - G[i][j]);
                }
            }
        }
        iteration++;
    } while (maxdiff > stopdiff);

    /* Stop timing. */
    if(gettimeofday(&end, 0) != 0) {
        fprintf(stderr, "Could not stop timer. Exiting...\n");
        exit(1);
    }

    elapsed_time = (end.tv_sec + (end.tv_usec / 1000000.0)) -
        (start.tv_sec + (start.tv_usec / 1000000.0));

    fprintf(stderr, "SOR took %10.3f seconds\n", elapsed_time);
    printf("Used %5d iterations, maximum difference is %10.6f, allowed difference is %10.6f\n",
        iteration, maxdiff, stopdiff);

    if (verbose)
        print_grid(G, N);

    return 0;
}