@*
HDF5.tmpl

Created by Graham Dennis on 2008-03-28.

Copyright (c) 2008-2012, Graham Dennis

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 2 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.

*@
@extends xpdeint.ScriptElement

@from xpdeint.Geometry.SplitUniformDimensionRepresentation import SplitUniformDimensionRepresentation
@from xpdeint.CallOnceGuards import callOnceGuard
@from itertools import chain

@attr $uselib = ['hdf5']

@@callOnceGuard
@def includes
  @#
  @super
  @#
#define H5_USE_16_API
#include <hdf5.h>

#if !defined(HAVE_H5LEXISTS)
htri_t H5Lexists(hid_t loc_id, const char *name, hid_t lapl_id)
{
  H5E_auto_t error_func;
  void* error_client_data;
  // Squelch errors generated by H5Gget_objinfo. It will report errors when it can't find an object
  // but that's the purpose of calling it.
  H5Eget_auto(&error_func, &error_client_data);
  H5Eset_auto(NULL, NULL);
  herr_t err = H5Gget_objinfo(loc_id, name, false, NULL);
  H5Eset_auto(error_func, error_client_data);
  if (err >= 0)
    return true;
  else
    return false;
}
#endif

#define H5T_NATIVE_REAL H5T_NATIVE_${{'double': 'DOUBLE', 'single': 'FLOAT'}[$precision]}
  @#
@end def

@def processData(dict)
  @set $field = dict['field']
  @set $basis = dict['basis']
  @set $operation = dict['operation']
  @assert operation in ['read', 'write']
  @set $variables = dict['variables']
  @set $dimensionOffsets = dict.get('dimensionOffsets', {})
  @set $componentCount = 0
  @for $variable in $variables
    @set $componentCount += len($variable.vector.components)
    @if $variable.vector.type == 'complex'
      @set $componentCount += len($variable.vector.components)
    @end if
  @end for
  @set dict['componentCount'] = componentCount
/* Create the data space */
  @set $dimensionCount = len(field.dimensions)
  @# File dim reps must be in dimension order as that is the order we desire for write-out
  @set $fileDimReps = [dim.inBasis(basis) for dim in field.dimensions]
  @set $memDimReps = field.inBasis(basis)
  @# Construct a list of (fileDimNum, memDimNum, dimRep) tuples. This is necessary
  @# for the case where we are using a distributed MPI driver with FFT's
  @# and the first two dimensions are transformed. In this situation, the
  @# first and second dimensions are transposed.
  @#
  @set dimRepOrdering = [(fileDimNum, memDimReps.index(dimRep), dimRep) \
                         for fileDimNum, dimRep in enumerate(fileDimReps)]
  @#
hsize_t file_start[${dimensionCount}] = {${', '.join(dimRep.localOffset for dimRep in fileDimReps)}};
hsize_t mem_dims[${dimensionCount+1}] = {${', '.join(chain((dimRep.localLattice for dimRep in memDimReps), ['1']))}};
hsize_t mem_start[${dimensionCount+1}] = {${', '.join(['0']*(dimensionCount+1))}};
hsize_t mem_stride[${dimensionCount+1}] = {${', '.join(['1']*(dimensionCount+1))}};
hsize_t mem_count[${dimensionCount+1}] = {${', '.join(chain((dimRep.localLattice for dimRep in memDimReps), ['1']))}};

  @for fileDimNum, memDimNum, dimRep in dimRepOrdering
    @if dimRep.name in dimensionOffsets and not isinstance(dimRep, SplitUniformDimensionRepresentation)
      @set offset = dimensionOffsets[dimRep.name]
if (file_start[${fileDimNum}] < ${offset}) {
  if (mem_count[${memDimNum}] < ${offset} - file_start[${fileDimNum}])
    mem_count[${memDimNum}] = 0;
  else {
    mem_count[${memDimNum}] -= ${offset} - file_start[${fileDimNum}];
    mem_start[${memDimNum}] += ${offset} - file_start[${fileDimNum}];
  }
  file_start[${fileDimNum}] = 0;
} else {
  file_start[${fileDimNum}] -= ${offset};
}

if (mem_count[$memDimNum] > file_dims[$fileDimNum]) {
  mem_count[$memDimNum] = file_dims[$fileDimNum];
}
    @end if
  @end for

hid_t mem_dataspace;
  @for variable in variables
    @set $components = $variable.separatedComponents
mem_dims[${dimensionCount}] = ${len(components)};
mem_dataspace = H5Screate_simple(${dimensionCount+1}, mem_dims, NULL);
mem_stride[${dimensionCount}] = ${len(components)};

    @capture writeLoopContents
      @for offset, componentName in components
mem_start[${dimensionCount}] = $offset;
H5Sselect_hyperslab(mem_dataspace, H5S_SELECT_SET, mem_start, mem_stride, mem_count, NULL);
        @# 
        @#  This looks like a typo because 'mem_stride' and 'mem_count' are used with 'file_start' are used here.
        @#  But it isn't a typo. The idea here is that the selection we want to make in the file has the same
        @#  number of elements in each dimension and the same stride as in memory (but ignoring the last dimension).
        @#  The only difference is the starting position for the selection.
        @#
        @if fileDimReps
          @# We can only do a selection in the file if the output data is more than zero-dimensional
H5Sselect_hyperslab(file_dataspace, H5S_SELECT_SET, file_start, mem_stride, mem_count, NULL);
        @end if

if (dataset_${componentName})
  H5D${operation}(dataset_${componentName}, H5T_NATIVE_REAL, mem_dataspace, file_dataspace, H5P_DEFAULT, ${variable.arrayName});
      @end for
    @end capture
    @#
    @# Permit the driver to modify the writeLoopContents
    @set $featureOrdering = ['Driver']
    @set $dict = {'writeLoopContents': writeLoopContents,
                  'dimRepOrdering': dimRepOrdering}
    @silent $insertCodeForFeatures('writeDataHDF5ModifyLoopContents', featureOrdering, dict)
    @set $writeLoopContents = dict['writeLoopContents']
    @#
    @# The object passed as the first argument to this next call is a 
// Select hyperslabs of memory and file data spaces for data transfer operation
${splitUniformDataSelect(dimRepOrdering[:], writeLoopContents, dimensionOffsets)}@slurp

H5Sclose(mem_dataspace);
  @end for
@end def

@def splitUniformDataSelect(remainingDimReps, writeLoopContents, dimensionOffsets)
  @if not remainingDimReps
${writeLoopContents}@slurp
  @else
    @set fileDimNum, memDimNum, dimRep = remainingDimReps.pop(0)
    @if isinstance(dimRep, SplitUniformDimensionRepresentation)
for (bool _positive_${dimRep.name} = true; ; _positive_${dimRep.name} = false) {
  if (_positive_${dimRep.name}) {
    mem_start[$memDimNum] = 0;
    if ($dimRep.localOffset >= (($dimRep.globalLattice-1)/2 +1)) // No positive values are stored in this rank.
      continue;
    mem_count[$memDimNum] = MIN((($dimRep.globalLattice-1)/2 +1) - $dimRep.localOffset, $dimRep.localLattice);
    file_start[$fileDimNum] = file_dims[$fileDimNum]/2 + $dimRep.localOffset;
      @if dimRep.name in dimensionOffsets
    if ($dimRep.localOffset > ((file_dims[$fileDimNum]-1)/2 +1))
      continue;
    mem_count[$memDimNum] = MIN(mem_count[$memDimNum], ((file_dims[$fileDimNum]-1)/2 +1) - $dimRep.localOffset);
      @end if
  } else {
    if (($dimRep.globalLattice-1)/2+1 < $dimRep.localOffset) // Only negative values are stored in this rank.
      mem_start[$memDimNum] = 0;
    else
      mem_start[$memDimNum] = ($dimRep.globalLattice-1)/2+1 - $dimRep.localOffset;
    file_start[$fileDimNum] = mem_start[$memDimNum] + $dimRep.localOffset - (($dimRep.globalLattice-1)/2+1);
    if ($dimRep.localLattice <= mem_start[$memDimNum]) // No negative values are stored in this rank.
      break; // end loop over this dimension
    mem_count[$memDimNum] = $dimRep.localLattice - mem_start[$memDimNum]; // To the end of this dimension
      @if dimRep.name in dimensionOffsets
    if ((long)file_start[$fileDimNum] > ${dimensionOffsets[dimRep.name]})
      file_start[$fileDimNum] -= ${dimensionOffsets[dimRep.name]};
    else {
      mem_start[$memDimNum] += ${dimensionOffsets[dimRep.name]} - file_start[$fileDimNum];
      if (mem_count[$memDimNum] > ${dimensionOffsets[dimRep.name]} - file_start[$fileDimNum])
        mem_count[$memDimNum] -= ${dimensionOffsets[dimRep.name]} - file_start[$fileDimNum];
      else
        break; // end loop over this dimension
      file_start[$fileDimNum] = 0;
    }
      @end if
  }
  
  ${splitUniformDataSelect(remainingDimReps, writeLoopContents, dimensionOffsets), autoIndent=True}@slurp
  
  if (!_positive_${dimRep.name})
    break;
}
    @else
${splitUniformDataSelect(remainingDimReps, writeLoopContents, dimensionOffsets)}@slurp
    @end if
  @end if
@end def
