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

/* Test case from John Bent (ROMIO req #835)
 * Aggregation code was not handling certain access patterns when collective
 * buffering forced */

/* Uses nonblocking collective I/O.*/

#include <unistd.h>
#include <stdlib.h>
#include <mpi.h>
#include <stdio.h>
#include <string.h>

#define NUM_OBJS 4
#define OBJ_SIZE 1048576

extern char *optarg;
extern int optind, opterr, optopt;


char *prog = NULL;
int debug = 0;

static void Usage(int line)
{
    int rank;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    if (rank == 0) {
        fprintf(stderr,
                "Usage (line %d): %s [-d] [-h] -f filename\n"
                "\t-d for debugging\n"
                "\t-h to turn on the hints to force collective aggregation\n", line, prog);
    }
    exit(0);
}

static void fatal_error(int mpi_ret, MPI_Status * mpi_stat, const char *msg)
{
    fprintf(stderr, "Fatal error %s: %d\n", msg, mpi_ret);
    MPI_Abort(MPI_COMM_WORLD, -1);
}

static void print_hints(int rank, MPI_File * mfh)
{
    MPI_Info info;
    int nkeys;
    int i, dummy_int;
    char key[1024];
    char value[1024];

    MPI_Barrier(MPI_COMM_WORLD);
    if (rank == 0) {
        MPI_File_get_info(*mfh, &info);
        MPI_Info_get_nkeys(info, &nkeys);

        printf("HINTS:\n");
        for (i = 0; i < nkeys; i++) {
            MPI_Info_get_nthkey(info, i, key);
            printf("%35s -> ", key);
            MPI_Info_get(info, key, 1024, value, &dummy_int);
            printf("%s\n", value);
        }
        MPI_Info_free(&info);
    }
    MPI_Barrier(MPI_COMM_WORLD);
}

static void fill_buffer(char *buffer, int bufsize, int rank, MPI_Offset offset)
{
    memset((void *) buffer, 0, bufsize);
    snprintf(buffer, bufsize, "Hello from %d at %lld\n", rank, offset);
}

static MPI_Offset get_offset(int rank, int num_objs, int obj_size, int which_obj)
{
    MPI_Offset offset;
    offset = (MPI_Offset) rank *num_objs * obj_size + which_obj * obj_size;
    return offset;
}

static void write_file(char *target, int rank, MPI_Info * info)
{
    MPI_File wfh;
    MPI_Request *request;
    MPI_Status *mpi_stat;
    int mpi_ret;
    int i;
    char **buffer;

    request = (MPI_Request *) malloc(NUM_OBJS * sizeof(MPI_Request));
    mpi_stat = (MPI_Status *) malloc(NUM_OBJS * sizeof(MPI_Status));
    buffer = (char **) malloc(NUM_OBJS * sizeof(char *));

    if (debug)
        printf("%d writing file %s\n", rank, target);

    if ((mpi_ret = MPI_File_open(MPI_COMM_WORLD, target,
                                 MPI_MODE_WRONLY | MPI_MODE_CREATE, *info, &wfh))
        != MPI_SUCCESS) {
        fatal_error(mpi_ret, NULL, "open for write");
    }

    /* nonblocking collective write */
    for (i = 0; i < NUM_OBJS; i++) {
        MPI_Offset offset = get_offset(rank, NUM_OBJS, OBJ_SIZE, i);
        buffer[i] = (char *) malloc(OBJ_SIZE);
        fill_buffer(buffer[i], OBJ_SIZE, rank, offset);
        if (debug)
            printf("%s", buffer[i]);
        if ((mpi_ret = MPI_File_iwrite_at_all(wfh, offset, buffer[i], OBJ_SIZE,
                                              MPI_CHAR, &request[i]))
            != MPI_SUCCESS) {
            fatal_error(mpi_ret, NULL, "write");
        }
    }

    if (debug)
        print_hints(rank, &wfh);

    MPI_Waitall(NUM_OBJS, request, mpi_stat);

    if ((mpi_ret = MPI_File_close(&wfh)) != MPI_SUCCESS) {
        fatal_error(mpi_ret, NULL, "close for write");
    }
    if (debug)
        printf("%d wrote file %s\n", rank, target);

    for (i = 0; i < NUM_OBJS; i++)
        free(buffer[i]);
    free(buffer);
    free(mpi_stat);
    free(request);
}

static int reduce_corruptions(int corrupt_blocks)
{
    int mpi_ret;
    int sum;
    if ((mpi_ret = MPI_Reduce(&corrupt_blocks, &sum, 1, MPI_INT, MPI_SUM, 0,
                              MPI_COMM_WORLD)) != MPI_SUCCESS) {
        fatal_error(mpi_ret, NULL, "MPI_Reduce");
    }
    return sum;
}

static void read_file(char *target, int rank, MPI_Info * info, int *corrupt_blocks)
{
    MPI_File rfh;
    MPI_Offset *offset;
    MPI_Request *request;
    MPI_Status *mpi_stat;
    int mpi_ret;
    int i;
    char **buffer;
    char **verify_buf = NULL;

    offset = (MPI_Offset *) malloc(NUM_OBJS * sizeof(MPI_Offset));
    request = (MPI_Request *) malloc(NUM_OBJS * sizeof(MPI_Request));
    mpi_stat = (MPI_Status *) malloc(NUM_OBJS * sizeof(MPI_Status));
    buffer = (char **) malloc(NUM_OBJS * sizeof(char *));
    verify_buf = (char **) malloc(NUM_OBJS * sizeof(char *));

    if (debug)
        printf("%d reading file %s\n", rank, target);

    if ((mpi_ret = MPI_File_open(MPI_COMM_WORLD, target, MPI_MODE_RDONLY,
                                 *info, &rfh)) != MPI_SUCCESS) {
        fatal_error(mpi_ret, NULL, "open for read");
    }

    /* nonblocking collective read */
    for (i = 0; i < NUM_OBJS; i++) {
        offset[i] = get_offset(rank, NUM_OBJS, OBJ_SIZE, i);
        buffer[i] = (char *) malloc(OBJ_SIZE);
        verify_buf[i] = (char *) malloc(OBJ_SIZE);
        fill_buffer(verify_buf[i], OBJ_SIZE, rank, offset[i]);
        if (debug)
            printf("Expecting %s", verify_buf[i]);
        if ((mpi_ret = MPI_File_iread_at_all(rfh, offset[i], buffer[i],
                                             OBJ_SIZE, MPI_CHAR, &request[i]))
            != MPI_SUCCESS) {
            fatal_error(mpi_ret, NULL, "read");
        }
    }

    MPI_Waitall(NUM_OBJS, request, mpi_stat);

    /* verification */
    for (i = 0; i < NUM_OBJS; i++) {
        if (memcmp(verify_buf[i], buffer[i], OBJ_SIZE) != 0) {
            (*corrupt_blocks)++;
            printf("Corruption at %lld\n", offset[i]);
            if (debug) {
                printf("\tExpecting %s\n" "\tRecieved  %s\n", verify_buf[i], buffer[i]);
            }
        }
    }

    if ((mpi_ret = MPI_File_close(&rfh)) != MPI_SUCCESS) {
        fatal_error(mpi_ret, NULL, "close for read");
    }

    for (i = 0; i < NUM_OBJS; i++) {
        free(verify_buf[i]);
        free(buffer[i]);
    }
    free(verify_buf);
    free(buffer);
    free(mpi_stat);
    free(request);
    free(offset);
}

static void set_hints(MPI_Info * info)
{
    MPI_Info_set(*info, "romio_cb_write", "enable");
    MPI_Info_set(*info, "romio_no_indep_rw", "1");
    MPI_Info_set(*info, "cb_nodes", "1");
    MPI_Info_set(*info, "cb_buffer_size", "4194304");
}

/*
void
set_hints(MPI_Info *info, char *hints) {
    char *delimiter = " ";
    char *hints_cp  = strdup(hints);
    char *key = strtok(hints_cp, delimiter);
    char *val;
    while (key) {
        val = strtok(NULL, delimiter);
        if (debug) printf("HINT: %s = %s\n", key, val);
        if (! val) {
            Usage(__LINE__);
        }
        MPI_Info_set(*info, key, val);
        key = strtok(NULL, delimiter);
    }
    free(hints_cp);
}
*/

int main(int argc, char *argv[])
{
    int nproc = 1, rank = 0;
    char *target = NULL;
    int c;
    MPI_Info info;
    int mpi_ret;
    int corrupt_blocks = 0;

    MPI_Init(&argc, &argv);
    MPI_Comm_size(MPI_COMM_WORLD, &nproc);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);

    if ((mpi_ret = MPI_Info_create(&info)) != MPI_SUCCESS) {
        if (rank == 0)
            fatal_error(mpi_ret, NULL, "MPI_info_create.\n");
    }

    prog = strdup(argv[0]);

    if (argc > 1) {
        while ((c = getopt(argc, argv, "df:h")) != EOF) {
            switch (c) {
            case 'd':
                debug = 1;
                break;
            case 'f':
                target = strdup(optarg);
                break;
            case 'h':
                set_hints(&info);
                break;
            default:
                Usage(__LINE__);
            }
        }
        if (!target) {
            Usage(__LINE__);
        }
    }
    else {
        target = "testfile";
        set_hints(&info);
    }

    write_file(target, rank, &info);
    read_file(target, rank, &info, &corrupt_blocks);

    corrupt_blocks = reduce_corruptions(corrupt_blocks);
    if (rank == 0) {
        if (corrupt_blocks == 0) {
            fprintf(stdout, " No Errors\n");
        }
        else {
            fprintf(stdout, "%d/%d blocks corrupt\n", corrupt_blocks, nproc * NUM_OBJS);
        }
    }
    MPI_Info_free(&info);

    MPI_Finalize();
    free(prog);
    exit(0);
}