/*
 * heat diffusion example, as discussed in textbook.
 *
 * command line arguments are number of points, maximum number of
 *   iterations, convergence threshold, and optional flag "values"
 *   to print final values.
 * timing information goes to stdout, final values to stderr.
 *
 * parallel version with OpenMP, more complex SPMD-like approach.
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <omp.h>

#define LEFTVAL 1.0
#define RIGHTVAL 10.0

void parse_arguments(int argc, char *argv[], 
        int *nx, int *maxsteps, double *threshold, int *print_flag);
void initialize_end_points(double uk[], double ukp1[], int nx);
void initialize_local_sections(double uk[], double ukp1[], 
        int loop_start, int loop_end);
void print_values(double uk[], int nx);

int main(int argc, char *argv[]) {

    int nx;
    int maxsteps;
    double threshold;
    int print_flag;

    double *uk; 
    double *ukp1;
    double *temp; 
    double dx, dt;
    double start_time, end_time;
    double maxdiff, diff;
    int step;

    int nthreads;

    parse_arguments(argc, argv, &nx, &maxsteps, &threshold, &print_flag);

    start_time = omp_get_wtime();

    #pragma omp parallel 
    {
        #pragma omp single 
        {
            nthreads = omp_get_num_threads();
        }
    }
    if ((nx % nthreads) != 0) {
        fprintf(stderr, "Number of threads must evenly divide %d\n", nx);
        exit(EXIT_FAILURE);
    }

    uk = malloc(sizeof(double) * nx);
    ukp1 = malloc(sizeof(double) * nx);
    if (!uk || !ukp1) {
        fprintf(stderr, "Unable to allocate memory\n");
        return EXIT_FAILURE;
    }
    dx = 1.0/nx;
    dt = 0.5*dx*dx;
    maxdiff = threshold;

    /* initialize(uk, ukp1, nx); */
    initialize_end_points(uk, ukp1, nx);
    /* remainder of initialize in parallel section */

    #pragma omp parallel private(diff)
    {
        int step_local; 
        double maxdiff_local;

        int my_id;
        int loop_start;
        int loop_end;
        int i;

        my_id = omp_get_thread_num();
        loop_start = (my_id == 0) ? 1 : my_id * (nx/nthreads);
        loop_end = (my_id == nthreads-1) ? nx-1 : (my_id + 1) * (nx/nthreads);

        initialize_local_sections(uk, ukp1, loop_start, loop_end);
        #pragma omp barrier

        for (step_local = 0; (step_local < maxsteps) && (maxdiff >= threshold);
                ++step_local) {

            /* compute new values */
            /* for (i = 1; i < nx-1; ++i) */
            for (i = loop_start; i < loop_end; ++i) {
                ukp1[i]=uk[i]+ (dt/(dx*dx))*(uk[i+1]-2*uk[i]+uk[i-1]);
            }

            /* check for convergence */
            #pragma omp barrier
            #pragma omp single
            {
                maxdiff = 0.0;
            }
            maxdiff_local = 0.0;
            /* for (i = 1; i < nx-1; ++i) */
            for (i = loop_start; i < loop_end; ++i) {
                diff = fabs(uk[i] - ukp1[i]);
                if (diff > maxdiff_local) maxdiff_local = diff;
            }
            #pragma omp critical
            {
                if (maxdiff_local > maxdiff) maxdiff = maxdiff_local;
            }

            /* "copy" ukp1 to uk by swapping pointers */
            #pragma omp barrier
            #pragma omp single
            {
                temp = ukp1; ukp1 = uk; uk = temp;
            }
        }

        /* save count of steps in shared variable */
        #pragma omp single
        {
            step = step_local;
        }
    }

    end_time = omp_get_wtime();

    if (print_flag) {
        print_values(uk, nx);
    }

    printf("OpenMP program, SPMD-style (%d threads):\n", nthreads);
    printf("nx = %d, maxsteps = %d, threshold = %g\n", nx, maxsteps, threshold);
    if (maxdiff < threshold) {
        printf("converged in %d iterations\n", step);
    }
    else {
        printf("failed to converge in %d iterations, maxdiff = %g\n", 
                step, maxdiff);
    }
    printf("execution time = %g\n", end_time - start_time);

    /* clean up and end */
    return EXIT_SUCCESS;
}

void parse_arguments(int argc, char *argv[], 
        int *nx, int *maxsteps, double *threshold, int *print_flag) {
    char *usage_msg = 
    "usage is %s points max_iterations convergence_threshold [\"values\"]\n";

    if (argc < 4) {
        fprintf(stderr, usage_msg, argv[0]);
        exit(EXIT_FAILURE);
    }
    *nx = 0;
    *maxsteps = 0;
    *threshold = 0.0;
    *nx = atoi(argv[1]);
    *maxsteps = atoi(argv[2]);
    *threshold = atof(argv[3]);
    *print_flag = 0;
    if ((*nx <= 0) || (*maxsteps <= 0) || (*threshold <= 0)) {
        fprintf(stderr, usage_msg, argv[0]);
        exit(EXIT_FAILURE);
    }
    if ((argc > 4) && (strcmp(argv[4], "values") == 0)) {
        *print_flag = 1;
    }
}
void initialize_end_points(double uk[], double ukp1[], int nx) {
    uk[0] = LEFTVAL; uk[nx-1] = RIGHTVAL;
    ukp1[0] = uk[0];
    ukp1[nx-1] = uk[nx-1];
}

void initialize_local_sections(double uk[], double ukp1[], 
        int loop_start, int loop_end) {
    int i;
    for (i = loop_start; i < loop_end; ++i) {
        uk[i] = 0.0;
        ukp1[i] = uk[i];
    }
}

void print_values(double uk[], int nx) {
    int i;
    for (i = 0; i < nx; ++i) {
        fprintf(stderr, "uk[%010d] = %14.10f\n", i, uk[i]);
    }
}