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

#include <stdio.h>
#include <stdint.h>
#include <mpi.h>
#include <math.h>
#include <stdbool.h>
#include "../gol.h"

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

#define NORTH_INDEX 0
#define EAST_INDEX 1
#define SOUTH_INDEX 2
#define WEST_INDEX 3

#define NORTH_ROW_TAG 0
#define EAST_COLUMN_TAG 1
#define SOUTH_ROW_TAG 2
#define WEST_COLUMN_TAG 3

#define NORTH_EAST_CELL_TAG 4
#define SOUTH_EAST_CELL_TAG 5
#define SOUTH_WEST_CELL_TAG 6
#define NORTH_WEST_CELL_TAG 7

#define CHUNK_BOARD_TAG 8

#define ROOT_RANK 0
#define FPS 60

int createGridCommunicators(MPI_Comm *cartComm, MPI_Comm *rowComm, MPI_Comm *colComm, int nProc) {
    int gridN = (int) sqrt(nProc);
    int dimensions[2] = {gridN, gridN};
    int periods[2] = {true, true}; // Cyclic on row-column

    MPI_Cart_create(MPI_COMM_WORLD, 2, dimensions, periods, 1, cartComm);

    /* Create row communicator */
    int remainDims[2] = {false, true};
    MPI_Cart_sub(*cartComm, remainDims, rowComm);

    /* Create column communicator */
    remainDims[0] = true; // rows
    remainDims[1] = false; // columns
    MPI_Cart_sub(*cartComm, remainDims, colComm);
    return gridN;
}

int *divideBoard(int n, int chunkN, int nProc) {
    int *indexes = calloc((size_t) nProc * 2, sizeof(int));
    for (int i = 0, y = 0, x = 0; i < nProc; ++i) {
        indexes[i * 2] = y;
        indexes[i * 2 + 1] = x;

        x += chunkN;
        if (x >= n) {
            x = 0;
            y += chunkN;
        }
    }
    return indexes;
}

void initChunkBoard(int8_t *chunkBoard, int chunkN) {
    for (int i = 0; i < chunkN; ++i) {
        for (int j = 0; j < chunkN; ++j) {
            chunkBoard[INDEX_2D_TO_1D(i, j, chunkN)] = rand() % 2;
        }
    }
}

void shareAndBuildEnvelope(int8_t *chunkBoardMyEnvelope, int8_t *chunkBoardEnvelope, MPI_Comm rowComm,
                           MPI_Comm colComm, int gridN, int coordinates[2], int chunkN, int chunkM) {
    int coordinateY = coordinates[0];
    int coordinateX = coordinates[1];
    MPI_Request requests[16] = {0};
    int iRequest = 0;

    // North
    {
        int8_t *chunkBoardMyEnvelopeNorth = &chunkBoardMyEnvelope[INDEX_2D_TO_1D(NORTH_INDEX, 0, chunkN)];
        int8_t *chunkBoardEnvelopeNorth = &chunkBoardEnvelope[INDEX_2D_TO_1D(NORTH_INDEX, 1, chunkM)];
        int destSource = (coordinateY - 1) < 0 ? (gridN - 1) : (coordinateY - 1);

        MPI_Isend(chunkBoardMyEnvelopeNorth, chunkN, MPI_INT8_T, destSource, NORTH_ROW_TAG, colComm,
                  &requests[iRequest++]);
        /* Neighbour send south row, which correspond to north envelope */
        MPI_Irecv(chunkBoardEnvelopeNorth, chunkN, MPI_INT8_T, destSource, SOUTH_ROW_TAG, colComm,
                  &requests[iRequest++]);
    }

    // East
    {
        int8_t *chunkBoardMyEnvelopeEast = &chunkBoardMyEnvelope[INDEX_2D_TO_1D(EAST_INDEX, 0, chunkN)];
        int8_t *chunkBoardEnvelopeEast = &chunkBoardEnvelope[INDEX_2D_TO_1D(EAST_INDEX, 1, chunkM)];
        int destSource = (coordinateX + 1) % gridN;

        MPI_Isend(chunkBoardMyEnvelopeEast, chunkN, MPI_INT8_T, destSource, EAST_COLUMN_TAG, rowComm,
                  &requests[iRequest++]);
        /* Neighbour send west column, which correspond to east envelope */
        MPI_Irecv(chunkBoardEnvelopeEast, chunkN, MPI_INT8_T, destSource, WEST_COLUMN_TAG, rowComm,
                  &requests[iRequest++]);
    }

    // South
    {
        int8_t *chunkBoardMyEnvelopeSouth = &chunkBoardMyEnvelope[INDEX_2D_TO_1D(SOUTH_INDEX, 0, chunkN)];
        int8_t *chunkBoardEnvelopeSouth = &chunkBoardEnvelope[INDEX_2D_TO_1D(SOUTH_INDEX, 1, chunkM)];
        int destSource = (coordinateY + 1) % gridN;

        MPI_Isend(chunkBoardMyEnvelopeSouth, chunkN, MPI_INT8_T, destSource, SOUTH_ROW_TAG, colComm,
                  &requests[iRequest++]);
        /* Neighbour send north row, which correspond to south envelope */
        MPI_Irecv(chunkBoardEnvelopeSouth, chunkN, MPI_INT8_T, destSource, NORTH_ROW_TAG, colComm,
                  &requests[iRequest++]);
    }

    // West
    {
        int8_t *chunkBoardMyEnvelopeWest = &chunkBoardMyEnvelope[INDEX_2D_TO_1D(WEST_INDEX, 0, chunkN)];
        int8_t *chunkBoardEnvelopeWest = &chunkBoardEnvelope[INDEX_2D_TO_1D(WEST_INDEX, 1, chunkM)];
        int destSource = (coordinateX - 1) < 0 ? (gridN - 1) : (coordinateX - 1);

        MPI_Isend(chunkBoardMyEnvelopeWest, chunkN, MPI_INT8_T, destSource, WEST_COLUMN_TAG, rowComm,
                  &requests[iRequest++]);
        /* Neighbour send east column, which correspond to west envelope */
        MPI_Irecv(chunkBoardEnvelopeWest, chunkN, MPI_INT8_T, destSource, EAST_COLUMN_TAG, rowComm,
                  &requests[iRequest++]);
    }

    int8_t missingCells[4] = {0};

    // North-East
    {
        int8_t *chunkBoardMyEnvelopeNorthEast = &chunkBoardMyEnvelope[INDEX_2D_TO_1D(NORTH_INDEX, chunkN - 1, chunkN)];
        int destSrcY = (coordinateY - 1) < 0 ? gridN - 1 : coordinateY - 1;
        int destSrcX = (coordinateX + 1) % gridN;
        int destSource = INDEX_2D_TO_1D(destSrcY, destSrcX, gridN);

        MPI_Isend(chunkBoardMyEnvelopeNorthEast, 1, MPI_INT8_T, destSource, NORTH_EAST_CELL_TAG, MPI_COMM_WORLD,
                  &requests[iRequest++]);
        /* Neighbour send south-west cell, which correspond to north-east cell */
        MPI_Irecv(&missingCells[1], 1, MPI_INT8_T, destSource, SOUTH_WEST_CELL_TAG, MPI_COMM_WORLD,
                  &requests[iRequest++]);
    }

    // South-East
    {
        int8_t *chunkBoardMyEnvelopeSouthEast = &chunkBoardMyEnvelope[INDEX_2D_TO_1D(SOUTH_INDEX, chunkN - 1, chunkN)];
        int destSrcY = (coordinateY + 1) % gridN;
        int destSrcX = (coordinateX + 1) % gridN;
        int destSource = INDEX_2D_TO_1D(destSrcY, destSrcX, gridN);

        MPI_Isend(chunkBoardMyEnvelopeSouthEast, 1, MPI_INT8_T, destSource, SOUTH_EAST_CELL_TAG, MPI_COMM_WORLD,
                  &requests[iRequest++]);
        /* Neighbour send north-west cell, which correspond to south-east cell */
        MPI_Irecv(&missingCells[2], 1, MPI_INT8_T, destSource, NORTH_WEST_CELL_TAG, MPI_COMM_WORLD,
                  &requests[iRequest++]);
    }

    // South-West
    {
        int8_t *chunkBoardMyEnvelopeSouthWest = &chunkBoardMyEnvelope[INDEX_2D_TO_1D(SOUTH_INDEX, 0, chunkN)];
        int destSrcY = (coordinateY + 1) % gridN;
        int destSrcX = (coordinateX - 1) < 0 ? gridN - 1 : coordinateX - 1;
        int destSource = INDEX_2D_TO_1D(destSrcY, destSrcX, gridN);

        MPI_Isend(chunkBoardMyEnvelopeSouthWest, 1, MPI_INT8_T, destSource, SOUTH_WEST_CELL_TAG, MPI_COMM_WORLD,
                  &requests[iRequest++]);
        /* Neighbour send north-east cell, which correspond to south-west cell */
        MPI_Irecv(&missingCells[3], 1, MPI_INT8_T, destSource, NORTH_EAST_CELL_TAG, MPI_COMM_WORLD,
                  &requests[iRequest++]);
    }

    // North-West
    {
        int8_t *chunkBoardMyEnvelopeNorthWest = &chunkBoardMyEnvelope[INDEX_2D_TO_1D(NORTH_INDEX, 0, chunkN)];
        int destSrcY = (coordinateY - 1) < 0 ? gridN - 1 : coordinateY - 1;
        int destSrcX = (coordinateX - 1) < 0 ? gridN - 1 : coordinateX - 1;
        int destSource = INDEX_2D_TO_1D(destSrcY, destSrcX, gridN);

        MPI_Isend(chunkBoardMyEnvelopeNorthWest, 1, MPI_INT8_T, destSource, NORTH_WEST_CELL_TAG, MPI_COMM_WORLD,
                  &requests[iRequest++]);
        /* Neighbour send south-east cell, which correspond to north-west cell */
        MPI_Irecv(&missingCells[0], 1, MPI_INT8_T, destSource, SOUTH_EAST_CELL_TAG, MPI_COMM_WORLD,
                  &requests[iRequest]);
    }

    MPI_Waitall(16, requests, MPI_STATUSES_IGNORE);

    /* Copy missing cells into the envelope */
    chunkBoardEnvelope[INDEX_2D_TO_1D(NORTH_INDEX, chunkN, chunkM)] = chunkBoardEnvelope[INDEX_2D_TO_1D(
            EAST_INDEX, 0, chunkM)] = missingCells[1];
    chunkBoardEnvelope[INDEX_2D_TO_1D(SOUTH_INDEX, chunkN, chunkM)] = chunkBoardEnvelope[INDEX_2D_TO_1D(EAST_INDEX,
                                                                                                        chunkN,
                                                                                                        chunkM)] = missingCells[2];
    chunkBoardEnvelope[INDEX_2D_TO_1D(SOUTH_INDEX, 0, chunkM)] = chunkBoardEnvelope[INDEX_2D_TO_1D(WEST_INDEX,
                                                                                                   chunkN,
                                                                                                   chunkM)] = missingCells[3];
    chunkBoardEnvelope[INDEX_2D_TO_1D(NORTH_INDEX, 0, chunkM)] = chunkBoardEnvelope[INDEX_2D_TO_1D(WEST_INDEX,
                                                                                                   0,
                                                                                                   chunkM)] = missingCells[0];
}

void chunkBoardToBoard(int8_t *board, int n, const int8_t *chunkBoard, int chunkN, const int *indexes, int rank) {
    int y = indexes[rank * 2];
    int x = indexes[rank * 2 + 1];

    for (int i = 0; i < chunkN; ++i) {
        for (int j = 0; j < chunkN; ++j) {
            board[INDEX_2D_TO_1D(y + i, x + j, n)] = chunkBoard[INDEX_2D_TO_1D(i, j, chunkN)];
        }
    }
}

int main(int argc, char *argv[]) {
    if (argc < 3) {
        printf("%s <deviceName> <boardN>\n", argv[0]);
        exit(0);
    }
    int myRank;
    int nProc;
    double start = MPI_Wtime();

    /* MPI Initialization */
    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &myRank);
    MPI_Comm_size(MPI_COMM_WORLD, &nProc);
    srand((unsigned int) myRank);

    MPI_Comm cartComm, rowComm, colComm;
    int gridN = createGridCommunicators(&cartComm, &rowComm, &colComm, nProc);

    int myCartRank;
    MPI_Comm_rank(cartComm, &myCartRank);

    int coordinates[2] = {0};
    MPI_Cart_coords(cartComm, myCartRank, 2, coordinates);

    /* Futhark Initialization */
    struct futhark_context_config *contextConfig = futhark_context_config_new();
    int nbDevices = atoi(argv[1]);
    char device[3] = {0};
    snprintf(device, sizeof(device), "#%d", myRank % nbDevices);
    futhark_context_config_set_device(contextConfig, device);
    struct futhark_context *futharkContext = futhark_context_new(contextConfig);

    /* GoL Initialization */
    int boardN = atoi(argv[2]);
    int boardNN = boardN * boardN;
    int chunkN = (int) (boardN / sqrt(nProc));
    int chunkNN = chunkN * chunkN;
    int chunkM = chunkN + 2;
    int *indexes = divideBoard(boardN, chunkN, nProc);

    int8_t *board = myRank == ROOT_RANK ? calloc((size_t) boardNN, sizeof(int8_t)) : NULL;
    int8_t *chunkBoard = calloc((size_t) chunkNN, sizeof(int8_t));
    int8_t *chunkBoardMyEnvelope = calloc(((size_t) (4 * chunkN)), sizeof(int8_t));
    int8_t *chunkBoardEnvelope = calloc(((size_t) (4 * chunkM)), sizeof(int8_t));

    initChunkBoard(chunkBoard, chunkN);

    struct futhark_i8_2d *futChunkBoard = futhark_new_i8_2d(futharkContext, chunkBoard, chunkN, chunkN);
    futhark_context_sync(futharkContext);
    struct futhark_i8_2d *futChunkBoardMyEnvelope;
    futhark_entry_get_envelope(futharkContext, &futChunkBoardMyEnvelope, futChunkBoard);
    futhark_context_sync(futharkContext);
    futhark_values_i8_2d(futharkContext, futChunkBoardMyEnvelope, chunkBoardMyEnvelope);
    futhark_context_sync(futharkContext);

    shareAndBuildEnvelope(chunkBoardMyEnvelope, chunkBoardEnvelope, rowComm, colComm, gridN, coordinates, chunkN,
                          chunkM);

    struct futhark_i8_2d *futChunkBoardEnvelope = futhark_new_i8_2d(futharkContext, chunkBoardEnvelope, 4, chunkM);
    futhark_context_sync(futharkContext);
    struct futhark_i8_2d *futNextChunkBoard;
    futhark_entry_next_chunk_board(futharkContext, &futNextChunkBoard, futChunkBoard, futChunkBoardEnvelope);
    futhark_context_sync(futharkContext);
    futhark_values_i8_2d(futharkContext, futNextChunkBoard, chunkBoard);
    futhark_context_sync(futharkContext);

    if (myRank == ROOT_RANK) {
        chunkBoardToBoard(board, boardN, chunkBoard, chunkN, indexes, myRank);
        int8_t *tmpChunkBoard = calloc((size_t) chunkNN, sizeof(int8_t));
        MPI_Status status = {0};
        for (int i = 0; i < nProc - 1; ++i) {
            MPI_Recv(tmpChunkBoard, chunkNN, MPI_INT8_T, MPI_ANY_SOURCE, CHUNK_BOARD_TAG, MPI_COMM_WORLD, &status);
            chunkBoardToBoard(board, boardN, tmpChunkBoard, chunkN, indexes, status.MPI_SOURCE);
        }
        free(tmpChunkBoard);
    } else {
        MPI_Send(chunkBoard, chunkNN, MPI_INT8_T, ROOT_RANK, CHUNK_BOARD_TAG, MPI_COMM_WORLD);
    }


    futhark_context_sync(futharkContext);
    futhark_free_i8_2d(futharkContext, futChunkBoard);
    futhark_free_i8_2d(futharkContext, futChunkBoardMyEnvelope);
    futhark_free_i8_2d(futharkContext, futChunkBoardEnvelope);
    futhark_free_i8_2d(futharkContext, futNextChunkBoard);


    free(indexes);
    free(chunkBoard);
    free(chunkBoardEnvelope);
    free(chunkBoardMyEnvelope);

    if (myRank == ROOT_RANK) {
        free(board);
    }

    futhark_context_free(futharkContext);
    futhark_context_config_free(contextConfig);

    MPI_Comm_free(&cartComm);
    MPI_Comm_free(&rowComm);
    MPI_Comm_free(&colComm);

    double finish = MPI_Wtime();
    if (myRank == ROOT_RANK) {
        printf("%d;%f\n", nProc, finish - start);
    }
    MPI_Finalize();
    return 0;
}