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

#define VERBOSE

double **y = NULL;
int tasks, rank, steps, start;
#ifdef VERBOSE
FILE *file;
#endif

/*
 * Connect the string at both ends by sending the calculated border values and
 * receiving the outer borders calculated by the adjacent nodes
 */
void interact(int state) {
    MPI_Status status;

    // Left border
    if( rank ) {
        MPI_Sendrecv(y[state] + 1, 1, MPI_DOUBLE, rank - 1, 0, y[state], 1,
                MPI_DOUBLE, rank - 1, MPI_ANY_TAG, MPI_COMM_WORLD, &status);
    }

    // Right border
    if( rank < tasks - 1 ) {
        MPI_Sendrecv(y[state] + steps, 1, MPI_DOUBLE, rank + 1, 0,
                y[state] + steps + 1, 1, MPI_DOUBLE, rank + 1, MPI_ANY_TAG,
                MPI_COMM_WORLD, &status);
    }
}

/*
 * Initialize the left- and right values of the string section. The outer
 * values are sent by other nodes.
 */
void init_borders() {
    // The string is attached at both ends
    if( !rank )
        y[2][1] = y[1][1] = y[0][1];
    else if( rank == tasks - 1 )
        y[2][steps] = y[1][steps] = y[0][steps];

    interact(0);
}

/*
 * Fill the assigned string section with a sinus shape.
 */
void init_sinus(double l, double dx, double n) {
    for( int i = 0; i < steps; i++ )
        y[0][i + 1] = sin(n * M_PI * (i + start) * dx / l);

    init_borders();
}

/*
 * Fill the assigned string section with a plucked shape.
 */
void init_plucked(double l, double dx, double xp) {
    double x;

    for( int i = 0; i < steps; i++ ) {
        x = (i + start) * dx;
        y[0][i + 1] = x < xp ? x / xp : (l - x) / (l - xp);
    }

    init_borders();
}

#ifdef VERBOSE
#define PRINT_FORMAT "%e"

/*
 * Print all calculated values in the string section of one of the three
 * currently known time states.
 */
void print(int index) {
    for( int i = 1; i <= steps; i++ ) {
        if( i > 1 )
            fprintf(file, " ");

        fprintf(file, PRINT_FORMAT, y[index][i]);
    }

    fprintf(file, "\n");
}
#endif

/*
 * Calculate and optionally print the string section states for the specified
 * number of time steps.
 */
void calculate_steps(int time, double dx, double tau) {
    int i, t, prev = 0, current = 1, next = 2,
        s = rank ? 1 : 2, e = rank == tasks - 1 ? steps - 1 : steps;

    // Calculate the position over the entire string at time dt using the
    // position at t = 0 and the information that y(x, -dt) == y(x, dt)
    for( i = s; i <= e; i++ ) {
        y[current][i] = y[0][i] + .5 * tau * tau * (y[0][i - 1] - 2 * y[0][i]
                  + y[0][i + 1]);
    }

    interact(current);

    #ifdef VERBOSE
    // Print x values (which are the same for every y-series)
    for( i = 0; i < steps; i++ ) {
        if( i )
            fprintf(file, " ");

        fprintf(file, PRINT_FORMAT, (i + start) * dx);
    }

    fprintf(file, "\n");

    // Print init states
    print(prev);
    print(current);
    #endif

    // Iterate over the length of the string for each time interval step
    for( t = 2; t < time; t++ ) {
        for( i = s; i <= e; i++ ) {
            y[next][i] = 2 * y[current][i] - y[prev][i] + tau * tau
                * (y[current][i-1] - 2 * y[current][i] + y[current][i+1]);
        }

        interact(next);

        #ifdef VERBOSE
        print(next);
        #endif

        prev = current;
        current = next;
        next = (next + 1) % 3;
    }
}

/*
 * Main method, initiates the program.
 */
int main(int argc, char **argv) {
    int error, time, i, total_steps;
    void (*init_method)(double, double, double);
    double l, tau, constant, dx;

    // Parse arguments
    if( argc < 7 ) {
        printf("Usage: %s INIT_METHOD TIME_STEPS LENGTH DX TAU N|XP\n",
               argv[0]);
        return EXIT_FAILURE;
    }

    if( !strcmp(argv[1], "sinus") ) {
        init_method = init_sinus;
    } else if( !strcmp(argv[1], "plucked") ) {
        init_method = init_plucked;
    } else {
        printf("Unknown initialization method '%s'.\n", argv[1]);
        return EXIT_FAILURE;
    }

    time = atoi(argv[2]);
    l = atof(argv[3]);
    dx = atof(argv[4]);
    tau = atof(argv[5]);
    constant = atof(argv[6]);

    // Initialize and setup MPI
    if( (error = MPI_Init(&argc, &argv)) != MPI_SUCCESS ) {
        printf("MPI_Init failed (error %d)\n", error);
        MPI_Abort(MPI_COMM_WORLD, error);
        return EXIT_FAILURE;
    }

    MPI_Comm_size(MPI_COMM_WORLD, &tasks);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);

    #ifdef VERBOSE
    // Open file for output
    char filename[9];
    sprintf(filename, "out%d.txt", rank);
    file = fopen(filename, "w");
    #endif

    // Calculate the indices of the string to calculate
    total_steps = (int)ceil(l / dx) + 1;
    steps = total_steps / tasks;
    start = steps * rank;

    // Fill in the gap caused by rounding errors in the last node. Also,
    // allocate the left and right outer border values that are calculated by
    // other nodes
    if( rank == tasks - 1 )
        steps += total_steps % steps;

    // Allocate a 3xN two-dimensional area for this part of the string that
    // functions as a cyclic buffer for three states of the string (previous,
    // current, next)
    y = (double **)malloc(3 * sizeof(double *));

    for( i = 0; i < 3; i++ )
        y[i] = (double *)malloc((steps + 2) * sizeof(double));

    // Initialize the string values (sinus/plucked) and calculate the specified
    // number of steps
    (*init_method)(l, dx, constant);
    calculate_steps(time, dx, tau);

    // Cleanup
    for( i = 0; i < 3; i++ )
        free(y[i]);

    free(y);

    MPI_Finalize();
    #ifdef VERBOSE
    fclose(file);
    #endif

    return EXIT_SUCCESS;
}
