/*
 * Copyright (C) by Argonne National Laboratory
 *     See COPYRIGHT in top-level directory
 */

#include "mpi.h"
#include "stdio.h"
#include "stdlib.h"
#include "mpitest.h"
#include <assert.h>
#include <string.h>

#define LAST_TEST 14
#define RMA_SIZE  2048
#define OFFSET_1  7
#define OFFSET_2  83
#define OFFSET_3  157

#define PUT_VAL 0xdcba97
#define ACC_VAL 10771134

/*
 * Additional tests for lock contention.  These are designed to exercise
 * some of the optimizations within MPICH, but all are valid MPI programs.
 * Tests structure includes
 *    lock local (must happen at this time since application can use load
 *                store after the lock)
 *    send message to partner
 *                                  receive message
 *                                  send ack
 *    receive ack
 *    Provide a delay so that
 *      the partner will see the
 *      conflict
 *                                  partner executes:
 *                                  lock         // Note: this may block
 *                                     rma operations (see below)
 *                                  unlock
 *
 *    unlock                        send back to partner
 *    receive from partner
 *    check for correct data
 *
 * The delay may be implemented as a ring of message communication; this
 * is likely to automatically scale the time to what is needed
 */

/* Define a datatype to be used with */
int stride = 11;
int veccount = 7;
MPI_Datatype vectype;
/* Define long RMA ops size */
int longcount = 512;
int medcount = 127;
int mednum = 4;

void RMATest(int i, MPI_Win win, int primary, int *srcbuf, int srcbufsize, int *getbuf,
             int getbufsize);
int RMACheck(int i, int *buf, MPI_Aint bufsize);
int RMACheckGet(int i, MPI_Win win, int *getbuf, MPI_Aint getsize);
void RMATestInit(int i, int *buf, MPI_Aint bufsize);

int main(int argc, char *argv[])
{
    int errs = 0;
    MPI_Win win;
    int *rmabuffer = 0, *getbuf = 0;
    MPI_Aint bufsize = 0, getbufsize = 0;
    int primary, partner, next, wrank, wsize, i;
    int ntest = LAST_TEST;
    int *srcbuf;

    MTest_Init(&argc, &argv);

    /* Determine who is responsible for each part of the test */
    MPI_Comm_rank(MPI_COMM_WORLD, &wrank);
    MPI_Comm_size(MPI_COMM_WORLD, &wsize);
    if (wsize < 3) {
        fprintf(stderr, "This test requires at least 3 processes\n");
        MPI_Abort(MPI_COMM_WORLD, 1);
    }

    primary = 0;
    partner = 1;
    next = wrank + 1;
    if (next == partner)
        next++;
    if (next >= wsize) {
        next = 0;
        if (next == partner)
            next++;
    }

    /* Determine the last test to run (by default, run them all) */
    for (i = 1; i < argc; i++) {
        if (strcmp("-ntest", argv[i]) == 0) {
            i++;
            if (i < argc) {
                ntest = atoi(argv[i]);
            } else {
                fprintf(stderr, "Missing value for -ntest\n");
                MPI_Abort(MPI_COMM_WORLD, 1);
            }
        }
    }

    MPI_Type_vector(veccount, 1, stride, MPI_INT, &vectype);
    MPI_Type_commit(&vectype);

    /* Create the RMA window */
    bufsize = 0;
    if (wrank == primary) {
        bufsize = RMA_SIZE;
        MPI_Alloc_mem(bufsize * sizeof(int), MPI_INFO_NULL, &rmabuffer);
    } else if (wrank == partner) {
        getbufsize = RMA_SIZE;
        getbuf = (int *) malloc(getbufsize * sizeof(int));
        if (!getbuf) {
            fprintf(stderr, "Unable to allocated %d bytes for getbuf\n", (int) getbufsize);
            MPI_Abort(MPI_COMM_WORLD, 1);
        }
    }
    srcbuf = malloc(RMA_SIZE * sizeof(*srcbuf));
    assert(srcbuf);

    MPI_Win_create(rmabuffer, bufsize * sizeof(int), sizeof(int), MPI_INFO_NULL,
                   MPI_COMM_WORLD, &win);

    /* Run a sequence of tests */
    for (i = 0; i <= ntest; i++) {
        if (wrank == primary) {
            MTestPrintfMsg(0, "Test %d\n", i);
            /* Because this lock is local, it must return only when the
             * lock is acquired */
            MPI_Win_lock(MPI_LOCK_EXCLUSIVE, 0, primary, win);
            RMATestInit(i, rmabuffer, bufsize);
            MPI_Send(MPI_BOTTOM, 0, MPI_INT, partner, i, MPI_COMM_WORLD);
            MPI_Send(MPI_BOTTOM, 0, MPI_INT, next, i, MPI_COMM_WORLD);
            MPI_Recv(MPI_BOTTOM, 0, MPI_INT, MPI_ANY_SOURCE, i, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            MPI_Win_unlock(primary, win);
            MPI_Recv(MPI_BOTTOM, 0, MPI_INT, partner, i, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            errs += RMACheck(i, rmabuffer, bufsize);
        } else if (wrank == partner) {
            MPI_Recv(MPI_BOTTOM, 0, MPI_INT, primary, i, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            MPI_Win_lock(MPI_LOCK_EXCLUSIVE, 0, primary, win);
            RMATest(i, win, primary, srcbuf, RMA_SIZE, getbuf, getbufsize);
            MPI_Win_unlock(primary, win);
            errs += RMACheckGet(i, win, getbuf, getbufsize);
            MPI_Send(MPI_BOTTOM, 0, MPI_INT, primary, i, MPI_COMM_WORLD);
        } else {
            MPI_Recv(MPI_BOTTOM, 0, MPI_INT, MPI_ANY_SOURCE, i, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            MPI_Send(MPI_BOTTOM, 0, MPI_INT, next, i, MPI_COMM_WORLD);
        }
    }

    if (rmabuffer) {
        MPI_Free_mem(rmabuffer);
    }
    if (getbuf) {
        free(getbuf);
    }
    free(srcbuf);
    MPI_Win_free(&win);
    MPI_Type_free(&vectype);

    MTest_Finalize(errs);
    return MTestReturnValue(errs);
}

/* Perform the tests.
 *
 * The srcbuf must be passed in because the buffer must remain valid
 * until the subsequent unlock call. */
void RMATest(int i, MPI_Win win, int primary, int *srcbuf, int srcbufsize, int *getbuf,
             int getbufsize)
{
    int j, k;
    int *source = srcbuf;
    assert(srcbufsize == RMA_SIZE);

    for (j = 0; j < srcbufsize; j++)
        source[j] = -j;

    switch (i) {
        case 0:        /* Single short put (1 word at OFFSET_1) */
            source[0] = PUT_VAL;
            MPI_Put(source, 1, MPI_INT, primary, OFFSET_1, 1, MPI_INT, win);
            break;
        case 1:        /* Single short accumulate (1 word of value 17 at OFFSET_2) */
            source[0] = ACC_VAL;
            MPI_Accumulate(source, 1, MPI_INT, primary, OFFSET_2, 1, MPI_INT, MPI_SUM, win);
            break;
        case 2:        /* Single short get (1 word at OFFSET_3) */
            getbuf[0] = -1;
            MPI_Get(getbuf, 1, MPI_INT, primary, OFFSET_3, 1, MPI_INT, win);
            break;
        case 3:        /* Datatype single put (strided put) */
            for (j = 0; j < veccount; j++) {
                source[j * stride] = PUT_VAL + j;
            }
            MPI_Put(source, 1, vectype, primary, OFFSET_1, 1, vectype, win);
            break;
        case 4:        /* Datatype single accumulate (strided acc) */
            for (j = 0; j < veccount; j++) {
                source[j * stride] = ACC_VAL + j;
            }
            MPI_Accumulate(source, 1, vectype, primary, OFFSET_2, 1, vectype, MPI_SUM, win);
            break;
        case 5:        /* Datatype single get (strided get) */
            for (j = 0; j < veccount; j++) {
                getbuf[j] = -j;
            }
            MPI_Get(getbuf, 1, vectype, primary, OFFSET_3, 1, vectype, win);
            break;
        case 6:        /* a few small puts (like strided put, but 1 word at a time) */
            for (j = 0; j < veccount; j++) {
                source[j * stride] = PUT_VAL + j;
            }
            for (j = 0; j < veccount; j++) {
                MPI_Put(source + j * stride, 1, MPI_INT, primary,
                        OFFSET_1 + j * stride, 1, MPI_INT, win);
            }
            break;
        case 7:        /* a few small accumulates (like strided acc, but 1 word at a time) */
            for (j = 0; j < veccount; j++) {
                source[j * stride] = ACC_VAL + j;
            }
            for (j = 0; j < veccount; j++) {
                MPI_Accumulate(source + j * stride, 1, MPI_INT, primary,
                               OFFSET_2 + j * stride, 1, MPI_INT, MPI_SUM, win);
            }
            break;
        case 8:        /* a few small gets (like strided get, but 1 word at a time) */
            for (j = 0; j < veccount; j++) {
                getbuf[j * stride] = -j;
            }
            for (j = 0; j < veccount; j++) {
                MPI_Get(getbuf + j * stride, 1, MPI_INT, primary,
                        OFFSET_3 + j * stride, 1, MPI_INT, win);
            }
            break;
        case 9:        /* Single long put (OFFSET_1) */
            for (j = 0; j < longcount; j++)
                source[j] = j;
            MPI_Put(source, longcount, MPI_INT, primary, OFFSET_1, longcount, MPI_INT, win);
            break;
        case 10:       /* Single long accumulate (OFFSET_2) */
            for (j = 0; j < longcount; j++)
                source[j] = j;
            MPI_Accumulate(source, longcount, MPI_INT, primary,
                           OFFSET_2, longcount, MPI_INT, MPI_SUM, win);
            break;
        case 11:       /* Single long get (OFFSET_3) */
            for (j = 0; j < longcount; j++)
                getbuf[j] = -j;
            MPI_Get(getbuf, longcount, MPI_INT, primary, OFFSET_3, longcount, MPI_INT, win);
            break;
        case 12:       /* a few long puts (start at OFFSET_1, medcount) */
            for (j = 0; j < mednum; j++) {
                for (k = 0; k < medcount; k++) {
                    source[j * medcount + k] = j * 2 * medcount + k;
                }
                MPI_Put(source + j * medcount, medcount, MPI_INT, primary,
                        OFFSET_1 + j * 2 * medcount, medcount, MPI_INT, win);
            }
            break;
        case 13:       /* a few long accumulates (start at OFFSET_2, medcount) */
            for (j = 0; j < mednum; j++) {
                for (k = 0; k < medcount; k++) {
                    source[j * medcount + k] = ACC_VAL + j * 2 * medcount + k;
                }
                MPI_Accumulate(source + j * medcount, medcount, MPI_INT, primary,
                               OFFSET_2 + j * 2 * medcount, medcount, MPI_INT, MPI_SUM, win);
            }
            break;
        case 14:       /* a few long gets (start at OFFSET_3, medcount) */
            for (j = 0; j < mednum; j++) {
                for (k = 0; k < medcount; k++) {
                    getbuf[j * medcount + k] = -(j * medcount + k);
                }
                MPI_Get(getbuf + j * medcount, medcount, MPI_INT, primary,
                        OFFSET_3 + j * 2 * medcount, medcount, MPI_INT, win);
            }
            break;
    }
}

int RMACheck(int i, int *buf, MPI_Aint bufsize)
{
    int j, k;
    int errs = 0;

    switch (i) {
        case 0:        /* Single short put (1 word at OFFSET_1) */
            if (buf[OFFSET_1] != PUT_VAL) {
                errs++;
                printf("case 0: value is %d should be %d\n", buf[OFFSET_1], PUT_VAL);
            }
            break;
        case 1:        /* Single short accumulate (1 word of value 17 at OFFSET_2) */
            if (buf[OFFSET_2] != ACC_VAL + OFFSET_2) {
                errs++;
                printf("case 1: value is %d should be %d\n", buf[OFFSET_2], ACC_VAL + OFFSET_2);
            }
            break;
        case 2:        /* Single short get (1 word at OFFSET_3) */
            /* See RMACheckGet */
            break;
        case 3:        /* Datatype single put (strided put) */
        case 6:        /* a few small puts (like strided put, but 1 word at a time) */
            for (j = 0; j < veccount; j++) {
                int pos = j * stride + OFFSET_1;
                if (buf[pos] != PUT_VAL + j) {
                    errs++;
                    printf("case %d: buf[%d] is %d should be %d\n", i, pos, buf[pos], PUT_VAL + j);
                }
            }
            break;
        case 4:        /* Datatype single accumulate (strided acc) */
        case 7:        /* a few small accumulates (like strided acc, but 1 word at a time) */
            for (j = 0; j < veccount; j++) {
                int pos = j * stride + OFFSET_2;
                if (buf[pos] != ACC_VAL + j + OFFSET_2 + j * stride) {
                    errs++;
                    printf("case %d: buf[%d] is %d should be %d\n", i, pos, buf[pos],
                           ACC_VAL + j + OFFSET_2 + j * stride);
                }
            }
            break;
        case 5:        /* Datatype single get (strided get) */
        case 8:        /* a few small gets (like strided get, but 1 word at a time) */
            /* See RMACheckGet */
            break;
        case 9:        /* Single long put (OFFSET_1) */
            for (j = 0; j < longcount; j++) {
                if (buf[OFFSET_1 + j] != j) {
                    errs++;
                    printf("case 9: value is %d should be %d\n", buf[OFFSET_1 + j], OFFSET_1 + j);
                }
            }
            break;
        case 10:       /* Single long accumulate (OFFSET_2) */
            for (j = 0; j < longcount; j++) {
                if (buf[OFFSET_2 + j] != OFFSET_2 + j + j) {
                    errs++;
                    printf("case 10: value is %d should be %d\n", buf[OFFSET_2 + j],
                           OFFSET_2 + j + j);
                }
            }
            break;
        case 11:       /* Single long get (OFFSET_3) */
            /* See RMACheckGet */
            break;
        case 12:       /* a few long puts (start at OFFSET_1, medcount) */
            for (j = 0; j < mednum; j++) {
                for (k = 0; k < medcount; k++) {
                    if (buf[OFFSET_1 + j * 2 * medcount + k] != j * 2 * medcount + k) {
                        errs++;
                        printf("case 12: value is %d should be %d\n",
                               buf[OFFSET_1 + j * 2 * medcount + k], j * 2 * medcount + k);
                    }
                }
            }
            break;
        case 13:       /* a few long accumulates (start at OFFSET_2, medcount) */
            for (j = 0; j < mednum; j++) {
                for (k = 0; k < medcount; k++) {
                    if (buf[OFFSET_2 + j * 2 * medcount + k] !=
                        OFFSET_2 + 2 * j * 2 * medcount + 2 * k + ACC_VAL) {
                        errs++;
                        printf("case 13: value is %d should be %d\n",
                               buf[OFFSET_2 + j * 2 * medcount + k],
                               OFFSET_2 + 2 * j * 2 * medcount + k + ACC_VAL);
                    }
                }
            }
            break;
        case 14:       /* a few long gets (start at OFFSET_3, medcount) */
            /* See RMACheckGet */
            break;
        default:
            fprintf(stderr, "Unrecognized case %d\n", i);
            errs++;
            break;
    }
    return errs;
}

int RMACheckGet(int i, MPI_Win win, int *getbuf, MPI_Aint getsize)
{
    int errs = 0;
    int j, k;

    /* */
    switch (i) {
        case 0:        /* Single short put (1 word at OFFSET_1) */
            break;
        case 1:        /* Single short accumulate (1 word of value 17 at OFFSET_2) */
            break;
        case 2:        /* Single short get (1 word at OFFSET_3) */
            if (getbuf[0] != OFFSET_3) {
                errs++;
                printf("case 2: value is %d should be %d\n", getbuf[0], OFFSET_3);
            }
            break;
        case 3:        /* Datatype single put (strided put) */
            break;
        case 4:        /* Datatype single accumulate (strided acc) */
            break;
        case 5:        /* Datatype single get (strided get) */
        case 8:        /* a few small gets (like strided get, but 1 word at a time) */
            for (j = 0; j < veccount; j++) {
                if (getbuf[j * stride] != OFFSET_3 + j * stride) {
                    errs++;
                    printf("case %d: value is %d should be %d\n", i,
                           getbuf[j * stride], OFFSET_3 + j * stride);
                }
            }

            break;
        case 6:        /* a few small puts (like strided put, but 1 word at a time) */
            break;
        case 7:        /* a few small accumulates (like strided acc, but 1 word at a time) */
            break;
        case 9:        /* Single long put (OFFSET_1) */
            break;
        case 10:       /* Single long accumulate (OFFSET_2) */
            break;
        case 11:       /* Single long get (OFFSET_3) */
            for (j = 0; j < longcount; j++) {
                if (getbuf[j] != OFFSET_3 + j) {
                    errs++;
                    printf("case 11: value is %d should be %d\n", getbuf[j], OFFSET_3 + j);
                }
            }
            break;
        case 12:       /* a few long puts (start at OFFSET_1, medcount) */
            break;
        case 13:       /* a few long accumulates (start at OFFSET_2, medcount) */
            break;
        case 14:       /* a few long gets (start at OFFSET_3, medcount) */
            for (j = 0; j < mednum; j++) {
                for (k = 0; k < medcount; k++) {
                    if (getbuf[j * medcount + k] != OFFSET_3 + j * 2 * medcount + k) {
                        errs++;
                        printf("case 14: buf[%d] value is %d should be %d\n",
                               j * medcount + k,
                               getbuf[j * medcount + k], OFFSET_3 + j * 2 * medcount + k);
                    }
                }
            }
            break;
        default:
            fprintf(stderr, "Unrecognized case %d\n", i);
            errs++;
            break;
    }
    return errs;
}


void RMATestInit(int i, int *buf, MPI_Aint bufsize)
{
    int j;
    for (j = 0; j < bufsize; j++) {
        buf[j] = j;
    }
}