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

These have an incredibly verbose header, plus four extensions.  See Paper,
section 4 for details.

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.

We check badpix; if all of them are 1 of a given spectrum,
it's "empty".  Quoting Bernd Husemann:

	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 rscdef
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-internal obsid is (base accref)-(x-index)-(y-index);
		# (base accref) is the standard accref minus the extension.
		
		idPattern = "%s-%%02d-%%02d"%(rscdef.getInputsRelativePath(
			self.sourceToken)[:-12])

		# 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
					pixelOffset = 0
				elif setup=="V1200":
					snrWL = 4000
					pixelOffset = 1000
				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])
				thisMeta["xindex"] = raInd+pixelOffset
				thisMeta["yindex"] = decInd+pixelOffset

				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"))