/**
 * Author: Baptiste Coudray
 * School: HEPIA
 * Class: ITI-3
 * Year 2020-2021
 */

#include <stdio.h>
#include <stdint.h>
#include <mpi.h>
#include <stdbool.h>
#include <unistd.h>
#include "gol.h"
#include "gfx.h"
#include "../futhark_mpi/dispatch.h"

#define INDEX_2D_TO_1D(y, x, nb_columns) ((y) * (nb_columns) + (x))

#define ROOT_RANK 0
#define FPS 60

void init_chunk_board(chunk_info_t *ci) {
    for (int i = 0; i < ci->dimensions[0]; ++i) {
        for (int j = 0; j < ci->dimensions[1]; ++j) {
            ((int8_t *) ci->data)[INDEX_2D_TO_1D(i, j, ci->dimensions[1])] = rand() % 2;
        }
    }
}

void compute_next_chunk_board(struct dispatch_context *dc, struct futhark_context *fc, chunk_info_t *ci) {
    envelope_t outer_envelope = get_outer_envelope(dc, fc, 1);
    struct futhark_i8_2d *fut_north = futhark_new_i8_2d(fc, outer_envelope.north.data,
                                                        outer_envelope.north.dimensions[0],
                                                        outer_envelope.north.dimensions[1]);
    struct futhark_i8_2d *fut_north_east = futhark_new_i8_2d(fc, outer_envelope.north_east.data,
                                                             outer_envelope.north_east.dimensions[0],
                                                             outer_envelope.north_east.dimensions[1]);

    struct futhark_i8_2d *fut_east = futhark_new_i8_2d(fc, outer_envelope.east.data,
                                                       outer_envelope.east.dimensions[0],
                                                       outer_envelope.east.dimensions[1]);

    struct futhark_i8_2d *fut_south_east = futhark_new_i8_2d(fc, outer_envelope.south_east.data,
                                                             outer_envelope.south_east.dimensions[0],
                                                             outer_envelope.south_east.dimensions[1]);
    struct futhark_i8_2d *fut_south = futhark_new_i8_2d(fc, outer_envelope.south.data,
                                                        outer_envelope.south.dimensions[0],
                                                        outer_envelope.south.dimensions[1]);

    struct futhark_i8_2d *fut_south_west = futhark_new_i8_2d(fc, outer_envelope.south_west.data,
                                                             outer_envelope.south_west.dimensions[0],
                                                             outer_envelope.south_west.dimensions[1]);

    struct futhark_i8_2d *fut_west = futhark_new_i8_2d(fc, outer_envelope.west.data,
                                                       outer_envelope.west.dimensions[0],
                                                       outer_envelope.west.dimensions[1]);

    struct futhark_i8_2d *fut_north_west = futhark_new_i8_2d(fc, outer_envelope.north_west.data,
                                                             outer_envelope.north_west.dimensions[0],
                                                             outer_envelope.north_west.dimensions[1]);

    struct futhark_i8_2d *fut_chunk_board = futhark_new_i8_2d(fc, ci->data, ci->dimensions[0], ci->dimensions[1]);
    futhark_context_sync(fc);
    envelope_free(&outer_envelope);

    struct futhark_i8_2d *fut_next_chunk_board;
    futhark_entry_next_chunk_board(fc, &fut_next_chunk_board, fut_chunk_board, fut_north_west, fut_north,
                                   fut_north_east, fut_west, fut_east, fut_south_west, fut_south, fut_south_east);

    futhark_context_sync(fc);

    futhark_free_i8_2d(fc, fut_chunk_board);
    futhark_free_i8_2d(fc, fut_north_west);
    futhark_free_i8_2d(fc, fut_north);
    futhark_free_i8_2d(fc, fut_north_east);
    futhark_free_i8_2d(fc, fut_west);
    futhark_free_i8_2d(fc, fut_east);
    futhark_free_i8_2d(fc, fut_south_west);
    futhark_free_i8_2d(fc, fut_south);
    futhark_free_i8_2d(fc, fut_south_east);

    futhark_values_i8_2d(fc, fut_next_chunk_board, ci->data);
    futhark_context_sync(fc);
    futhark_free_i8_2d(fc, fut_next_chunk_board);
}

int main(int argc, char *argv[]) {
    if(argc < 3) {
        printf("usage: mpirun -n <nb_proc> %s <height> <width>\n", argv[0]);
        return EXIT_FAILURE;
    }

    MPI_Init(&argc, &argv);
    int my_rank;
    int world_size;
    MPI_Comm_rank(MPI_COMM_WORLD, &my_rank);
    MPI_Comm_size(MPI_COMM_WORLD, &world_size);

    struct futhark_context_config *fut_config = futhark_context_config_new();
    struct futhark_context *fut_context = futhark_context_new(fut_config);

    int board_n = atoi(argv[1]);
    int board_m = atoi(argv[2]);
    int board_dimensions[2] = {board_n, board_m};

    struct dispatch_context *disp_context = dispatch_context_new(board_dimensions, MPI_INT8_T, 2);
    chunk_info_t ci = get_chunk_info(disp_context);
    init_chunk_board(&ci);

    /* GFX Initialization */
    struct gfx_context_t *gfx_context =
            my_rank == ROOT_RANK ? gfx_create("Game of Life", (uint) board_dimensions[1], (uint) board_dimensions[0])
                                 : NULL;
    if (my_rank == ROOT_RANK && !gfx_context) {
        fprintf(stderr, "Graphic mode initialization failed!\n");
        return EXIT_FAILURE;
    }
    if (my_rank == ROOT_RANK) {
        SDL_ShowCursor(SDL_ENABLE);
    }

    bool exit = false;
    while (!exit) {
        /* GoL Iteration */
        compute_next_chunk_board(disp_context, fut_context, &ci);

        /* Recover GoL board */
        int8_t *board = get_data(disp_context);

        if (my_rank == ROOT_RANK) {
            SDL_PumpEvents();
            SDL_Event event;
            SDL_PollEvent(&event);

            exit = gfx_keypressed() == SDLK_ESCAPE ||
                   (event.type == SDL_WINDOWEVENT && event.window.event == SDL_WINDOWEVENT_CLOSE);

            gfx_clear(gfx_context, COLOR_BLACK);
            for (int y = 0; y < board_dimensions[0]; ++y) {
                for (int x = 0; x < board_dimensions[1]; ++x) {
                    int cell = (int) board[INDEX_2D_TO_1D(y, x, board_dimensions[1])];
                    gfx_putpixel(gfx_context, x, y, MAKE_COLOR(cell * 255, cell * 255, cell * 255));
                }
            }
            gfx_present(gfx_context);
            free(board);
        }

        MPI_Bcast(&exit, 1, MPI_C_BOOL, ROOT_RANK, MPI_COMM_WORLD);
        usleep(1000000 / FPS);
    }

    if (my_rank == ROOT_RANK) {
        gfx_destroy(gfx_context);
    }
    dispatch_context_free(disp_context);
    futhark_context_config_free(fut_config);
    futhark_context_free(fut_context);
    return MPI_Finalize();
}