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

#include <stdio.h>
#include "stdlib.h"
#include <mpi.h>
#include <unistd.h>
#include "elementary.h"
#include "../futhark_mpi/dispatch.h"
#include "gfx.h"

#define INDEX_2D_TO_1D(y, x, nb_columns) ((y) * (nb_columns) + (x))
#define INDEX_3D_TO_1D(y, x, z, nb_columns, nb_depths) ((x) + (nb_columns) * ((y) + (nb_depths) * (z)))

#define ROOT_RANK 0
#define FPS (60*800)

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

void compute_next_chunk_elems(struct dispatch_context *dc, struct futhark_context *fc, chunk_info_t *ci) {
    struct futhark_i8_1d *fut_chunk_with_envelope = get_chunk_with_envelope(dc, fc, 1, futhark_new_i8_1d);

    struct futhark_i8_1d *fut_next_chunk_elems;
    futhark_entry_next_chunk_elems(fc, &fut_next_chunk_elems, fut_chunk_with_envelope);
    futhark_context_sync(fc);

    futhark_values_i8_1d(fc, fut_next_chunk_elems, ci->data);
    futhark_context_sync(fc);

    futhark_free_i8_1d(fc, fut_chunk_with_envelope);
    futhark_free_i8_1d(fc, fut_next_chunk_elems);
}

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);
    srand(time(NULL) + my_rank);

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

    int elems_n = atoi(argv[2]);
    int elems_dimensions[1] = {elems_n};
    int history_lines = atoi(argv[1]);

    struct dispatch_context *disp_context = dispatch_context_new(elems_dimensions, MPI_INT8_T, 1);
    chunk_info_t ci = get_chunk_info(disp_context);
    init_chunk_elems(&ci);
    int8_t *history = my_rank == ROOT_RANK ? calloc((size_t) history_lines * (size_t) elems_n, sizeof(int8_t))
                                           : NULL;

    /* GFX Initialization */
    struct gfx_context_t *gfx_context =
            my_rank == ROOT_RANK ? gfx_create("Simple Cellular Automaton", (uint) elems_dimensions[0], (uint) history_lines)
                                 : 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;
    int i = 0;
    while (!exit) {
        /* AC Iteration */
        compute_next_chunk_elems(disp_context, fut_context, &ci);

        /* Recover AC line */
        int8_t *line = get_data(disp_context);
        if (my_rank == ROOT_RANK) {
            memcpy(&history[INDEX_2D_TO_1D(i, 0, elems_n)], line, (size_t) elems_n * sizeof(int8_t));

            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 < history_lines; ++y) {
                for (int x = 0; x < elems_n; ++x) {
                    int cell = (int) history[INDEX_2D_TO_1D(y, x, elems_n)];
                    gfx_putpixel(gfx_context, x, y, MAKE_COLOR(cell * 255, cell * 255, cell * 255));
                }
            }
            gfx_present(gfx_context);
            free(line);
        }

        MPI_Bcast(&exit, 1, MPI_C_BOOL, ROOT_RANK, MPI_COMM_WORLD);
        usleep(1000000 / FPS);
        i = (i + 1) % history_lines;
        if (my_rank == ROOT_RANK && i == 0) {
            sleep(10000);
            memset(history, 0, (size_t) history_lines * (size_t) elems_n * sizeof(int8_t));
        }
    }

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