"""
Helpers for generating boosters for HDF5 data.

HDF5 is fairly complex, and directgrammar is too long as is.  Also, I don't
want to require h5py as a fixed dependency of DaCHS; if it's not
there, you should still be able to use other sorts of direct grammars.

TODO: Most of the hdf interface functions return 0 on success, != on
failure; we'd like a macro to catch these.
"""

#c Copyright 2008-2022, the GAVO project <gavo@ari.uni-heidelberg.de>
#c
#c This program is free software, covered by the GNU GPL.  See the
#c COPYING file in the source distribution.


import h5py

from gavo import base
from gavo import utils

# See directgrammar.HDF5vaexCodeGenerator for how this import works
from gavo.grammars import directgrammar


class HDF5vaexCodeGenerator(directgrammar._NumpyMetaCodeGenerator):
	"""A code generator for boosters importing HDF5 files in VAEX convention.

	These have one array per column in a "columns" group; the actual data
	is in a "data" group.

	Our strategy when parsing from them is to read CHUNK_SIZE items
	at a time into the corresponding arrays, and then iterating
	over these chunks, building the records.
	"""
	_h5typemap = {
		"byte": "H5T_NATIVE_CHAR",
		"int8": "H5T_NATIVE_CHAR",
		"ubyte": "H5T_NATIVE_UCHAR",
		"uint8": "H5T_NATIVE_UCHAR",
    "short": "H5T_NATIVE_SHORT",
		"int16": "H5T_NATIVE_SHORT",
    "ushort": "H5T_NATIVE_USHORT",
    "uint16": "H5T_NATIVE_USHORT",
		"cint": "H5T_NATIVE_INT", 
		"int32": "H5T_NATIVE_INT",
		"uint": "H5T_NATIVE_UINT", 
		"uintc": "H5T_NATIVE_UINT", 
		"int64": "H5T_NATIVE_LLONG", 
		"uint64": "H5T_NATIVE_ULLONG", 
		"ulonglong": "H5T_NATIVE_ULLONG",
		"single": "H5T_NATIVE_FLOAT",
		"float32": "H5T_NATIVE_FLOAT",
		"double": "H5T_NATIVE_DOUBLE",
		"float64": "H5T_NATIVE_DOUBLE",
	}

	def __init__(self, grammar, tableDef):
		directgrammar._NumpyMetaCodeGenerator.__init__(self, grammar, tableDef)

		if self.grammar.parent.sources is None:
			raise base.StructureError("Cannot make HDF5 vaex booster without"
				" a sources element on the embedding data.")
		try:
			sampleSource = next(self.grammar.parent.sources.iterSources())
			hdf = h5py.File(sampleSource)
		except StopIteration:
			raise base.StructureError("Building an HDF5 booster requires"
				" at least one matching source.")

		try:
			self.dsName = self.grammar.getProperty("dataset", "table")
			cols = hdf[self.dsName]["columns"]
		except KeyError:
			raise base.StructureError(f"Cannot access dataset {self.dsName} in"
				f" {sampleSource}.  Override the grammar's dataset property"
				" to point it to the right dataset.")

		self.inputColumns = []
		for name in cols:
			numpyType = cols[name]["data"].dtype
			if numpyType.ndim!=0:
				raise NotImplementedError("Cannot handle arrays in vaex"
					" HDF5 yet.  Please contact the authors.")
			try:
				cType = self.numpyTypes[str(numpyType)]
			except KeyError:
				raise NotImplementedError("Cannot handle the numpy type"
					f" {numpyType} yet.  Please contact the authors.")
			self.inputColumns.append((name, numpyType, cType))
		hdf.close()

		nameMap = {}
		if self.grammar.mapKeys:
			nameMap = self.grammar.mapKeys.maps

		self.byName = utils.CaseSemisensitiveDict(
			(nameMap.get(t[0], t[0]), t)
			for t in self.inputColumns)

		
	def _getStructDefinition(self):
		"""returns a definition for a C structure for holding one record.
		"""
		res = ["typedef struct InRec_s {"]
		for name, _, cType in self.inputColumns:
			res.append(f"  {cType} {name};")
		res.append("} InRec;")
		return "\n".join(res)

	def _getArrayDecls(self):
		"""returns the declarations for the buffer arrays.
		"""
		return "\n".join(
			f"  {cType} arr_{name}[CHUNK_SIZE];"
			for name, _, cType in self.inputColumns)

	def _getRefillCode(self):
		"""returns code to (up to) CHUNK_SIZE items into the column
		arrays.
		"""
		res = [	
			"offsets[0] = total_read;",
			"if (total_read+chunk_size[0]>nrecs) {",
			"	chunk_size[0] = nrecs-total_read;",
			"	HDFGUARD(H5Sselect_hyperslab(memspace, H5S_SELECT_SET,"
			"		null_offset, NULL, chunk_size, NULL));",
			"}",
			"HDFGUARD(H5Sselect_hyperslab(colspace, H5S_SELECT_SET,"
			"  offsets, NULL, chunk_size, NULL));",]

		for index, (name, npType, cType) in enumerate(self.inputColumns):
			memTypeId = self._h5typemap[str(npType)]
			res.append(
				f"HDFGUARD(H5Dread(datasets[{index}], {memTypeId}, memspace,"
				f" colspace, H5P_DEFAULT, arr_{name}));")
		
		return "\n".join(res)

	def _getDatasetsCode(self):
		"""returns code to create the datasets for the columns; the
		result is kept in a datasets array.
		"""
		res = []
		for index, (name, _, cType) in enumerate(self.inputColumns):
			col_ds_name = "/".join([self.dsName, "columns", name, "data"])
			res.extend([
				f"  datasets[{index}] = H5Dopen2("
					f'input_id, "{col_ds_name}", H5P_DEFAULT);',
				f"  if (datasets[{index}]<0) abort();"])
		return "\n".join(res)
	
	def _getTupleBuilder(self):
		"""returns code to build the InRec struct passed to getTuple.
		"""
		res = []
		for index, (name, _, cType) in enumerate(self.inputColumns):
			res.append(f"cur_rec.{name} = arr_{name}[index_in_chunk];");
		return "\n".join(res)

	def getPreamble(self):
		n_cols = len(self.inputColumns)

		return directgrammar._NumpyMetaCodeGenerator.getPreamble(self
			)+[
			"#include <inttypes.h>",
			"#include <stdlib.h>",
			"#include <hdf5.h>",
			f"#define N_COLS {n_cols}",
			# how many rows to process in one go?
			"#define CHUNK_SIZE {}".format(
				self.grammar.getProperty("chunkSize", 5000)),  
			self._getStructDefinition(),
			"#define HDFGUARD(x) if ((x)<0) abort();",]

	def getPrototype(self):
		return "Field *getTuple(InRec *data, int rowIndex)"
	
	def getFooter(self):
		return "\n".join([
			directgrammar.COMMON_MAIN_HEADER,
			"	hid_t input_id;"
			"	hid_t datasets[N_COLS];",
			"	hid_t colspace;",
			"	size_t nrecs;",
			"	int rank;",
			"	hsize_t dims[1], offsets[1];",
			"	hsize_t null_offset[1] = {0};",
			"	hsize_t chunk_size[1] = {CHUNK_SIZE};",
			"	hid_t memspace = H5Screate_simple(1, chunk_size, NULL);",
			# index_in_chunk set so we immediately read a chunk at startup
			"	size_t index_in_chunk = CHUNK_SIZE;",
			"	size_t total_read = 0;",
			"	InRec cur_rec;",
			"",
			directgrammar.COMMON_MAIN_INTRO,
			"	input_id = H5Fopen(argv[1], H5F_ACC_RDONLY, H5P_DEFAULT);",
			"	if (input_id<0) abort();",
			"	HDFGUARD(H5Sselect_hyperslab(memspace, H5S_SELECT_SET,"
			" 		null_offset, NULL, chunk_size, NULL));",
			self._getArrayDecls(),
			self._getDatasetsCode(),
			'	colspace = H5Dget_space(datasets[0]);',
			'	rank = H5Sget_simple_extent_ndims(colspace);',
			'	if (rank>1) abort();',
			'	H5Sget_simple_extent_dims(colspace, dims, NULL);',
			'	nrecs = dims[0];',

			"	while (total_read<nrecs) {",
			directgrammar.LOOP_BODY_INTRO,
			"			if (index_in_chunk>=CHUNK_SIZE) { ",
			utils.fixIndentation(self._getRefillCode(), "				"),
			"			index_in_chunk = 0;",
			"		}",
			utils.fixIndentation(self._getTupleBuilder(), "		"),
			"		total_read += 1;",
			"		index_in_chunk += 1;",
			"		tuple = getTuple(&cur_rec, total_read-1);",
			directgrammar.LOOP_BODY_FOOT,
			"	}",
			directgrammar.COMMON_MAIN_FOOT,])

	def getItemParser(self, item, index):
		nameForItem = directgrammar.getNameForItem(item)
		typeMacro = directgrammar._getMakeMacro(item)
		res = [
			f"/* {item.description} ({item.type}) */",]
		if item.name not in self.byName:
			res.append(
				f"MAKE_NULL({nameForItem});"
				f"/* {typeMacro}({nameForItem}, FILL IN VALUE); */")

		else:
			origName, _, castTo = self.byName[item.name]
			# there is a source column in the HDF5, make a default map
			res.append(
				f"{typeMacro}({nameForItem}, (({castTo})(data->{origName})));")

		return res
			
