/* 
 * File:   convert-mpi.c
 * Author: Mark Howison
 *
 * Description
 * -----------
 * This program converts data stored in the Analyze 7.5 file format to
 * the H5Block file format. It can also scale the data by an integer
 * factor and add padding and chunking in all 3 dimensions.
 *
 * Created on Thu May 28 13:24:52 PDT 2009
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>
#include <mpi.h>
#include <H5Part.h>
#include <H5Block.h>
#include "analyzeIO.h"

#define H5_VERBOSITY 3
#define ONE_MEGABYTE 1048576 

#define ERROR(msg,my_rank) do{\
    fprintf (stderr, "rank %d: %s (%d): %s!\n",\
            my_rank, __FILE__, __LINE__, msg);\
    fflush (stderr);\
    exit (EXIT_FAILURE);\
}while(0);

int main (int argc, char** argv)
{
    int     i,j,k;
    int     indx0,indx1;
    int     ci,cj,ck;
    int     sx,sy,sz;
    int     px,py,pz;
    int     x,y,z;
    int     rank,nprocs;
    int     decomp[3];
    char    chunk_str[64];
    
    /* Analyze 7.5 */
    int     ndims;
    int     dims[8];
    int     npts=0;
    int*    tdims;
    float*  pdims;
    int     dType;
    int     bytesPerVoxel;
    void*   sd=NULL;
    char*   infile;
    char*   outfile;
    float*  inData;
    
    /* H5Block */
    H5PartFile*     h5file;
    h5part_int64_t  nelems;
    h5part_int64_t  chunk[3];
    h5part_int64_t  slab[3];
    h5part_int64_t  start[3];
    h5part_int64_t  end[3];
    h5part_float32_t* outData;

    MPI_Init (&argc, &argv);
    MPI_Comm_rank (MPI_COMM_WORLD, &rank);
    MPI_Comm_size (MPI_COMM_WORLD, &nprocs);
    
    if (argc < 12) {
        if (rank == 0) {
            printf ("Usage: ./convert-mpi ");
            printf ("input.hdr output.h5 sx sy sz px py pz cx cy cz\n");
            printf ("\tsx,sy,sz = scaling factors\n");
            printf ("\tpx,py,pz = padding\n");
            printf ("\tcx,cy,cz = chunking\n");
        }
        MPI_Finalize();
        exit (EXIT_SUCCESS);
    } else {
        infile = argv[1];
        outfile = argv[2];
        sx = atoi (argv[3]);
        sy = atoi (argv[4]);
        sz = atoi (argv[5]);
        px = atoi (argv[6]);
        py = atoi (argv[7]);
        pz = atoi (argv[8]);
        chunk[0] = atoi(argv[9]);
        chunk[1] = atoi(argv[10]);
        chunk[2] = atoi(argv[11]);
    }
    
    if (chunk[0] <= 0 || chunk[1] <= 0 || chunk[2] <= 0) {
        ERROR("Invalid chunk dimensions: each must be > 0",rank);
    }

    H5PartSetVerbosityLevel(H5_VERBOSITY);
    
    /************************************************************************/
    /* Read in the Analyze 7.5 data                                         */
    /************************************************************************/
    
    printf ("Reading from Analyze 7.5 file %s\n", infile);
    
    analyzeLoadData (infile, &ndims,
            &tdims, &pdims, &dType, &bytesPerVoxel, &sd);
    
    for (i=0;i<ndims;i++) {
        if (i==0) {
            npts = tdims[i];
        } else {
            npts *= tdims[i];
        }
        dims[i] = tdims[i];
        //pdims[i] = pdims[i];
    }

    if (rank == 0) printf ("Converting values to floats... ");
    
    inData = convertToFloat (sd, dType, npts);

    if (rank == 0) {
        printf ("done.\n");
        printf ("Input dimensions: (%d,%d,%d)\n", dims[0], dims[1], dims[2]);
    }

    if (rank == 0) {
        printf ("Chunk dimensions: (%ld,%ld,%ld)\n",
                chunk[0],
                chunk[1],
                chunk[2]);
    }
    
    /************************************************************************/
    /* Scaling and padding                                                  */
    /************************************************************************/

    x = (sx*dims[0] + 2*px);
    y = (sy*dims[1] + 2*py);
    z = (sz*dims[2] + 2*pz);

    if (rank == 0) {
        printf("Scaled/padded dimensions: (%d,%d,%d)\n", x, y, z);
    }

    assert (x % chunk[0] == 0);
    assert (y % chunk[1] == 0);
    assert (z % chunk[2] == 0);
    decomp[0] = x / chunk[0];
    decomp[1] = y / chunk[1];
    decomp[2] = z / chunk[2];
    if (rank == 0) {
        printf("Chunk decomposition: (%d,%d,%d)\n",
                decomp[0],
                decomp[1],
                decomp[2]);
    }

    if (decomp[0]*decomp[1]*decomp[2] != nprocs) {
        ERROR("Number of chunks != number of procs",rank);
    }

    ci = rank % decomp[0];
    cj = (rank / decomp[0]) % decomp[1];
    ck = rank / (decomp[0] * decomp[1]);

    start[0] = ci * chunk[0];
    start[1] = cj * chunk[1];
    start[2] = ck * chunk[2];

    end[0] = start[0] + chunk[0];
    end[1] = start[1] + chunk[1];
    end[2] = start[2] + chunk[2];

    if (start[0] < px) start[0] = px;
    if (start[1] < py) start[1] = py;
    if (start[2] < pz) start[2] = pz;

    if (end[0] > x - px) end[0] = x - px;
    if (end[1] > y - py) end[1] = y - py;
    if (end[2] > z - pz) end[2] = z - pz;

    if (rank == 0) printf("Creating chunks... ");

    nelems = chunk[0]*chunk[1]*chunk[2];
    outData = (h5part_float32_t*) malloc(nelems*sizeof(h5part_float32_t));
    if (!outData) {
        ERROR("Could not allocate memory for buffer",rank);
    }
    for (i=0;i<nelems;i++) outData[i] = 0.F;
    
    for (k=start[2]; k<end[2]; k++) {
        z = (k - pz) / sz;
        for (j=start[1]; j<end[1]; j++) {
            y = (j - py) / sy;
            indx1 = (k - ck*chunk[2])*chunk[1]*chunk[0]
                  + (j - cj*chunk[1])*chunk[0]
                  + (start[0] - ci*chunk[0]);
            for (i=start[0]; i<end[0]; i++) {
                x = (i - px) / sx;
                indx0 = z*dims[1]*dims[0] + y*dims[0] + x;
                outData[indx1++] = inData[indx0];
            }
        }
    }

    if (rank == 0) printf("done.\n");
    
    /************************************************************************/
    /* Write out the H5Block file                                           */
    /************************************************************************/
    
    if (rank == 0) printf("Opening H5Part file '%s' for writing.\n", outfile);
    
    h5file = H5PartOpenFileParallelAlign(
            outfile,
            H5PART_WRITE,
            MPI_COMM_WORLD,
            ONE_MEGABYTE);
    if (!h5file) ERROR("Could not open file",rank);
    
    H5PartSetStep (h5file, 0);
    
    if (rank == 0) printf ("Writing data... ");

    // use slabs as the HDF "chunk" dimension, to improve read access
    // to ghost zones
    //slab[0] = chunk[0];
    //slab[1] = chunk[1];
    //slab[2] = 1;
    slab[0] = 32;
    slab[1] = 32;
    slab[2] = 32;

    H5BlockDefine3DChunk (h5file, slab);

    start[0] = ci * chunk[0];
    start[1] = cj * chunk[1];
    start[2] = ck * chunk[2];

    end[0] = start[0] + chunk[0] - 1;
    end[1] = start[1] + chunk[1] - 1;
    end[2] = start[2] + chunk[2] - 1;

    H5BlockDefine3DFieldLayout(
            h5file,
            start[0], end[0],
            start[1], end[1],
            start[2], end[2]);
   
    H5BlockWriteFieldFloat32 (h5file, "Analyze7.5", outData);

    sprintf (chunk_str, "%ldx%ldx%ld", chunk[0], chunk[1], chunk[2]);

    H5BlockWriteFieldAttribString (
            h5file,
            "Analyze7.5",
            "ChunkDims",
            chunk_str);

    H5PartCloseFile (h5file);
    
    if (rank == 0) printf("done.\nSuccess!\n");
    
    MPI_Finalize();
    return (EXIT_SUCCESS);
}