/*
 * C program for Conway's ``game of life''.
 *
 * Input is from command line and (possibly) a file:
 * 
 * Command-line arguments specify 
 *   (*) input source (more below)
 *   (*) number of steps
 *   (*) how often to print new board configuration (P means print results 
 *       every P steps)
 *   (*) (optionally) output file for new board configurations
 * Input source can either be the name of an input file containing an 
 *   initial configuration or the keyword "random" and the following
 *   parameters:
 *   (*) size of board
 *   (*) fraction of cells that should initially be "live" 
 *   (*) seed for generating random sequence
 * Output is the number of "live" cells initially and after each step,
 *   and optionally (if an output file is specified) new board configurations
 *
 * If an input file is specified, it contains a representation of the initial 
 *  board configuration:  N (size of board) and N*N values (each 0 or 1).
 *
 * Parallel version using MPI (and distributing the board among processes)..
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdbool.h>

#include "timer.h"

#include <mpi.h>

/* 
 * data structure for two-dimensional array of "short" 
 * (what we really want here is bool, but MPI doesn't seem to support
 * sending those, so "short" seems like an okay substitute)
 */
typedef short array_type;
#define MPI_ARRAY_TYPE MPI_SHORT
typedef struct twoD_array {
    long rows;
    long cols;
    array_type ** elems;
} twoD_array_t;


/* function declarations (comments with code below) */
twoD_array_t * build_array(twoD_array_t * a, long rows, long cols);
void free_array(twoD_array_t * a);
bool read_board(FILE* infile, twoD_array_t *board);
bool random_board(long size, double fraction, long seed, twoD_array_t *board);
void update_board(twoD_array_t *board, twoD_array_t *new_board);
void print_board_p0(FILE* outfile, twoD_array_t *board);
void print_board_notp0(twoD_array_t *board);
void clear_border(twoD_array_t *board);
long live_cells(twoD_array_t *board);
long rows_per_process(long size);
long local_rows(int proc_id, long size);
long local_start_row(int proc_id, long size);
long local_end_row(int proc_id, long size);
int process_owning_row(int row, long size);

/* global variables */
int nprocs;
int myid;

/* message tags */
#define INITIALIZE_TAG  1
#define EXCHANGE_TAG    2
#define PRINT_TAG       3

/* main */
int main(int argc, char* argv[]) {
    if (MPI_Init(&argc, &argv) != MPI_SUCCESS) {
        fprintf(stderr, "MPI initialization error\n"); exit(EXIT_FAILURE);
    }
    MPI_Comm_size(MPI_COMM_WORLD, &nprocs);
    MPI_Comm_rank(MPI_COMM_WORLD, &myid);

    long steps, print_interval, size;
    FILE *outfile = NULL;
    bool have_outfile = false;

    twoD_array_t board1, board2;
    twoD_array_t *board = &board1;
    twoD_array_t *new_board = &board2;

    double start_time, init_done_time, compute_done_time, end_time;
    long live_count;

    char* usage_fmt = 
        "usage:  %s input_args num_steps print_interval [outfile]\n"
        "  (input_args is infile or 'random' boardsize fraction seed)\n";
    char* end_ptr_for_strtol;

    start_time = get_time();

    if (argc < 4) {
        if (myid == 0) {
            fprintf(stderr, usage_fmt, argv[0]);
        }
        MPI_Finalize(); exit(EXIT_FAILURE);
    }

    /* process command-line arguments and initialize board */

    int steps_index; /* index of "steps" argument */
    if (strcmp(argv[1], "random") != 0) {

        /* from file */
        FILE* infile;
        infile = fopen(argv[1], "r");
        if (infile == NULL) {
            fprintf(stderr, "unable to open input file %s\n", argv[1]);
            /* MPI_Abort here since we could maybe fail in only one process */
            MPI_Abort(MPI_COMM_WORLD, EXIT_FAILURE);
        }
        steps_index = 2;
        if (!read_board(infile, board)) {
            /* MPI_Abort here since we could maybe fail in only one process */
            MPI_Abort(MPI_COMM_WORLD, EXIT_FAILURE);
        }
        size = board->cols-2;
        fclose(infile);
        if (myid == 0) {
            printf("Input:  board read from file %s\n", argv[1]);
        }
    } else {

        /* with randomly-generated data */
        long seed;
        double fraction;

        if (argc < 7) {
            if (myid == 0) {
                fprintf(stderr, usage_fmt, argv[0]);
            }
            MPI_Finalize(); exit(EXIT_FAILURE);
        }
        size = strtol(argv[2], &end_ptr_for_strtol, 10);
        if ((*end_ptr_for_strtol != '\0') || (size <= 0)) {
            if (myid == 0) {
                fprintf(stderr, usage_fmt, argv[0]);
            }
            MPI_Finalize(); exit(EXIT_FAILURE);
        }
        fraction = strtod(argv[3], &end_ptr_for_strtol);
        if ((*end_ptr_for_strtol != '\0') || (fraction < 0) || (fraction > 1)) {
            if (myid == 0) {
                fprintf(stderr, usage_fmt, argv[0]);
            }
            MPI_Finalize(); exit(EXIT_FAILURE);
        }
        seed = strtol(argv[4], &end_ptr_for_strtol, 10);
        if (*end_ptr_for_strtol != '\0') {
            if (myid == 0) {
                fprintf(stderr, usage_fmt, argv[0]);
            }
            MPI_Finalize(); exit(EXIT_FAILURE);
        }
        steps_index = 5;
        if (!random_board(size, fraction, seed, board)) {
            /* MPI_Abort here since we could maybe fail in only one process */
            MPI_Abort(MPI_COMM_WORLD, EXIT_FAILURE);
        }

        if (myid == 0) {
            printf("Input:  "
                    "'random' board of size %ld with fraction %f seed %ld\n",
                    size, fraction, seed);
        }
    }

    steps = strtol(argv[steps_index], &end_ptr_for_strtol, 10);
    if (*end_ptr_for_strtol != '\0') {
        if (myid == 0) {
            fprintf(stderr, usage_fmt, argv[0]);
        }
        MPI_Finalize(); exit(EXIT_FAILURE);
    }
    print_interval = strtol(argv[steps_index+1], &end_ptr_for_strtol, 10);
    if (*end_ptr_for_strtol != '\0') {
        if (myid == 0) {
            fprintf(stderr, usage_fmt, argv[0]);
        }
        MPI_Finalize(); exit(EXIT_FAILURE);
    }

    if (argc >= steps_index+3) {
        if (myid == 0) {
            outfile = fopen(argv[steps_index+2], "w");
            if (outfile == NULL) {
                fprintf(stderr, "cannot open output file %s\n", 
                        argv[steps_index+2]);
                /* MPI_Abort here since we could fail in only one process */
                MPI_Abort(MPI_COMM_WORLD, EXIT_FAILURE);
            }
        }
        have_outfile = true;
    }

    live_count = live_cells(board);
    if (myid == 0) {
        printf("\n");
        printf("%ld live cells initially\n", live_count);
    }

    /* create "new board" and clear borders */
    if (build_array(new_board, local_rows(myid, size)+2, size+2) == NULL) {
        fprintf(stderr, "unable to allocate space for board of size %ld\n",
                size);
        /* MPI_Abort here since we could fail in only one process */
        MPI_Abort(MPI_COMM_WORLD, EXIT_FAILURE);
    }
    clear_border(new_board);

    /* print initial configuration */
    if (have_outfile) {
        if (myid == 0) {
            fprintf(outfile, "Initial board\n\n");
            print_board_p0(outfile, board);
            fprintf(outfile, "\n");
        }
        else {
            print_board_notp0(board);
        }
    }

    init_done_time = get_time();
    compute_done_time = init_done_time; /* mainly to pacify compiler */

    /* loop to update board and print */
    for (long step = 1; step <= steps; ++step) {
        /* update (results in new_board) */
        update_board(board, new_board);
        compute_done_time = get_time();
        /* print */
        if (have_outfile && ((step % print_interval) == 0)) {
            if (myid == 0) {
                fprintf(outfile, "Board after step %ld\n\n", step);
                print_board_p0(outfile, new_board);
                fprintf(outfile, "\n");
            }
            else {
                print_board_notp0(new_board);
            }
        }
        live_count = live_cells(new_board);
        if (myid == 0) {
            printf("%ld live cells after step %ld\n", live_count, step);
        }
        /* swap old and new boards */
        {
            twoD_array_t *temp = board;
            board = new_board;
            new_board = temp;
        }
    }

    end_time = get_time();

    /* print timing information:
     * writes this to standard error so it can be easily separated from
     * rest of output (which could be long)
     */
    if (myid == 0) {
        printf("\nTimes for MPI program with %d processes, "
                "board size %ld, %ld steps, print interval %ld:\n",
                nprocs, size, steps, print_interval);
        printf("Total %g seconds\n", end_time - start_time);
        printf("Not counting initialization or final printing %g seconds\n",
                compute_done_time - init_done_time);
    }

    /* tidy up and return */
    free_array(board);
    free_array(new_board);
    MPI_Finalize();
    return EXIT_SUCCESS;
}

/* 
 * constructs twoD_array structure.  returns NULL if unable to allocate
 * space for elements, pointer to structure otherwise.
 */
twoD_array_t * build_array(twoD_array_t * a, long rows, long cols) {
    array_type * temp;
    a->rows = rows;
    a->cols = cols;
    if ((a->elems = malloc(rows * sizeof(a->elems[0]))) == NULL) {
        return NULL;
    }
    if ((temp = malloc(rows * cols * sizeof(temp[0]))) == NULL) {
        free (a->elems);
        return NULL;
    }
    for (long row = 0; row < rows; ++row, temp+=cols) {
        a->elems[row] = temp;
    }
    return a;
}

/* frees space pointed to by twoD_array structure */
void free_array(twoD_array_t * a) {
    free(a->elems[0]);
    free(a->elems);
}

/*
 * sets unused "edge" cells to 0 
 */
void clear_border(twoD_array_t *board) {
    for (long c = 0; c < board->cols; ++c) {
        board->elems[0][c] = 0;
        board->elems[board->rows-1][c] = 0;
    }
    for (long r = 0; r < board->rows; ++r) {
        board->elems[r][0] = 0;
        board->elems[r][board->cols-1] = 0;
    }
}

/*
 * reads initial configuration from infile 
 * returns 0 if all is well, otherwise prints error message and returns
 *   a positive non-zero value.
 * all processes read the whole file and store only "their" values.
 */
bool read_board(FILE* infile, twoD_array_t *board) {

    long size;
    if (fscanf(infile, "%ld", &size) != 1) {
        fprintf(stderr, "unable to read size of board\n");
        return false;
    }

    long start_row = local_start_row(myid, size);
    long end_row = local_end_row(myid, size);
    int temp;

    if (build_array(board, local_rows(myid, size)+2, size+2) == NULL) {
        fprintf(stderr, "unable to allocate space for board of size %ld\n",
                size);
        return false;
    }
    for (long i = 1; i <= size; ++i) {
        for (long j = 1; j <= size; ++j) {
            if (fscanf(infile, "%d", &temp) != 1) {
                fprintf(stderr, "unable to read values for board\n");
                return false;
            }
            if ((temp == 0) || (temp == 1)) {
                if ((start_row <= i) && (i < end_row)) {
                    board->elems[i-start_row+1][j] = temp;
                }
            }
            else {
                fprintf(stderr, "unable to read values for board\n");
                return false;
            }
        }
    }
    clear_border(board);
    return true;
}

/*
 * generates random board configuration for given size, fraction, and seed.
 * returns true if all is well, otherwise prints error message and returns
 *   false.
 *
 * all processes generate the full sequence of numbers and store only 
 *   "their" values.
 */
bool random_board(long size, double fraction, long seed, twoD_array_t *board) {

    long start_row = local_start_row(myid, size);
    long end_row = local_end_row(myid, size);

    if (build_array(board, local_rows(myid, size)+2, size+2) == NULL) {
        fprintf(stderr, "unable to allocate space for board of size %ld\n",
                size);
        return false;
    }

    srand(seed);
    for (long i = 1; i <= size; ++i) {
        for (long j = 1; j <= size; ++j) {
            bool live = ((double) rand() / (double) (RAND_MAX)) < fraction;
            if ((start_row <= i) && (i < end_row)) {
                board->elems[i-start_row+1][j] = live ? 1 : 0;
            }
        }
    }
    clear_border(board);
    return true;
}

/*
 * updates board configuration
 */
void update_board(twoD_array_t *board, twoD_array_t *new_board) {
    long size;
    MPI_Status status;
    MPI_Request req_recv_above, req_recv_below, req_send_above, req_send_below;

    size = board->cols-2;

    /* exchange information with neighbors -- using asynchronous
     * communcation to avoid possible deadlock if messages are too big */
    /* initiate communication */
    if (myid != 0) {
        /* receive bottom row from neighbor "above" */
        MPI_Irecv(&(board->elems[0][1]), 
                size, MPI_ARRAY_TYPE, myid-1, EXCHANGE_TAG, MPI_COMM_WORLD,
                &req_recv_above);
    }
    if (myid != (nprocs-1)) {
        /* receive top row from neighbor "below" */
        MPI_Irecv(&(board->elems[local_rows(myid, size)+1][1]), 
                size, MPI_ARRAY_TYPE, myid+1, EXCHANGE_TAG, MPI_COMM_WORLD,
                &req_recv_below);
    }
    if (myid != 0) {
        /* send top row to neighbor "above" */
        MPI_Isend(&(board->elems[1][1]), 
                size, MPI_ARRAY_TYPE, myid-1, EXCHANGE_TAG, MPI_COMM_WORLD,
                &req_send_above);
    }
    if (myid != (nprocs-1)) {
        /* send bottom row to neighbor "below" */
        MPI_Isend(&(board->elems[local_rows(myid, size)][1]), 
                size, MPI_ARRAY_TYPE, myid+1, EXCHANGE_TAG, MPI_COMM_WORLD,
                &req_send_below);
    }
    /* wait for communication to complete */
    if (myid != 0) {
        /* receive bottom row from neighbor "above" */
        MPI_Wait(&req_recv_above, &status);
    }
    if (myid != (nprocs-1)) {
        /* receive top row from neighbor "below" */
        MPI_Wait(&req_recv_below, &status);
    }
    if (myid != 0) {
        /* send top row to neighbor "above" */
        MPI_Wait(&req_send_above, &status);
    }
    if (myid != (nprocs-1)) {
        /* send bottom row to neighbor "below" */
        MPI_Wait(&req_send_below, &status);
    }

    /* update board */
    for (int i = 1; i <= board->rows-2; ++i) {
        for (int j = 1; j <= board->cols-2; ++j) {
            int nbrs = 0;
            nbrs += board->elems[i-1][j-1];
            nbrs += board->elems[i-1][j];
            nbrs += board->elems[i-1][j+1];
            nbrs += board->elems[i][j-1];
            nbrs += board->elems[i][j+1];
            nbrs += board->elems[i+1][j-1];
            nbrs += board->elems[i+1][j];
            nbrs += board->elems[i+1][j+1];

            if (board->elems[i][j] == 1) {
                if ((nbrs == 2) || (nbrs == 3))
                    new_board->elems[i][j] = 1;
                else
                    new_board->elems[i][j] = 0;
            }
            else {
                if (nbrs == 3)
                    new_board->elems[i][j] = 1;
                else
                    new_board->elems[i][j] = 0;
            }
        }
    }
}

/*
 * prints current board configuration (code for process 0 -- actual print)
 */
void print_board_p0(FILE* outfile, twoD_array_t *board) {

    long size = board->cols-2;

    /* process 0 prints, accepting values from other processes */
    array_type * temprow = malloc(size * sizeof(array_type));
    if (temprow == NULL) {
        fprintf(stderr, "unable to allocate space for printing board\n");
        /* FIXME do this more gracefully? */
        MPI_Abort(MPI_COMM_WORLD, EXIT_FAILURE);
    }

    MPI_Status status;
    for (long i = 1; i <= size; ++i) {
        int process_for_row = process_owning_row(i, size);
        if (process_for_row != 0) {
            MPI_Recv(temprow, size, MPI_ARRAY_TYPE, process_for_row,
                    PRINT_TAG, MPI_COMM_WORLD, &status);
        }
        for (long j = 1; j <= size; ++j) {
            array_type temp;
            if (process_for_row == 0)
                temp = board->elems[i][j];
            else
                temp = temprow[j-1];
            if (temp == 0)
                fprintf(outfile, ". ");
            else
                fprintf(outfile, "1 ");
        }
        fprintf(outfile, "\n");
    }
    free(temprow);
}
/*
 * prints current board configuration (code for processes other than 0)
 */
void print_board_notp0(twoD_array_t *board) {

    long size = board->cols-2;

    /* processes send values to process 0 */
    for (long i = 1; i <= local_rows(myid, size); ++i) {
        MPI_Send(&(board->elems[i][1]), 
                size, MPI_ARRAY_TYPE, 0, PRINT_TAG, MPI_COMM_WORLD);
    }
}

/* 
 * counts "live" cells 
 */
long live_cells(twoD_array_t *board) {
    long local_count = 0;
    for (long i = 1; i <= board->rows-2; ++i) {
        for (long j = 1; j <= board->cols-2; ++j) {
            local_count += board->elems[i][j];
        }
    }
    long count;
    MPI_Allreduce(&local_count, &count, 1, MPI_LONG, MPI_SUM, MPI_COMM_WORLD);
    return count;
}


/* returns rows per process */
long rows_per_process(long size) {
    return (size + (nprocs-1)) / nprocs;
}
/* returns rows for this process */
long local_rows(int proc_id, long size) {
    if (proc_id == (nprocs-1))  {
        return size - ((nprocs-1) * rows_per_process(size));
    }
    else {
        return rows_per_process(size);
    }
}
/* returns local start row (absolute row number) */
long local_start_row(int proc_id, long size) {
    return proc_id * rows_per_process(size) + 1;
}
/* returns local end row (absolute row number) */
long local_end_row(int proc_id, long size) {
    if (proc_id == (nprocs-1))
        return size + 1;
    else
        return local_start_row(proc_id, size) + rows_per_process(size);
}
/* returns ID of process owning row */
int process_owning_row(int row, long size) {
    return (row-1) / rows_per_process(size);
}