"""
A custom grammar for the fits/dat combination in Robert Schmidt's infinite
light curve.

The lightcurve is assembled from three seperate items:

* the raw amplification patterns (MAGPAT*.gz), which are just dumps of
  FORTRAN arrays.  They are used to build images.
* An offsets file containing horizontal offsets between consecutive MAGPATS
* the lightcurves -- these also give translations from local indices into
  MAGPATS to global line numbers.

The complexity with the various indices is because the border regions of the
patterns are bogus.  Therefore, they have been computed with overlaps, and
the local indices in the lightcurves address the valid parts (plus,
horizontally, the offsets).

The first two items are loaded into the datapack since the data structures
are inited for all magpats at once.
"""

import base64
import glob
import gzip
import itertools
import os
import re
import subprocess

import numpy


from gavo import api
from gavo.grammars.customgrammar import CustomRowIterator


def _getPixarray(fName, imageDimension):
	if fName.endswith(".gz"):
		# numpy.fromfile actually uses the C file API, so I have to bite the
		# bullet and unzip stuff manually.  I'll be reckless here, as this
		# code will essentially never run.
		subprocess.check_call(["gzip", "-d", fName])
		fName = fName[:-3]

	f = open(fName, "rb")
	f.read(4)  # Skip weird header
	arr = numpy.fromfile(f,
		dtype=("uint16", (imageDimension,imageDimension)))
	f.close()
	subprocess.check_call(["gzip", fName])
	return arr[0]


def _iterPixelLines(values, lowerCutoff, upperCutoff):
	"""yields imageData lines from values
	"""
	normalizer = (upperCutoff-lowerCutoff)/256.
	pixels = (numpy.floor(numpy.clip(values, lowerCutoff, upperCutoff-1)-
		lowerCutoff)/normalizer).astype("uint8")
	for ind, line in enumerate(pixels):
		yield line.tostring()


def _readOffsetsFile(sourcePath, firstMagpat):
	offsets = {}
	curOffset = 0
	for magpatInd, ln in enumerate(open(sourcePath)):
		curOffset -= int(ln.split()[3])
		offsets[firstMagpat+magpatInd] = curOffset
	return offsets


def _getLightcurvePats(grammar):
	lightcurvePats = []
	for field in grammar.rd.getById("data"):
		try:
			srcPath = field.getProperty("srcPath")
		except api.NotFoundError:  # not a field with srcPath, has no light curve
			continue
		lightcurvePats.append((field.name,
				os.path.join(grammar.rd.resdir, srcPath)))
	return lightcurvePats


class DataPack(object):
	def __init__(self, offsets, lightcurvePats, imageDimension):
		self.offsets, self.lightcurvePats = offsets, lightcurvePats
		self.imageDimension = imageDimension


def makeDataPack(grammar):
	try:
		offsetPath = os.path.join(grammar.rd.resdir,
			grammar.getProperty("offsetsFile"))
		imageDimension = int(grammar.getProperty("imageDimension"))
	except KeyError:
		raise api.StructError("Your custom grammar needs at least"
			" offsetsFile and imageDimension properties")
	return DataPack(_readOffsetsFile(offsetPath, 1), _getLightcurvePats(
		grammar), imageDimension)


class RowIterator(CustomRowIterator):
	def _readLightcurve(self, fName):
		f = open(fName)
		vals = []
		for ln in f:
			if not ln.strip():
				break
			s = ln.split()
			vals.append((int(s[0]), int(s[1]), float(s[2])))
		f.close()
		return vals

	def _readLightcurvesFor(self, sourceNumber):
		"""reads all lightcuves for source number and returns a dict mapping
		local row numbers to dictionaries of lightcurve values.
		"""
		lcValues = {}
		lciter = iter(self.grammar.dataPack.lightcurvePats)
		destName, srcPat = next(lciter)
		thisCurve = self._readLightcurve(srcPat%sourceNumber)
		startGlobal, startLine = thisCurve[0][0], thisCurve[0][1]
		for globalInd, localInd, value in thisCurve:
			lcValues[localInd] = (globalInd, {destName: value})

		for destName, srcPat in lciter:
			thisCurve = self._readLightcurve(srcPat%sourceNumber)
			if (startGlobal, startLine)!=(thisCurve[0][0], thisCurve[0][1]):
				raise Error("At source number %d: start lines don't agree"
					" for %s"%(sourceNumber, destName))
			for globalInd, localInd, value in thisCurve:
				lcValues[localInd][1][destName] = value

		return lcValues

	def _getSourceNumber(self, sourceToken):
		return int(re.search("\d\d\d\d", sourceToken).group())

	def _iterRows(self):
		sourceNumber = self._getSourceNumber(self.sourceToken)
		lightcurves = self._readLightcurvesFor(sourceNumber)
		upperCutoff = int(self.grammar.parent.getProperty("upperCutoff"))
		lowerCutoff = int(self.grammar.parent.getProperty("lowerCutoff"))

		# ROI is horizontal; vertical ROI is by not having localInds in
		# lightcurves
		roiWidth = int(self.grammar.parent.getProperty("roiWidth"))
		roiStart = (self.grammar.dataPack.imageDimension//2-roiWidth//2-
			self.grammar.dataPack.offsets[sourceNumber])
		roiEnd = roiStart+roiWidth
		arr = _getPixarray(self.sourceToken, self.grammar.dataPack.imageDimension)
		roi = arr[roiStart:roiEnd].transpose()

		for localInd, ln in enumerate(
				_iterPixelLines(roi, lowerCutoff, upperCutoff)):

			if localInd in lightcurves:
				globalInd, mags = lightcurves[localInd]
				res = {
					"line": globalInd,
					"data": base64.b64encode(ln),
					"srcname": os.path.basename(self.sourceToken)}
				res.update(mags)
				yield res