/*
 * Copyright (C) NHR@FAU, University Erlangen-Nuremberg.
 * All rights reserved. This file is part of nusif-solver.
 * Use of this source code is governed by a MIT style
 * license that can be found in the LICENSE file.
 */
#include <stdlib.h>

#include "comm.h"

#ifdef _MPI

static void gatherArray(
    Comm* c, int cnt, int* rcvCounts, int* displs, double* src, double* dst)
{
    double* sendbuffer = src + (c->imaxLocal + 2);

    if (c->rank == 0) {
        sendbuffer = src;
    }

    MPI_Gatherv(sendbuffer,
        cnt,
        MPI_DOUBLE,
        dst,
        rcvCounts,
        displs,
        MPI_DOUBLE,
        0,
        MPI_COMM_WORLD);
}
#endif // defined _MPI

// exported subroutines
int commIsBoundary(Comm* c, int direction)
{
#ifdef _MPI
    switch (direction) {
    case L:
        return 1;
        break;
    case R:
        return 1;
        break;
    case B:
        return c->rank == 0;
        break;
    case T:
        return c->rank == (c->size - 1);
        break;
    }
#endif

    return 1;
}

void commExchange(Comm* c, double* grid)
{
#ifdef _MPI
    MPI_Request requests[4] = { MPI_REQUEST_NULL,
        MPI_REQUEST_NULL,
        MPI_REQUEST_NULL,
        MPI_REQUEST_NULL };

    /* exchange ghost cells with top neighbor */
    if (c->rank + 1 < c->size) {
        int top     = c->rank + 1;
        double* src = grid + (c->jmaxLocal) * (c->imaxLocal + 2) + 1;
        double* dst = grid + (c->jmaxLocal + 1) * (c->imaxLocal + 2) + 1;

        MPI_Isend(src, c->imaxLocal, MPI_DOUBLE, top, 1, MPI_COMM_WORLD, &requests[0]);
        MPI_Irecv(dst, c->imaxLocal, MPI_DOUBLE, top, 2, MPI_COMM_WORLD, &requests[1]);
    }

    /* exchange ghost cells with bottom neighbor */
    if (c->rank > 0) {
        int bottom  = c->rank - 1;
        double* src = grid + (c->imaxLocal + 2) + 1;
        double* dst = grid + 1;

        MPI_Isend(src, c->imaxLocal, MPI_DOUBLE, bottom, 2, MPI_COMM_WORLD, &requests[2]);
        MPI_Irecv(dst, c->imaxLocal, MPI_DOUBLE, bottom, 1, MPI_COMM_WORLD, &requests[3]);
    }

    MPI_Waitall(4, requests, MPI_STATUSES_IGNORE);
#endif
}

void commShift(Comm* c, double* f, double* g)
{
#ifdef _MPI
    MPI_Request requests[2] = { MPI_REQUEST_NULL, MPI_REQUEST_NULL };

    /* shift G */
    /* receive ghost cells from bottom neighbor */
    if (c->rank > 0) {
        int bottom = c->rank - 1;
        MPI_Irecv(g + 1,
            c->imaxLocal,
            MPI_DOUBLE,
            bottom,
            0,
            MPI_COMM_WORLD,
            &requests[0]);
    }

    if (c->rank + 1 < c->size) {
        int top     = c->rank + 1;
        double* buf = g + (c->jmaxLocal) * (c->imaxLocal + 2) + 1;
        /* send ghost cells to top neighbor */
        MPI_Isend(buf, c->imaxLocal, MPI_DOUBLE, top, 0, MPI_COMM_WORLD, &requests[1]);
    }

    MPI_Waitall(2, requests, MPI_STATUSES_IGNORE);
#endif
}

void commCollectResult(Comm* c,
    double* ug,
    double* vg,
    double* pg,
    double* u,
    double* v,
    double* p,
    int imax,
    int jmax)
{
#ifdef _MPI
    int *rcvCounts, *displs;
    int cnt = c->jmaxLocal * (imax + 2);

    if (c->rank == 0) {
        rcvCounts = (int*)malloc(c->size * sizeof(int));
        displs    = (int*)malloc(c->size * sizeof(int));
    }

    if (c->rank == 0 && c->size == 1) {
        cnt = (c->jmaxLocal + 2) * (imax + 2);
    } else if (c->rank == 0 || c->rank == (c->size - 1)) {
        cnt = (c->jmaxLocal + 1) * (imax + 2);
    }

    MPI_Gather(&cnt, 1, MPI_INTEGER, rcvCounts, 1, MPI_INTEGER, 0, MPI_COMM_WORLD);

    if (c->rank == 0) {
        displs[0]  = 0;
        int cursor = rcvCounts[0];

        for (int i = 1; i < c->size; i++) {
            displs[i] = cursor;
            cursor += rcvCounts[i];
        }
    }

    gatherArray(c, cnt, rcvCounts, displs, p, pg);
    gatherArray(c, cnt, rcvCounts, displs, u, ug);
    gatherArray(c, cnt, rcvCounts, displs, v, vg);
#endif
}

void commPartition(Comm* c, int jmax, int imax)
{
#ifdef _MPI
    c->imaxLocal = imax;
    c->jmaxLocal = sizeOfRank(c->coords[JDIM], c->size, jmax);

    c->neighbours[BOTTOM] = c->rank == 0 ? -1 : c->rank - 1;
    c->neighbours[TOP]    = c->rank == (c->size - 1) ? -1 : c->rank + 1;
    c->neighbours[LEFT]   = -1;
    c->neighbours[RIGHT]  = -1;

    c->coords[IDIM] = 0;
    c->coords[JDIM] = c->rank;

    c->dims[IDIM] = 1;
    c->dims[JDIM] = c->size;
#else
    c->imaxLocal = imax;
    c->jmaxLocal = jmax;
#endif
}

void commUpdateDatatypes(Comm* oldcomm, Comm* newcomm, int imaxLocal, int jmaxLocal)
{

#if defined _MPI
    newcomm->comm = MPI_COMM_NULL;
    int result    = MPI_Comm_dup(MPI_COMM_WORLD, &newcomm->comm);

    if (result == MPI_ERR_COMM) {
        printf("\nNull communicator. Duplication failed !!\n");
    }

    newcomm->rank = oldcomm->rank;
    newcomm->size = oldcomm->size;

    newcomm->imaxLocal = imaxLocal / 2;
    newcomm->jmaxLocal = jmaxLocal / 2;

    newcomm->neighbours[BOTTOM] = newcomm->rank == 0 ? -1 : newcomm->rank - 1;
    newcomm->neighbours[TOP]    = newcomm->rank == (newcomm->size - 1) ? -1 : newcomm->rank + 1;
    newcomm->neighbours[LEFT]   = -1;
    newcomm->neighbours[RIGHT]  = -1;

    newcomm->coords[IDIM] = 0;
    newcomm->coords[JDIM] = newcomm->rank;

    newcomm->dims[IDIM] = 1;
    newcomm->dims[JDIM] = newcomm->size;


#endif
    newcomm->imaxLocal = imaxLocal;
    newcomm->jmaxLocal = jmaxLocal;
}

void commFreeCommunicator(Comm* comm)
{
#ifdef _MPI
    MPI_Comm_free(&comm->comm);
#endif
}