/*
 mpicc -o x5 x5.c -lhdf5  ; mpiexec -n 2 ./x5
*/
#include <stdio.h>
#include <stdlib.h>
#include <hdf5.h>

typedef struct _Walker {
 int sign;
 double w;
 int age;
 int trapped;
 int elsancho;
 double *r;
  
} Walker;
                                                                                
typedef struct _Sim {
 int nelectrons;
 int nwalkers;
 Walker *w;
} Sim;
 
 
int hdf_writ(char *fname, Sim *s);
int hdf_read(char *fname, Sim *s);
 
int global_nwalkers;
int prev_nwalkers;
int local_nwalkers;

int main (int argc, char *argv[]) {
 
    int	i,j;
    int mpi_size, mpi_rank;
    Sim         univ_s,*s=NULL;
 
    MPI_Init(&argc, &argv);
    MPI_Comm_size(MPI_COMM_WORLD, &mpi_size);
    MPI_Comm_rank(MPI_COMM_WORLD, &mpi_rank);
 
/* FIXME : here I assume that all ranks have same numbers of walkers
           really need the running total from rank 0....mpi_rank 
*/
    local_nwalkers = 2;
    global_nwalkers = mpi_size*2; 
    prev_nwalkers = mpi_rank*2;

    univ_s.nelectrons = 5;
    univ_s.nwalkers = global_nwalkers;
    
    s = &univ_s;
    s->w = (Walker *)malloc((size_t)(local_nwalkers*sizeof(Walker)));
    s->w[0].r = (double *)malloc((size_t)(3*s->nelectrons*local_nwalkers*sizeof(double)));
    for(i=0;i<local_nwalkers;i++) {
     s->w[i].sign = mpi_rank+i;
     s->w[i].w = mpi_rank+i;
     s->w[i].age = mpi_rank+i;
     s->w[i].trapped = mpi_rank+i;
     s->w[i].elsancho = mpi_rank+i;
     if(i) s->w[i].r = s->w[i-1].r+3*s->nelectrons;
     for(j=0;j<3*s->nelectrons;j++) s->w[i].r[j] = i+j+1000*mpi_rank;
    }

    for(i=0;i<local_nwalkers;i++) {
     printf("init %d %d %f %d %d %d    %f %f %f\n",
      mpi_rank, s->w[i].sign, s->w[i].w,  s->w[i].age, s->w[i].trapped ,
      s->w[i].elsancho, s->w[i].r[0], s->w[i].r[1], s->w[i].r[2]);
    }

  
    hdf_write("x5.h5",s);
    free((char *)s->w[0].r);
    free((char *)s->w);
    s->nelectrons = 0;
    s->nwalkers = 0;
    hdf_read("x5.h5",s);
    hdf_write("x5o.h5",s);
    hdf_read("x5.h5",s);
    hdf_write("x5p.h5",s);

    for(i=0;i<local_nwalkers;i++) {
     printf("fini %d %d %f %d %d %d    %f %f %f\n",
      mpi_rank, s->w[i].sign, s->w[i].w,  s->w[i].age, s->w[i].trapped ,
      s->w[i].elsancho, s->w[i].r[0], s->w[i].r[1], s->w[i].r[2]);
    }

    free((char *)s->w[0].r);
    free((char *)s->w);
    s->nelectrons = 0;
    s->nwalkers = 0;




    MPI_Finalize();
    return 0;

}

int hdf_write(char *fname, Sim *s) {
    hid_t       file_id, dataset;   
    hid_t       filespace, memspace;    
    hid_t	sim_tid, walker_tid;
    hsize_t     sdim[1];
    hsize_t     wdim[1], wcount[1], woffset[1];
    hsize_t     edim[2], ecount[2], eoffset[2];
    hid_t	plist_id,attr_id;
    int         i,j,one=2;
    herr_t	status;
    double	*data;
    int mpi_size, mpi_rank;

    MPI_Comm_size(MPI_COMM_WORLD, &mpi_size);
    MPI_Comm_rank(MPI_COMM_WORLD, &mpi_rank);  


   sdim[0]=1;

   wdim[0] = global_nwalkers;
   wcount[0] = local_nwalkers;
   woffset[0] = prev_nwalkers;

   edim[0] = global_nwalkers;
   edim[1] = 3*s->nelectrons;
   ecount[0] = local_nwalkers;
   ecount[1] = 3*s->nelectrons;;
   eoffset[0] = prev_nwalkers;
   eoffset[1] = 0;

/* sim */

    plist_id = H5Pcreate(H5P_FILE_ACCESS);
    H5Pset_fapl_mpio(plist_id, MPI_COMM_WORLD, MPI_INFO_NULL);
    file_id = H5Fcreate(fname, H5F_ACC_TRUNC, H5P_DEFAULT, plist_id);
    H5Pclose(plist_id); 


    sim_tid = H5Tcreate(H5T_COMPOUND, sizeof(Sim));
    H5Tinsert(sim_tid, "nelectrons", HOFFSET(Sim, nelectrons), H5T_NATIVE_INT);
    H5Tinsert(sim_tid, "nwalkers", HOFFSET(Sim, nwalkers), H5T_NATIVE_INT);
    filespace = H5Screate_simple(1, sdim, NULL); 
    memspace = H5Screate_simple(1, sdim, NULL);
    dataset = H5Dcreate(file_id, "sim", sim_tid, memspace, H5P_DEFAULT);
    if(!mpi_rank)
     H5Dwrite(dataset, sim_tid, memspace, filespace, H5P_DEFAULT, s);
    H5Tclose(sim_tid);
    H5Dclose(dataset);
    H5Sclose(filespace);
    H5Sclose(memspace);

/* walkers */
  
    walker_tid = H5Tcreate(H5T_COMPOUND, sizeof(Walker));
    H5Tinsert(walker_tid, "sign", HOFFSET(Walker, sign ), H5T_NATIVE_INT);
    H5Tinsert(walker_tid, "w", HOFFSET(Walker, w ), H5T_NATIVE_DOUBLE);
    H5Tinsert(walker_tid, "age", HOFFSET(Walker, age ), H5T_NATIVE_INT);
    H5Tinsert(walker_tid, "trapped", HOFFSET(Walker, trapped), H5T_NATIVE_INT);
    H5Tinsert(walker_tid, "elsancho", HOFFSET(Walker, elsancho ), H5T_NATIVE_INT);
    memspace = H5Screate_simple(1, wcount, NULL);
    filespace = H5Screate_simple(1, wdim, NULL); 
    dataset = H5Dcreate(file_id, "walkers", walker_tid, filespace, H5P_DEFAULT);
    H5Sclose(filespace);
    filespace = H5Dget_space(dataset);
    H5Sselect_hyperslab(filespace, H5S_SELECT_SET, woffset, NULL, wcount, NULL);
    H5Dwrite(dataset, walker_tid, memspace, filespace, H5P_DEFAULT, s->w);
    H5Tclose(walker_tid);
    H5Dclose(dataset);
    H5Sclose(filespace);
    H5Sclose(memspace);
    H5Fclose(file_id);

/* electrons */

    data = (double *) malloc((size_t)(3*sizeof(double)*ecount[0]*ecount[1]));
    for(i=0;i<ecount[0];i++) {
     for(j=0;j<3*ecount[1];j++) {
      data[i*ecount[1]+j] = s->w[i].r[j];
     }
    }

    plist_id = H5Pcreate(H5P_FILE_ACCESS);
    H5Pset_fapl_mpio(plist_id, MPI_COMM_WORLD, MPI_INFO_NULL);
    file_id = H5Fopen(fname, H5F_ACC_RDWR, plist_id);
    H5Pclose(plist_id);
    memspace = H5Screate_simple(2, ecount, NULL);
    filespace = H5Screate_simple(2, edim, NULL); 
    dataset = H5Dcreate(file_id, "electrons", H5T_NATIVE_DOUBLE, filespace,
			H5P_DEFAULT);
    H5Sclose(filespace);
    filespace = H5Dget_space(dataset);
    H5Sselect_hyperslab(filespace, H5S_SELECT_SET, eoffset, NULL, ecount, NULL);

    plist_id = H5Pcreate(H5P_DATASET_XFER);
    H5Pset_dxpl_mpio(plist_id, H5FD_MPIO_COLLECTIVE);
    H5Dwrite(dataset, H5T_NATIVE_DOUBLE, memspace, filespace,
		      plist_id, data);
    H5Pclose(plist_id);
    H5Dclose(dataset);
    H5Sclose(filespace);
    H5Sclose(memspace);
    H5Fclose(file_id);
    free((char *) data);

}

int hdf_read(char *fname, Sim *s) {
    hid_t       file_id, dataset;   
    hid_t       filespace, memspace;    
    hid_t	sim_tid, walker_tid;
    hsize_t     sdim[1];
    hsize_t     wdim[1], wcount[1], woffset[1];
    hsize_t     edim[2], ecount[2], eoffset[2];
    hid_t	plist_id,attr_id;
    int         i,j;
    herr_t	status;
    double	*data;
    Sim		*sp;
    int mpi_size, mpi_rank;

    MPI_Comm_size(MPI_COMM_WORLD, &mpi_size);
    MPI_Comm_rank(MPI_COMM_WORLD, &mpi_rank);  


   sdim[0]=1;

/* sim */
    plist_id = H5Pcreate(H5P_FILE_ACCESS);
    H5Pset_fapl_mpio(plist_id, MPI_COMM_WORLD, MPI_INFO_NULL);
    file_id = H5Fopen(fname, H5F_ACC_RDWR, plist_id);
    H5Pclose(plist_id);
   
    dataset = H5Dopen(file_id, "sim");
    sim_tid =  H5Dget_type(dataset);
    filespace = H5Dget_space(dataset);
    memspace = H5Dget_space(dataset);

    plist_id = H5Pcreate(H5P_DATASET_XFER);
    H5Pset_dxpl_mpio(plist_id, H5FD_MPIO_COLLECTIVE);
    H5Dread(dataset, sim_tid, memspace, filespace, plist_id, s);
    H5Pclose(plist_id);
    H5Dclose(dataset);
    H5Sclose(memspace);
    H5Sclose(filespace);
    H5Fclose(file_id);

   if(s->nwalkers < mpi_size) {
    wdim[0] = mpi_size;
    wcount[0] = 1;
    woffset[0] = s->nwalkers % mpi_rank;
   } else {   /* FIXME  need to account for uneven distribution
		 if mpi_size does not evenly divide the nwalkers
                 then need to distribute and communicate back zori 
                 who has what -- Can you do this Alan? */
                 
    wdim[0] = s->nwalkers;
    wcount[0]  = s->nwalkers/mpi_size;
    woffset[0] = (s->nwalkers/mpi_size)*mpi_rank;
   }

   edim[0] = wdim[0];
   edim[1] = 3*s->nelectrons;
   ecount[0] = wcount[0];
   ecount[1] = 3*s->nelectrons;;
   eoffset[0] = wcount[0]*mpi_rank;
   eoffset[1] = 0;


/*
   MPI_Barrier(MPI_COMM_WORLD);
   printf("%d s->nelectrons = %d s->nwalkers = %d\n",
	mpi_rank, s->nelectrons,s->nwalkers);
   printf("%d sdim[0] = %lld\n", mpi_rank, sdim[0]);
   printf("%d wdim[0] = %lld  wcount[0] = %lld  woffset[0] = %lld \n",
	 mpi_rank, wdim[0],wcount[0],woffset[0]);
   printf("%d edim[0] = %lld  ecount[0] = %lld  eoffset[0] = %lld \n",
	 mpi_rank, edim[0],ecount[0],eoffset[0]);
   printf("%d edim[1] = %lld  ecount[1] = %lld  eoffset[1] = %lld \n",
	 mpi_rank, edim[1],ecount[1],eoffset[1]);
   fflush(stdout);
   MPI_Barrier(MPI_COMM_WORLD);
*/

/* walker */
    s->w = (Walker *) malloc((size_t)(sizeof(Walker)*wdim[0]));
    plist_id = H5Pcreate(H5P_FILE_ACCESS);
    H5Pset_fapl_mpio(plist_id, MPI_COMM_WORLD, MPI_INFO_NULL);
    file_id = H5Fopen(fname, H5F_ACC_RDWR, plist_id);
    H5Pclose(plist_id);
   
    dataset = H5Dopen(file_id, "walkers");
    walker_tid =  H5Dget_type(dataset);
    filespace = H5Dget_space(dataset);
    memspace = H5Screate_simple(1, wcount, NULL);
    H5Sselect_hyperslab(filespace, H5S_SELECT_SET, woffset, NULL, wcount, NULL);
    plist_id = H5Pcreate(H5P_DATASET_XFER);
    H5Pset_dxpl_mpio(plist_id, H5FD_MPIO_COLLECTIVE);
    H5Dread(dataset, walker_tid, memspace, filespace, plist_id, s->w);
    H5Pclose(plist_id);
    H5Dclose(dataset);
    H5Sclose(memspace);
    H5Sclose(filespace);
    H5Fclose(file_id);


/* electrons */
    s->w[0].r = (double *)malloc((size_t)(3*s->nelectrons*wcount[0]*sizeof(double)));
    for(i=1;i<wcount[0];i++) {
     s->w[i].r = s->w[i-1].r+3*s->nelectrons;
    }
    plist_id = H5Pcreate(H5P_FILE_ACCESS);
    H5Pset_fapl_mpio(plist_id, MPI_COMM_WORLD, MPI_INFO_NULL);
    file_id = H5Fopen(fname, H5F_ACC_RDWR, plist_id);
    H5Pclose(plist_id);
   
    dataset = H5Dopen(file_id, "electrons");
    filespace = H5Dget_space (dataset);
    H5Sselect_hyperslab(filespace, H5S_SELECT_SET, eoffset, NULL, ecount, NULL);
    memspace = H5Screate_simple(2, ecount, NULL);
    plist_id = H5Pcreate(H5P_DATASET_XFER);
    H5Pset_dxpl_mpio(plist_id, H5FD_MPIO_COLLECTIVE);
    H5Dread(dataset, H5T_NATIVE_DOUBLE, memspace, filespace, plist_id, s->w[0].r);
    H5Pclose(plist_id);
    H5Dclose(dataset);
    H5Sclose(memspace);
    H5Sclose(filespace);
    H5Fclose(file_id);

}