/*@@
   @file      DumpVar.c
   @date      Fri May 21 1999
   @author    Thomas Radke
   @desc
              Routines for writing variables into HDF5 data or checkpoint files.
              These routines are used by other HDF5 I/O methods.
   @enddesc
   @version   $Id$
 @@*/


#include <stdlib.h>

#include "cctk.h"
#include "cctk_Parameters.h"
#include "CactusPUGH/PUGH/src/include/pugh.h"
#include "CactusPUGH/PUGHSlab/src/NewPUGHSlab.h"
#include "CactusBase/IOUtil/src/ioGH.h"
#include "ioHDF5UtilGH.h"

/* the rcs ID and its dummy function to use it */
static const char *rcsid = "$Header$";
CCTK_FILEVERSION(BetaThorns_IOHDF5Util_DumpVar_c)


/********************************************************************
 ********************    Macro Definitions   ************************
 ********************************************************************/
/* uncomment the following line to enable some debugging output */
/* #define IOHDF5UTIL_DEBUG 1 */

/* tag base for MPI messages (this may break on more than 2000 processors) */
#define MPITAGBASE 20000


/********************************************************************
 ********************    Internal Routines   ************************
 ********************************************************************/
static int WriteGS (const cGH *GH, const ioSlab *slab, const char *name,
                    hid_t file);
static int WriteGA (const cGH *GH, const ioSlab *slab, const char *name,
                    hid_t file);
static int GetLocalHyperslab (const cGH *GH, const ioSlab *slab, void **data,
                              int *free_data);
static void WriteData (const cGH *GH, const ioSlab *slab, const char *name,
                       const void *data, int proc, hid_t file);
#if defined(CCTK_MPI) && defined(H5_HAVE_PARALLEL)
static void WriteDataCollective (const cGH *GH, const ioSlab *slab,
                                 const char *name, const void *data,hid_t file);
#endif


/*@@
   @routine    IOHDF5Util_DumpVar
   @date       16 Apr 1999
   @author     Thomas Radke
   @desc
               Generic dump routine, just calls the type-specific dump routine.
   @enddesc

   @calls      WriteGS
               WriteGA

   @var        GH
   @vdesc      pointer to CCTK grid hierarchy
   @vtype      const cGH *
   @vio        in
   @endvar
   @var        slab
   @vdesc      reference to the I/O hyperslab description
   @vtype      const ioSlab *
   @vio        in
   @endvar
   @var        file
   @vdesc      the HDF5 file handle
   @vtype      hid_t
   @vio        in
   @endvar

   @returntype int
   @returndesc
               -1 if variable has unknown group type
               or return code of either @seeroutine WriteGS or
                 @seeroutine WriteGA
   @endreturndesc
@@*/
int IOHDF5Util_DumpVar (const cGH *GH, const ioSlab *slab, hid_t file)
{
  int gtype, retval;
  char *fullname, *objectname;


  /* build the unique name for the file object to write */
  fullname = CCTK_FullName (slab->vindex);
  objectname = (char *) malloc (strlen (fullname) + 80);
  sprintf (objectname, "%s timelevel %d at iteration %d",
           fullname, slab->timelevel, GH->cctk_iteration);

  /* check whether the object already exists */
  if (slab->check_exist && file >= 0)
  {
    H5E_BEGIN_TRY
    {
      H5Gunlink (file, objectname);
    } H5E_END_TRY;
  }

  /* branch to the appropriate dump routine */
  gtype = CCTK_GroupTypeFromVarI (slab->vindex);
  if (gtype == CCTK_SCALAR)
  {
    retval = WriteGS (GH, slab, objectname, file);
  }
  else if (gtype == CCTK_ARRAY || gtype == CCTK_GF)
  {
    retval = WriteGA (GH, slab, objectname, file);
  }
  else
  {
    CCTK_VWarn (1, __LINE__, __FILE__, CCTK_THORNSTRING,
                "Invalid group type %d for variable '%s'", gtype, fullname);
    retval = -1;
  }

  /* clean up */
  free (objectname);
  free (fullname);

  return (retval);
}


/********************************************************************
 ********************    Internal Routines   ************************
 ********************************************************************/
/*@@
   @routine    WriteGS
   @date       May 21 1999
   @author     Thomas Radke
   @desc
               Writes a CCTK_SCALAR type variable into a HDF5 file.
   @enddesc
   @var        GH
   @vdesc      pointer to CCTK grid hierarchy
   @vtype      const cGH *
   @vio        in
   @endvar
   @var        slab
   @vdesc      reference to the I/O hyperslab description
   @vtype      const ioSlab *
   @vio        in
   @endvar
   @var        name
   @vdesc      name of the dataset to write
   @vtype      const char *
   @vio        in
   @endvar
   @var        file
   @vdesc      the HDF5 file to dump to
   @vtype      hid_t
   @vio        in
   @endvar

   @returntype int
   @returndesc
               0 for success, or -1 if file handle is invalid
   @endreturndesc
@@*/
static int WriteGS (const cGH *GH, const ioSlab *slab, const char *name,
                    hid_t file)
{
  ioGH *ioUtilGH;
  ioHDF5UtilGH *myGH;
  hid_t dataset, hdf5type;


  ioUtilGH = (ioGH *) CCTK_GHExtension (GH, "IO");
  myGH = (ioHDF5UtilGH *) CCTK_GHExtension (GH, "IOHDF5Util");

  /* only I/O processors write data */
  if (CCTK_MyProc (GH) != ioUtilGH->ioproc)
  {
    return (0);
  }

  /* check if file handle is invalid */
  if (file < 0)
  {
    return (-1);
  }

  hdf5type = IOHDF5Util_DataType (myGH, slab->hdatatype);
  HDF5_ERROR (dataset = H5Dcreate (file, name, hdf5type, myGH->scalar_dataspace,
                                   H5P_DEFAULT));
  HDF5_ERROR (H5Dwrite (dataset, hdf5type, H5S_ALL, H5S_ALL, H5P_DEFAULT,
                        CCTK_VarDataPtrI (GH, slab->timelevel, slab->vindex)));

  IOHDF5Util_DumpCommonAttributes (GH, slab, dataset);

  HDF5_ERROR (H5Dclose (dataset));

  return (0);
}


/*@@
   @routine    WriteGA
   @date       May 21 1999
   @author     Paul Walker
   @desc
               Writes a grid array into a HDF5 file.
   @enddesc

   @calls      GetLocalHyperslab
               WriteData
               WriteDataCollective

   @var        GH
   @vdesc      pointer to CCTK grid hierarchy
   @vtype      const cGH *
   @vio        in
   @endvar
   @var        slab
   @vdesc      reference to the I/O hyperslab description
   @vtype      const ioSlab *
   @vio        in
   @endvar
   @var        name
   @vdesc      name of the dataset to write
   @vtype      const char *
   @vio        in
   @endvar
   @var        file
   @vdesc      the HDF5 file to dump to
   @vtype      hid_t
   @vio        in
   @endvar

   @returntype int
   @returndesc
               0 for success
               or returncode of @seeroutine GetLocalHyperslab
   @endreturndesc
@@*/
static int WriteGA (const cGH *GH, const ioSlab *slab, const char *name,
                    hid_t file)
{
  ioGH *ioUtilGH;
  ioHDF5UtilGH *myGH;
  int myproc, nprocs, retval;
  void *data;    /* The data pointer to dump ... */
  int free_data; /* and whether it needs freeing */
#ifdef CCTK_MPI
  void *tmpd;
  int i, j, incoming, outgoing;
  MPI_Comm comm;
  MPI_Status ms;
  MPI_Datatype mpitype;
#endif
  DECLARE_CCTK_PARAMETERS


  myGH = (ioHDF5UtilGH *) CCTK_GHExtension (GH, "IOHDF5Util");
  ioUtilGH = (ioGH *) CCTK_GHExtension (GH, "IO");

#ifdef CCTK_MPI
  mpitype = PUGH_MPIDataType (PUGH_pGH (GH), slab->hdatatype);
#endif

  /* get the pointer to the data we want to output (includes downsampling) */
  retval = GetLocalHyperslab (GH, slab, &data, &free_data);
  if (retval < 0)
  {
    return (retval);
  }

  myproc = CCTK_MyProc (GH);
  nprocs = CCTK_nProcs (GH);

#ifdef CCTK_MPI
#ifdef H5_HAVE_PARALLEL
  if (ioUtilGH->unchunked)
  {
    WriteDataCollective (GH, slab, name, data, file);
    if (free_data)
    {
      free (data);
    }
    return (0);
  }
#endif
#endif

  /* dump data held on I/O processor */
  if (myproc == ioUtilGH->ioproc)
  {
    WriteData (GH, slab, name, data, myproc, file);
  }
#ifdef CCTK_MPI
  comm = PUGH_pGH (GH)->PUGH_COMM_WORLD;
  if (myproc == ioUtilGH->ioproc)
  {
    /* dump data from all other processors */
    for (i = myproc+1; i < myproc+ioUtilGH->ioproc_every && i < nprocs; i++)
    {
      /* receive geometry (this assumes the geometry arrays
         to be contiguous starting at hoffset */
      CACTUS_MPI_ERROR (MPI_Recv (slab->hoffset, 3*slab->hdim, PUGH_MPI_INT, i,
                                  2*i + MPITAGBASE + 1, comm, &ms));

      incoming = 1;
      for (j = 0; j < slab->hdim; j++)
      {
        incoming *= slab->hsize_chunk[slab->hdim + j];
      }

      /* receive data */
      tmpd = NULL;
      if (incoming > 0)
      {
        tmpd = malloc (incoming * CCTK_VarTypeSize (slab->hdatatype));
        CACTUS_MPI_ERROR (MPI_Recv (tmpd, incoming, mpitype, i,
                                    2*i + MPITAGBASE, comm, &ms));
      }

      /* write data */
      WriteData (GH, slab, name, tmpd, i, file);

      if (tmpd)
      {
        free (tmpd);
      }

    } /* end loop over processors */

  }
  else
  {
    /*
     * send the geometry and data from non-I/O processors to the I/O processors
     */
    outgoing = 1;
    for (i = 0; i < slab->hdim; i++)
    {
      outgoing *= slab->hsize_chunk[slab->hdim + i];
    }

    /* send geometry (this assumes the geometry arrays
       to be contiguous starting at hoffset */
    CACTUS_MPI_ERROR (MPI_Send (slab->hoffset, 3*slab->hdim, PUGH_MPI_INT,
                                ioUtilGH->ioproc, 2*myproc + MPITAGBASE + 1,
                                comm));
    /* send data */
    if (outgoing > 0)
    {
      CACTUS_MPI_ERROR (MPI_Send (data, outgoing, mpitype,
                                  ioUtilGH->ioproc, 2*myproc + MPITAGBASE,
                                  comm));
#ifdef IOHDF5UTIL_DEBUG
      printf ("Processor %d sent %d data points\n", myproc, outgoing);
#endif
    }
  }

  /* wait for every processor to catch up */
  CCTK_Barrier (GH);
#endif

  /* free allocated resources */
  if (free_data)
  {
    free (data);
  }

  return (retval);
}


static int GetLocalHyperslab (const cGH *GH, const ioSlab *slab, void **data,
                              int *free_data)
{
  int i, retval;
  char *fullname;


#if 0
  mapping = Hyperslab_DefineLocalMappingByIndex (GH, vindex, slab->hdim,
                                                 slab->direction, slab->origin,
                                                 slab->extent, slab->downsample,
                                                 -1, NULL, slab->hsize_local,
                                                 slab->hsize_global,
                                                 slab->hoffset_global);
           const cGH          *GH,
           CCTK_INT            vindex,
           CCTK_INT            hdim,
           const CCTK_INT     *direction         /* vdim*hdim */,
           const CCTK_INT     *origin            /* vdim */,
           const CCTK_INT     *extent            /* hdim */,
           const CCTK_INT     *downsample        /* hdim */,
           CCTK_INT            table_handle,
           t_hslabConversionFn conversion_fn,
           CCTK_INT           *hsize_local,      /* hdim */
           CCTK_INT           *hsize_global,     /* hdim */
           CCTK_INT           *hoffset_global    /* hdim */);

  if (mapping < 0)
  {
    fullname = CCTK_FullName (vindex);
    CCTK_VWarn (1, __LINE__, __FILE__, CCTK_THORNSTRING,
                "Failed to extract hyperslab for variable '%s'", fullname);
    free (fullname);
  }

#else
{
  int *hsizes, *hsizes_global, *hsizes_offset; /* geometry information */

  /* allocate array to hold size information of request */
  hsizes        = (int *) malloc (3 * slab->hdim * sizeof (int));
  hsizes_global = hsizes + 1*slab->hdim;
  hsizes_offset = hsizes + 2*slab->hdim;

  retval = NewHyperslab_GetLocalHyperslab (GH, slab->vindex, slab->timelevel, slab->hdim,
                                      slab->hdatatype, NULL, slab->origin,
                                      slab->direction, slab->extent,
                                      slab->downsample, data, free_data,
                                      hsizes, hsizes_global, hsizes_offset);
  if (retval == 0)
  {
    for (i = 0; i < slab->hdim; i++)
    {
      slab->hoffset[i] = hsizes_offset[i];
      slab->hsize[i] = hsizes_global[i];
      slab->hsize_chunk[i] = hsizes[i];
    }
  }
  else
  {
    fullname = CCTK_FullName (slab->vindex);
    CCTK_VWarn (1, __LINE__, __FILE__, CCTK_THORNSTRING,
                "Failed to extract hyperslab for variable '%s'", fullname);
    free (fullname);
  }

  free (hsizes);
}
#endif

  return (retval);
}


/*@@
   @routine    WriteData
   @author     Thomas Radke
   @date       May 21 1999
   @desc
               All I/O processors dump the data of their group into the file,
               either as one chunk per processor or unchunked
               (depending on parameter "unchunked").
   @enddesc

   @calls      IOHDF5Util_DumpCommonAttributes

   @var        GH
   @vdesc      pointer to CCTK grid hierarchy
   @vtype      const cGH *
   @vio        in
   @endvar
   @var        slab
   @vdesc      reference to the I/O hyperslab description
   @vtype      const ioSlab *
   @vio        in
   @endvar
   @var        name
   @vdesc      name of the dataset to write
   @vtype      const char *
   @vio        in
   @endvar
   @var        data
   @vdesc      pointer to the chunk to dump
   @vtype      const void *
   @vio        in
   @endvar
   @var        proc
   @vdesc      the processor which's chunk is to be dumped
   @vtype      int
   @vio        in
   @endvar
   @var        file
   @vdesc      the HDF5 file handle
   @vtype      hid_t
   @vio        in
   @endvar
@@*/
static void WriteData (const cGH *GH, const ioSlab *slab, const char *name,
                       const void *data, int proc, hid_t file)
{
  int i, myproc;
  ioGH *ioUtilGH;
  ioHDF5UtilGH *myGH;
  hid_t hdf5type, group, dataset, memspace, filespace, plist;
  char *chunkname;
  hssize_t *chunk_origin;
  hsize_t *chunk_dims, *file_dims;
  hsize_t buffersize;
  DECLARE_CCTK_PARAMETERS


  /* immediately return if file handle is invalid */
  if (file < 0)
  {
    return;
  }

  ioUtilGH = (ioGH *) CCTK_GHExtension (GH, "IO");
  myGH = (ioHDF5UtilGH *) CCTK_GHExtension (GH, "IOHDF5Util");

  /* copy the size arrays from CCTK_INT to appropriate types
     note that HDF5 wants elements in reverse order */
  chunk_origin = (hssize_t *) malloc (slab->hdim * sizeof (hssize_t));
  chunk_dims   = (hsize_t *)  malloc (2*slab->hdim * sizeof (hsize_t));
  file_dims    = chunk_dims + slab->hdim;
  for (i = 0; i < slab->hdim; i++)
  {
    chunk_origin[i] = slab->hoffset[slab->hdim - 1 - i];
    file_dims   [i] = slab->hsize[slab->hdim - 1 - i];
    chunk_dims  [i] = slab->hsize_chunk[slab->hdim - 1 - i];
  }

  myproc = CCTK_MyProc (GH);

  memspace = -1;
  if (data)
  {
    /* create the memspace according to chunk dims */
    HDF5_ERROR (memspace = H5Screate_simple (slab->hdim, chunk_dims, NULL));
  }

  hdf5type = IOHDF5Util_DataType (myGH, slab->hdatatype);
  if (ioUtilGH->unchunked)
  {
    /* create the (global) filespace and set the hyperslab for the chunk */
    HDF5_ERROR (filespace = H5Screate_simple (slab->hdim, file_dims, NULL));
    HDF5_ERROR (H5Sselect_hyperslab (filespace, H5S_SELECT_SET, chunk_origin,
                                     NULL, chunk_dims, NULL));

    /* the I/O processor creates the dataset and adds the common attributes
       when writing its own data, otherwise the dataset is reopened */
    if (proc == myproc)
    {
      HDF5_ERROR (plist = H5Pcreate (H5P_DATASET_CREATE));
      /* enable compression for chunked dataset if compression was requested */
      if (compression_level)
      {
        HDF5_ERROR (H5Pset_chunk (plist, slab->hdim, chunk_dims));
        HDF5_ERROR (H5Pset_deflate (plist, compression_level));
      }
      HDF5_ERROR (dataset = H5Dcreate (file, name, hdf5type, filespace, plist));
      HDF5_ERROR (H5Pclose (plist));
      IOHDF5Util_DumpCommonAttributes (GH, slab, dataset);
    }
    else
    {
      HDF5_ERROR (dataset = H5Dopen (file, name));
    }

    if (memspace >= 0)
    {
      /* increase the buffer size if the default isn't sufficient */
      HDF5_ERROR (plist = H5Pcreate (H5P_DATASET_XFER));
      buffersize = H5Dget_storage_size (dataset);
      if (buffersize > H5Pget_buffer (plist, NULL, NULL))
      {
        HDF5_ERROR (H5Pset_buffer (plist, buffersize, NULL, NULL));
      }

      /* write the data */
      HDF5_ERROR (H5Dwrite (dataset, hdf5type, memspace, filespace, plist,
                            data));

      /* close the transfer property list */
      HDF5_ERROR (H5Pclose (plist));
    }
    /* close the file dataspace */
    HDF5_ERROR (H5Sclose (filespace));
  }
  else
  {
    /* the I/O processor creates the chunk group and adds common attributes */
    if (proc == myproc)
    {
      HDF5_ERROR (group = H5Gcreate (file, name, 0));
      IOHDF5Util_DumpCommonAttributes (GH, slab, group);
      HDF5_ERROR (H5Gclose (group));
    }

    dataset = -1;
    if (memspace >= 0)
    {
      /* now the chunk dataset for each processor is created within the group */
      chunkname = (char *) malloc (strlen (name) + 20);
      sprintf (chunkname, "%s/chunk%d", name, proc - myproc);
      HDF5_ERROR (plist = H5Pcreate (H5P_DATASET_CREATE));
      /* enable compression for chunked dataset if compression was requested */
      if (compression_level)
      {
        HDF5_ERROR (H5Pset_chunk (plist, slab->hdim, chunk_dims));
        HDF5_ERROR (H5Pset_deflate (plist, compression_level));
      }
      /* create the chunk dataset and dump the chunk data */
      HDF5_ERROR (dataset = H5Dcreate (file, chunkname, hdf5type, memspace,
                                       plist));
      HDF5_ERROR (H5Pclose (plist));
      HDF5_ERROR (H5Dwrite (dataset, hdf5type, H5S_ALL, H5S_ALL, H5P_DEFAULT,
                            data));

      /* add the "origin" attribute for the chunk */
      WRITE_ATTRIBUTE ("chunk_origin", slab->hoffset, dataset, myGH, slab->hdim,
                       HDF5_INT);

      free (chunkname);
    }
  }

  /* close the dataset and the memspace */
  if (dataset >= 0)
  {
    HDF5_ERROR (H5Dclose (dataset));
  }
  if (memspace >= 0)
  {
    HDF5_ERROR (H5Sclose (memspace));
  }

  /* free allocated resources */
  free (chunk_origin);
  free (chunk_dims);
}


#if defined(CCTK_MPI) && defined(H5_HAVE_PARALLEL)
/*@@
   @routine    WriteDataCollective
   @author     Thomas Radke
   @date       May 21 1999
   @desc
               All processors dump their data into an unchunked file.
   @enddesc

   @calls      IOHDF5Util_DumpCommonAttributes

   @var        GH
   @vdesc      pointer to CCTK grid hierarchy
   @vtype      const cGH *
   @vio        in
   @endvar
   @var        slab
   @vdesc      reference to the I/O hyperslab description
   @vtype      const ioSlab *
   @vio        in
   @endvar
   @var        name
   @vdesc      name of the dataset to write
   @vtype      const char *
   @vio        in
   @endvar
   @var        data
   @vdesc      pointer to the chunk to dump
   @vtype      const void *
   @vio        in
   @endvar
   @var        file
   @vdesc      the HDF5 file handle
   @vtype      hid_t
   @vio        in
   @endvar
@@*/
static void WriteDataCollective (const cGH *GH, const ioSlab *slab,
                                 const char *name, const void *data, hid_t file)
{
  int i, dim;
  hid_t hdf5type, dataset, memspace, filespace, plist;
  hssize_t *chunk_origin;
  hsize_t *chunk_dims, *file_dims;
  hsize_t buffersize;
  DECLARE_CCTK_PARAMETERS


  /* immediately return if file handle is invalid */
  if (file < 0)
  {
    return;
  }

  /* copy the size arrays from CCTK_INT to appropriate types
     note that HDF5 wants elements in reverse order */
  chunk_origin = (hssize_t *) malloc (slab->hdim * sizeof (hssize_t));
  chunk_dims   = (hsize_t *)  malloc (2*slab->hdim * sizeof (hsize_t));
  file_dims    = chunk_dims + slab->hdim;
  for (i = 0; i < slab->hdim; i++)
  {
    chunk_origin[i] = slab->hoffset[slab->hdim - 1 - i];
    chunk_dims  [i] = slab->hsize_chunk[slab->hdim - 1 - i];
    file_dims   [i] = slab->hsize_global[slab->hdim - 1 - i];
  }

  /* create the memspace according to chunk dims */
  HDF5_ERROR (memspace = H5Screate_simple (slab->hdim, chunk_dims, NULL));

  /* create the (global) filespace and set the hyperslab for the chunk */
  HDF5_ERROR (filespace = H5Screate_simple (slab->hdim, file_dims, NULL));
  HDF5_ERROR (H5Sselect_hyperslab (filespace, H5S_SELECT_SET, chunk_origin,
                                   NULL, chunk_dims, NULL));

  /* the I/O processor creates the dataset and adds the common attributes
     when writing its own data, otherwise the dataset is reopened */
  hdf5type = IOHDF5Util_DataType (myGH, slab->hdatatype);

  /* enable compression for chunked dataset if compression was requested */
  HDF5_ERROR (plist = H5Pcreate (H5P_DATASET_CREATE));
  if (compression_level)
  {
    HDF5_ERROR (H5Pset_chunk (plist, slab->hdim, chunk_dims));
    HDF5_ERROR (H5Pset_deflate (plist, compression_level));
  }
  HDF5_ERROR (dataset = H5Dcreate (file, name, hdf5type, filespace, plist));
  HDF5_ERROR (H5Pclose (plist));
  if (CCTK_MyProc (GH) == 0)
  {
    IOHDF5Util_DumpCommonAttributes (GH, slab, dataset);
  }

  /* increase the buffer size if the default isn't sufficient */
  HDF5_ERROR (plist = H5Pcreate (H5P_DATASET_XFER));
  buffersize = H5Dget_storage_size (dataset);
  if (buffersize > H5Pget_buffer (plist, NULL, NULL))
  {
    HDF5_ERROR (H5Pset_buffer (plist, buffersize, NULL, NULL));
  }

  /* write the data */
  HDF5_ERROR (H5Dwrite (dataset, hdf5type, memspace, filespace, plist, data));

  /* close resources */
  HDF5_ERROR (H5Pclose (plist));
  HDF5_ERROR (H5Sclose (filespace));
  HDF5_ERROR (H5Dclose (dataset));
  HDF5_ERROR (H5Sclose (memspace));

  /* free allocated resources */
  free (chunk_origin);
  free (chunk_dims);
}
#endif /* CCTK_MPI && H5_HAVE_PARALLEL */