"""
A grammar to read the data cubes from CALIFA.

These have an incredibly verbose header, plus three extensions.  From
a mail from Bernd Husemann 2012-04-11:

	The data cube has 4 FITS extensions:
	HDU 0: DATA  (observed spectra in flux units)
	HDU 1: ERROR (corresponding error spectra)
	HDU 2: ERRWEIGHT
	HDU 3: BADPIX (bad pixel mask 1: bad, 0: good)

	Please ignore the ERRWEIGHT for the time being.

This is a dispatched grammar; it *could* read the actual position/flux data
as well, but that's now done using res/booster, so it's suppressed here.

However, we need to check badpix; if all of them are 1 of a given spectrum,
it's "empty".  Again quoting Bernd:

	The actual field of view has a hexagonal shape, so that empty spectra are
	added to achieve a shape of a cube.
"""

import itertools
import math
import operator
import os

import pywcs

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

class RowIterator(CustomRowIterator):
	indices = 0,0

	def _getWCSMapper(self, refPix, refVal, cDelt):
		def mapper(val):
			return (val-refPix)*cDelt+refVal
		return mapper

	def _getWavelengthMapper(self, primaryHeader):
		return self._getWCSMapper(
			primaryHeader["CRPIX3"], primaryHeader["CRVAL3"],
			primaryHeader["CDELT3"])

	def _iterRows(self):
		hdus = pyfits.open(self.sourceToken)
		primary, err, _, badpix, fiberCov = hdus
		wlMapper = self._getWavelengthMapper(primary.header)
		posMapper = pywcs.WCS(primary.header, naxis=(1,2)).all_pix2sky
		baseName = os.path.basename(self.sourceToken)
		stem = baseName[:baseName.find(".")]
		if "V1200" in self.sourceToken:
			setup = "V1200"
		elif "V500" in self.sourceToken:
			setup = "V500"
		else:
			setup = "?SETUP?"

		baseMeta = {}
		for card in primary.header.ascard:
			try:
				baseMeta[card.key] = card.value
			except ValueError, msg:
				base.ui.notifyError(
					"Bad card in %s: %s"%(self.sourceToken, str(msg)))

		# the GAVO/CALIFA-internal obsid is (setup)/(object)-(x-index)-(y-index);
		# this definition is encoded both here and in res/booster.c.
		# (setup) is either V500 or V1200.
		idPattern = "%s/%s-%%02d-%%02d"%(setup, stem)

		# now, for each spectrum, make up an obsid and return the primary
		# keys plus the obsid to the metadata table, and then the spectrum
		# to the fluxes table
		# The array axes are (3/Spectrum, 2/Dec, 1/RA)
		specPoints, decPoints, raPoints = primary.data.shape

		for raInd in range(raPoints):
			for decInd in range(decPoints):
				self.indices = raInd, decInd

				if min(badpix.data[:, decInd, raInd])==1:
					# all pixels in this spectrum are bad, it's a filler, skip it
					continue

				curId = idPattern%(raInd, decInd)
				thisMeta = baseMeta.copy()
				thisMeta["obsId"] = curId
				thisMeta["setup"] = setup
				thisMeta["dstitle"] = "CALIFA %s %s/%s"%(thisMeta["OBJECT"],
					raInd, decInd)
				physPos = posMapper((raInd,), (decInd,), 0)
				thisMeta["RA"] = float(physPos[0][0])
				thisMeta["Dec"] = float(physPos[1][0])

				# the wavelength on which the SNR is computed depends on
				# whether we have a V500 or a V1200 spectrum; I can tell them
				# apart by what headers there are
				if setup=="V500":
					snrWL = 5000
				elif setup=="V1200":
					snrWL = 4000
				else:
					raise base.SourceParseError("Neither V500 nor V1200?",
						source=self.sourceToken)

				snrInd = int((snrWL-thisMeta["CRVAL3"])/thisMeta["CDELT3"]
					+thisMeta["CRPIX3"])
				thisMeta["computedSNR"] = (primary.data[snrInd, decInd, raInd]
					/err.data[snrInd, decInd, raInd])

				yield ("spectra", thisMeta)
				yield ("products", thisMeta)

######## Import of the actual data is now done via a booster.
#
#				dataIter = itertools.izip(primary.data[:, decInd, raInd],
#					err.data[:, decInd, raInd],
#					badpix.data[:, decInd, raInd])
#
#				for wlInd, (flux, fluxErr, bad) in enumerate(dataIter):
#					if bad:
#						yield ("fluxes", {
#							"obsId": curId,
#							"lambda": wlMapper(wlInd+1),
#							"flux": None,
#							"error": None})
#					else:
#						yield ("fluxes", {
#							"obsId": curId,
#							"lambda": wlMapper(wlInd+1),
#							"flux": flux,
#							"error": fluxErr,})

	def getLocator(self):
		return "ra/dec %s/%s"%self.indices


if __name__=="__main__":
	from gavo.user import logui
	from gavo import api
	logui.LoggingUI(api.ui)
	class StandinGrammar:
		rd = api.getRD("califa/q")
	g = StandinGrammar()

	r = RowIterator(g, "data/UGC10693.V1200.rscube.fits", None)
	for rec in r._iterRows():
		print dict((k,v) for k,v in rec[1].iteritems() if not k.startswith("PPAK") and not k.startswith("PIPE"))