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

#include <stdio.h>
#include <stdint.h>
#include "stdlib.h"
#include <mpi.h>
#include "../gol.h"
#include "../../futhark_mpi/dispatch.h"

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

#define ROOT_RANK 0

#define N_MEASURES 15

#if defined(FUTHARK_BACKEND_cuda) || defined(FUTHARK_BACKEND_opencl)
#define N_ITERATIONS 300
#else
#define N_ITERATIONS 100
#endif

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) {
            size_t k = INDEX_2D_TO_1D((size_t) i, (size_t) j, (size_t) ci->dimensions[1]);
            ((int8_t *) ci->data)[k] = rand() % 2;
        }
    }
}

double compute_next_chunk_board(struct dispatch_context *dc, struct futhark_context *fc, chunk_info_t *ci) {
    double start = MPI_Wtime();
    struct futhark_i8_2d *fut_chunk_with_envelope = get_chunk_with_envelope(dc, fc, 1, futhark_new_i8_2d);

    struct futhark_i8_2d *fut_next_chunk_board;
    for (int i = 0; i < N_ITERATIONS; ++i) {
        futhark_entry_next_chunk_board(fc, &fut_next_chunk_board, fut_chunk_with_envelope);
        futhark_context_sync(fc);
        if (i + 1 < N_ITERATIONS) {
            futhark_free_i8_2d(fc, fut_next_chunk_board);
        }
    }
    futhark_values_i8_2d(fc, fut_next_chunk_board, ci->data);
    futhark_context_sync(fc);
    futhark_free_i8_2d(fc, fut_next_chunk_board);
    futhark_free_i8_2d(fc, fut_chunk_with_envelope);
    double finish = MPI_Wtime();

    return finish - start;
}

int main(int argc, char *argv[]) {
    if (argc < 4) {
        printf("usage: mpirun -n <nb_proc> %s <nb_devices> <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();
    int nb_devices = atoi(argv[1]);
#if defined(FUTHARK_BACKEND_cuda) || defined(FUTHARK_BACKEND_opencl)
#if defined(FUTHARK_BACKEND_opencl)
    futhark_context_config_list_devices(fut_config);
#endif
    char device[4] = {0};
    snprintf(device, sizeof(device), "#%d", my_rank % nb_devices);
    futhark_context_config_set_device(fut_config, device);
#endif
    struct futhark_context *fut_context = futhark_context_new(fut_config);

    int board_n = atoi(argv[2]);
    int board_m = atoi(argv[3]);
    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);

    for (int i = 0; i < N_MEASURES; ++i) {
        double time = compute_next_chunk_board(disp_context, fut_context, &ci);
        if (my_rank == ROOT_RANK) {
            printf("%d;%d;%d;%d;%f\n", world_size, nb_devices, board_n, board_m, time);
        }
    }

    dispatch_context_free(disp_context);
    futhark_context_config_free(fut_config);
    futhark_context_free(fut_context);
    return MPI_Finalize();
}