"""
A grammar to parse split-order echelle spectra written by two versions of
MIDAS.

The first was for early FLASH spectra.

In them, only CDELT1 is usable from the normal WCS.  The rest is
in ASCII in the HISTORY cards.  There, for each order (i.e., for each
in NAXIS2), the start wavelength (in Angstrom) is given.  NPTOT gives
the number of valid pixels in that order.  NORDER is a running number of
questionable meaning.

The grammar yields the FITS header via getParameters, tuples of
order, wavelength (Angstrom), and uncalibrated flux as rowdicts.

The other format was employed later in the FLASH life cycle and by HEROS.  In
it, there are ECHFIT descriptors which give the coefficients of a polynomial
yielding the wavelengths as

{[a(1)+a(2)*m] + [a(3)+a(4)*m]*x + a(5)*x**2 + a(6)*x**3} / (1+a(7)*m

where x = 1 .. NAXIS1 and m = 1 .. NAXIS2.

That format also has DTIME, which is an array with a funny representation
of the observeration date.  Use the fourth element of that array,
which happens to be the MJD.

See also the README.



The grammar here is also funky in that it lets you pass either a string or
a dict as sourceToken.  The second form expects a dict containing
{ 'input_path':
	'order_min':
	'order_max':
} -- this lets the datalink function preselect what orders it wants
to see.  min and max may be None for open ranges.
"""

import re

from gavo import utils
from gavo.utils import pyfits
from gavo.grammars.customgrammar import CustomRowIterator


class RowIterator(CustomRowIterator):
	def __init__(self, grammar, sourceToken, *args, **kwargs):
		self.orderMin = self.orderMax = None
		if isinstance(sourceToken, dict):
			self.orderMin = sourceToken["order_min"]
			self.orderMax = sourceToken["order_max"]
			sourceToken = sourceToken["input_path"]

		CustomRowIterator.__init__(self, grammar, sourceToken, *args, **kwargs)
		hdulist = pyfits.open(self.sourceToken)
		self.header = hdulist[0].header
		self.data = hdulist[0].data[:]
		hdulist.close()
		self.esoDescriptors = utils.parseESODescriptors(self.header)

	def getParameters(self):
		return dict(iter(self.header.items()))

	def _iterOrdersFLASH(self):
		for index, (order, nptot, lambda0) in enumerate(
				zip(self.esoDescriptors["NORDER"],
					self.esoDescriptors["NPTOT"],
					self.esoDescriptors["WSTART"])):
			yield index, order, nptot, lambda0

	def _filterOrdersMetaKeys(self, aDict):
		return dict((k, aDict[k])
			for k in "n_orders order_min order_max total_length"
				" specMin specMax".split())

	def _getHEROSOrder(self, orderIndex):
		"""returns the Echelle order for the array index of an order.
		"""
		first, last = self.esoDescriptors["ORDER"]
		delta = (first-last)//abs(first-last)
		return first-orderIndex*delta

	# The following methods are used by the import_ordersmeta data
	# item.
	def _getOrdersMetaFLASH(self):
		spectBinSize = self.header["CDELT1"]
		n_orders, specMin, specMax = 0, utils.Supremum, utils.Infimum
		order_min, order_max = utils.Supremum, utils.Infimum
		total_length = 0

		for index, order, nptot, lambda0 in self._iterOrdersFLASH():
			n_orders += 1
			total_length += nptot
			specMin = min(specMin, lambda0)
			specMax = max(specMax, lambda0+index*spectBinSize)
			order_min = min(order_min, order)
			order_max = max(order_max, order)

		return self._filterOrdersMetaKeys(locals())
	
	def _getOrdersMetaHEROS(self):
		trafo = self._makeHerosTrafo(self.esoDescriptors["ECHFIT"])
		n_orders, specMin, specMax = 0, utils.Supremum, utils.Infimum
		order_min, order_max = utils.Supremum, utils.Infimum

		for orderIndex in range(self.header["NAXIS2"]):
			n_orders += 1
			specMin = min(specMin, trafo(1, orderIndex+1))
			specMax = max(specMax, trafo(self.header["NAXIS1"], orderIndex+1))
			order = self._getHEROSOrder(orderIndex)
			order_min = min(order_min, order)
			order_max = max(order_max, order)

		total_length = n_orders*self.header["NAXIS1"]
		return self._filterOrdersMetaKeys(locals())
		
	def getOrdersMeta(self):
		if "WSTART" in self.esoDescriptors:
			return self._getOrdersMetaFLASH()
		else:
			return self._getOrdersMetaHEROS()

	def _iterRowsFLASH(self):
		spectBinSize = self.header["CDELT1"]
		for index, order, nptot, lambda0 in self._iterOrdersFLASH():
			if self.orderMax is not None and order>self.orderMax:
				continue
			if self.orderMin is not None and order<self.orderMin:
				continue

			for spectralInd in range(nptot):
				yield (str(order), {
					"ecorder": order,
					"lambda": lambda0+spectralInd*spectBinSize,
					"flux": self.data[index, spectralInd]})

	def _makeHerosTrafo(self, coeffs):
		"""returns a transformation function between a HEROS order index
		and the wavelength.

		The second argument is the fits index of the order (i.e. 1..NAXIS2).
		"""
		# Otmar suspects that the last three values are unused
		a, b, c, d, e, f, g, _, _, _ = coeffs

		def wavelength(spectral, order):
			return (((a+b*order)
					+(c+d*order)*spectral
					+e*spectral**2
					+f*spectral**3)
       	/(1+g*order))

		return wavelength

	def _iterRowsHEROS(self):
		trafo = self._makeHerosTrafo(self.esoDescriptors["ECHFIT"])
		for orderIndex in range(self.header["NAXIS2"]):
			order = self._getHEROSOrder(orderIndex)
			if self.orderMax is not None and order>self.orderMax:
				continue
			if self.orderMin is not None and order<self.orderMin:
				continue

			for spectralIndex in range(self.header["NAXIS1"]):
				yield str(order), ({
					"ecorder": order,
					"lambda": trafo(spectralIndex+1, orderIndex+1),
					"flux": self.data[orderIndex, spectralIndex]})

	def _iterRows(self):
		if "WSTART" in self.esoDescriptors:
			return self._iterRowsFLASH()
		else:
			return self._iterRowsHEROS()


if __name__=="__main__":
	from gavo.grammars.customgrammar import CustomGrammar
	ri = RowIterator(CustomGrammar(None), "../data_raw/ls95/n-files/blue/n0044.mt")
	print(ri.getParameters())
	for row in ri:
		print(row)