/* -*- Mode: C; c-basic-offset:4 ; indent-tabs-mode:nil ; -*- */
/*
 *
 *  (C) 2015 by Argonne National Laboratory.
 *      See COPYRIGHT in top-level directory.
 */

/* This test is going to test the atomicity for "read-modify-write" in GACC
 * operations */

/* This test is similiar with atomic_rmw_fop.c.
 * There are three processes involved in this test: P0 (origin_shm), P1 (origin_am),
 * and P2 (dest). P0 and P1 issues multiple GACC with MPI_SUM and OP_COUNT integers
 * (value 1) to P2 via SHM and AM respectively. The correct results should be that the
 * results on P0 and P1 never be the same for intergers on the corresponding index
 * in [0...OP_COUNT-1].
 */

#include "mpi.h"
#include <stdio.h>
#include "mpitest.h"

#define OP_COUNT 10
#define AM_BUF_NUM  10
#define SHM_BUF_NUM 10000
#define WIN_BUF_NUM 1

#define LOOP_SIZE 15
#define CHECK_TAG 123

int rank, size;
int dest, origin_shm, origin_am;
int *orig_buf = NULL, *result_buf = NULL, *target_buf = NULL, *check_buf = NULL;
MPI_Win win;

void checkResults(int loop_k, int *errors)
{
    int i, j, m;
    MPI_Status status;

    if (rank != dest) {
        /* check results on P0 and P2 (origin) */
        if (rank == origin_am) {
            MPI_Send(result_buf, AM_BUF_NUM * OP_COUNT, MPI_INT, origin_shm, CHECK_TAG,
                     MPI_COMM_WORLD);
        } else if (rank == origin_shm) {
            MPI_Alloc_mem(sizeof(int) * AM_BUF_NUM * OP_COUNT, MPI_INFO_NULL, &check_buf);
            MPI_Recv(check_buf, AM_BUF_NUM * OP_COUNT, MPI_INT, origin_am, CHECK_TAG,
                     MPI_COMM_WORLD, &status);
            for (i = 0; i < AM_BUF_NUM; i++) {
                for (j = 0; j < SHM_BUF_NUM; j++) {
                    for (m = 0; m < OP_COUNT; m++) {
                        if (check_buf[i * OP_COUNT + m] == result_buf[j * OP_COUNT + m]) {
                            printf
                                ("LOOP=%d, rank=%d, FOP, both check_buf[%d] and result_buf[%d] equal to %d, expected to be different. \n",
                                 loop_k, rank, i * OP_COUNT + m, j * OP_COUNT + m,
                                 check_buf[i * OP_COUNT + m]);
                            (*errors)++;
                        }
                    }
                }
            }
            MPI_Free_mem(check_buf);
        }
    } else {
        MPI_Win_lock(MPI_LOCK_SHARED, rank, 0, win);
        /* check results on P1 */
        for (i = 0; i < OP_COUNT; i++) {
            if (target_buf[i] != AM_BUF_NUM + SHM_BUF_NUM) {
                printf("LOOP=%d, rank=%d, FOP, target_buf[%d] = %d, expected %d. \n",
                       loop_k, rank, i, target_buf[i], AM_BUF_NUM + SHM_BUF_NUM);
                (*errors)++;
            }
        }
        MPI_Win_unlock(rank, win);
    }
}

int main(int argc, char *argv[])
{
    int i, k;
    int errors = 0;
    int my_buf_num;
    MPI_Datatype origin_dtp, target_dtp;

    MTest_Init(&argc, &argv);

    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);
    if (size != 3) {
        /* run this test with three processes */
        goto exit_test;
    }

    MPI_Type_contiguous(OP_COUNT, MPI_INT, &origin_dtp);
    MPI_Type_commit(&origin_dtp);
    MPI_Type_contiguous(OP_COUNT, MPI_INT, &target_dtp);
    MPI_Type_commit(&target_dtp);

    /* this works when MPIR_PARAM_CH3_ODD_EVEN_CLIQUES is set */
    dest = 2;
    origin_shm = 0;
    origin_am = 1;

    if (rank == origin_am)
        my_buf_num = AM_BUF_NUM;
    else if (rank == origin_shm)
        my_buf_num = SHM_BUF_NUM;

    if (rank != dest) {
        MPI_Alloc_mem(sizeof(int) * my_buf_num * OP_COUNT, MPI_INFO_NULL, &orig_buf);
        MPI_Alloc_mem(sizeof(int) * my_buf_num * OP_COUNT, MPI_INFO_NULL, &result_buf);
    }

    MPI_Win_allocate(sizeof(int) * WIN_BUF_NUM * OP_COUNT, sizeof(int), MPI_INFO_NULL,
                     MPI_COMM_WORLD, &target_buf, &win);

    for (k = 0; k < LOOP_SIZE; k++) {

        /* ====== Part 1: test basic datatypes ======== */

        /* init buffers */
        if (rank != dest) {
            for (i = 0; i < my_buf_num * OP_COUNT; i++) {
                orig_buf[i] = 1;
                result_buf[i] = 0;
            }
        } else {
            MPI_Win_lock(MPI_LOCK_SHARED, rank, 0, win);
            for (i = 0; i < WIN_BUF_NUM * OP_COUNT; i++) {
                target_buf[i] = 0;
            }
            MPI_Win_unlock(rank, win);
        }

        MPI_Barrier(MPI_COMM_WORLD);

        MPI_Win_lock_all(0, win);
        if (rank != dest) {
            for (i = 0; i < my_buf_num; i++) {
                MPI_Get_accumulate(&(orig_buf[i * OP_COUNT]), OP_COUNT, MPI_INT,
                                   &(result_buf[i * OP_COUNT]), OP_COUNT, MPI_INT,
                                   dest, 0, OP_COUNT, MPI_INT, MPI_SUM, win);
                MPI_Win_flush(dest, win);
            }
        }
        MPI_Win_unlock_all(win);

        MPI_Barrier(MPI_COMM_WORLD);

        checkResults(k, &errors);

        /* ====== Part 2: test derived datatypes (origin derived, target derived) ======== */

        /* init buffers */
        if (rank != dest) {
            for (i = 0; i < my_buf_num * OP_COUNT; i++) {
                orig_buf[i] = 1;
                result_buf[i] = 0;
            }
        } else {
            MPI_Win_lock(MPI_LOCK_SHARED, rank, 0, win);
            for (i = 0; i < WIN_BUF_NUM * OP_COUNT; i++) {
                target_buf[i] = 0;
            }
            MPI_Win_unlock(rank, win);
        }

        MPI_Barrier(MPI_COMM_WORLD);

        MPI_Win_lock_all(0, win);
        if (rank != dest) {
            for (i = 0; i < my_buf_num; i++) {
                MPI_Get_accumulate(&(orig_buf[i * OP_COUNT]), 1, origin_dtp,
                                   &(result_buf[i * OP_COUNT]), 1, origin_dtp,
                                   dest, 0, 1, target_dtp, MPI_SUM, win);
                MPI_Win_flush(dest, win);
            }
        }
        MPI_Win_unlock_all(win);

        MPI_Barrier(MPI_COMM_WORLD);

        checkResults(k, &errors);

        /* ====== Part 3: test derived datatypes (origin basic, target derived) ======== */

        /* init buffers */
        if (rank != dest) {
            for (i = 0; i < my_buf_num * OP_COUNT; i++) {
                orig_buf[i] = 1;
                result_buf[i] = 0;
            }
        } else {
            MPI_Win_lock(MPI_LOCK_SHARED, rank, 0, win);
            for (i = 0; i < WIN_BUF_NUM * OP_COUNT; i++) {
                target_buf[i] = 0;
            }
            MPI_Win_unlock(rank, win);
        }

        MPI_Barrier(MPI_COMM_WORLD);

        MPI_Win_lock_all(0, win);
        if (rank != dest) {
            for (i = 0; i < my_buf_num; i++) {
                MPI_Get_accumulate(&(orig_buf[i * OP_COUNT]), OP_COUNT, MPI_INT,
                                   &(result_buf[i * OP_COUNT]), OP_COUNT, MPI_INT,
                                   dest, 0, 1, target_dtp, MPI_SUM, win);
                MPI_Win_flush(dest, win);
            }
        }
        MPI_Win_unlock_all(win);

        MPI_Barrier(MPI_COMM_WORLD);

        checkResults(k, &errors);

        /* ====== Part 4: test derived datatypes (origin derived target basic) ======== */

        /* init buffers */
        if (rank != dest) {
            for (i = 0; i < my_buf_num * OP_COUNT; i++) {
                orig_buf[i] = 1;
                result_buf[i] = 0;
            }
        } else {
            MPI_Win_lock(MPI_LOCK_SHARED, rank, 0, win);
            for (i = 0; i < WIN_BUF_NUM * OP_COUNT; i++) {
                target_buf[i] = 0;
            }
            MPI_Win_unlock(rank, win);
        }

        MPI_Barrier(MPI_COMM_WORLD);

        MPI_Win_lock_all(0, win);
        if (rank != dest) {
            for (i = 0; i < my_buf_num; i++) {
                MPI_Get_accumulate(&(orig_buf[i * OP_COUNT]), 1, origin_dtp,
                                   &(result_buf[i * OP_COUNT]), 1, origin_dtp,
                                   dest, 0, OP_COUNT, MPI_INT, MPI_SUM, win);
                MPI_Win_flush(dest, win);
            }
        }
        MPI_Win_unlock_all(win);

        MPI_Barrier(MPI_COMM_WORLD);

        checkResults(k, &errors);
    }

    MPI_Win_free(&win);

    if (rank == origin_am || rank == origin_shm) {
        MPI_Free_mem(orig_buf);
        MPI_Free_mem(result_buf);
    }

    MPI_Type_free(&origin_dtp);
    MPI_Type_free(&target_dtp);

  exit_test:
    MTest_Finalize(errors);
    return MTestReturnValue(errors);
}