"""
A grammar taking rows from a FITS table.
"""

#c Copyright 2008-2025, 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 itertools

from gavo import base
from gavo.grammars import common

from astropy import table


class FITSTableIterator(common.RowIterator):
	"""The row iterator for FITSTableGrammars.
	"""
	def _iterRows(self):
		data = table.Table.read(
			self.sourceToken,
			format="fits",
			hdu=self.grammar.hdu)
		names = [n for n in data.dtype.names]

		if self.grammar.lowerKeys:
			names = [n.lower() for n in names]

		mask = data.mask
		if mask is None:
			mask = itertools.cycle([[False]*len(names)])

		for raw_row, mask in zip(data, mask):
			row = {}
			for i in range(len(raw_row)):
				val = raw_row[i]
				if mask[i]:
					row[names[i]] = None
				elif self.grammar.nanIsNULL and val!=val:
					row[names[i]] = None
				else:
					row[names[i]] = val

			yield row


class FITSTableGrammar(common.Grammar):
	"""A grammar parsing from FITS tables.

	fitsTableGrammar result in typed records, i.e., values normally come
	in the types they are supposed to have.  Of course, that won't work
	for datetimes, STC-S regions, and the like.

	The keys of the result dictionaries are simpily the names given in
	the FITS.
	"""
	name_ = "fitsTableGrammar"

	_hduIndex = base.IntAttribute("hdu", default=1,
		description="Take the data from this extension (primary=0)."
		" Tabular data typically resides in the first extension.")
	_lowerKeys = base.BooleanAttribute("lowerKeys", default=False,
		description="Convert column names to lower case when building"
		" the rawdict.")
	_mapNaN = base.BooleanAttribute("nanIsNULL", default=False,
		description="Map NaN floats to NULL when building the rawdict.")

	rowIterator = FITSTableIterator
