"""
Writing data in FITS binary tables
"""

#c Copyright 2008-2024, 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 os
import tempfile
import time

import numpy

from gavo import base
from gavo import rsc
from gavo import utils
from gavo.formats import common
from gavo.utils import pyfits
from gavo.votable import common as votcommon


_fitsCodeMap = {
	"short": "I",
	"int": "J",
	"long": "K",
	"float": "E",
	"double": "D",
	"boolean": "L",
	"char": "A",
	"unicodeChar": "A",
}

_typeConstructors = {
	"short": int,
	"int": int,
	"long": int,
	"float": float,
	"double": float,
	"boolean": int,
	"char": str,
	"unicodeChar": str,
}


def _makeCharArray(values, colInd, colDesc):
	"""returns a pyfits-capable column array for strings stored in the
	colInd-th column of values.

	This is for char-valued columns, i.e., ASCII-only.
	"""
	arr = numpy.array([(v[colInd] or "").encode("ascii", "replace")
		for v in values],
		dtype=numpy.bytes_)
	return "%dA"%arr.itemsize, arr


def _makeStringArray(values, colInd, colDesc):
	"""returns a pyfits-capable column array for strings stored in the
	colInd-th column of values.

	This is for unicodeChar-valued columns; we'd like to do something
	sensible with them; this implementation is FIXME, because we're
	essentially discaring all non-ASCII and write a normal char column.
	"""
# Yeah, right now this is identical to _makeCharArray, but that's FIXME.
	arr = numpy.array([(v[colInd] or "").encode("ascii", "replace")
		for v in values],
		dtype=numpy.bytes_)
	return "%dA"%arr.itemsize, arr


def _getNullValue(colDesc):
	"""returns a null value we consider ok for a column described by colDesc.

	This is mainly intended for array serialisation.  It is also
	used for column construction for integer types.
	"""
	length = votcommon.getLength(colDesc["arraysize"])
	if length and length>1:
		# TODO: figure out nullvalue interpretation in FITS arrays.
		return None

	nullValue = colDesc["nullvalue"]
	if nullValue is None:
		# enter some reasonable defaults
		if (colDesc["datatype"]=="float"
			or colDesc["datatype"]=="double"):
			nullValue = float("NaN")

		elif colDesc["datatype"]=="text":
			nullValue = ""

	else:
		nullValue = _typeConstructors[colDesc["datatype"]](nullValue)
	
	return nullValue


def _computeTypeCode(colDesc):
	"""returns a FITS type code suitable for colDesc.
	"""
	baseCode = _fitsCodeMap[colDesc["datatype"]]
	length = votcommon.getLength(colDesc["arraysize"])

	if length==1:
		return length, baseCode

	elif length is None:
		# a variable-length array (other than a string).  We, in particular,
		# need this for polygons.
		return length, "P%s()"%(baseCode)

	else:
		return length, "%d%s"%(length, baseCode)


def _makeValueArray(values, colInd, colDesc):
	"""returns a pyfits-capable column array for non-string values
	stored in the colInd-th column of values.
	"""
	length, typecode = _computeTypeCode(colDesc)
	nan = float("NaN")

	if length is None:
		# That's a variable-length array that we'll fill from a
		# numpy object array
		return typecode, numpy.array([list(v[colInd]) for v in values],
			dtype=numpy.object)

	elif length>1:
		# numpy 1:1.12.1-3 and pyfits from astropy 1.3-8 apparently
		# have trouble with masked arrays here, so we exploit
		# that in the VO and for floats NaN and NULL is treated equivalently.
		nullArray = numpy.array([nan]*length)
		def mkval(v):
			if v is None:
				return nullArray
			else:
				return numpy.array(v)

	else:
		def mkval(v, nullValue=_getNullValue(colDesc)):
			if v is None:
				if nullValue is None:
					raise ValueError("While serializing a FITS table: NULL"
						" detected in column '%s' but no null value declared"%
						colDesc["name"])
				return nullValue
			else:
				return v

	arr = numpy.array([mkval(v[colInd]) for v in values])
	return typecode, arr


def _makeExtension(serMan):
	"""returns a pyfits hdu for the valuemappers.SerManager instance table.
	"""
	values = list(serMan.getMappedTuples())
	columns = []
	utypes = []
	descriptions = []

	for colInd, colDesc in enumerate(serMan):
		descriptions.append(utils.defuseFileName(colDesc["description"], False))
		if colDesc["datatype"]=="char":
			makeArray = _makeCharArray
		elif colDesc["datatype"]=="unicodeChar":
			makeArray = _makeStringArray
		else:
			makeArray = _makeValueArray
		
		typecode, arr = makeArray(values, colInd, colDesc)
		# only pass on null values to integral types; FITS doesn't support
		# anything else, and I'll just leave Nones in non-integral columns
		# to astropy.io.fits -- with all the mess involved, in particular
		# that \0 and NULL or NaN and NULL can't be told apart in FITS.
		nullValue = None
		if typecode in 'BIJK':
			nullValue = _getNullValue(colDesc)
		
		columns.append(pyfits.Column(name=str(colDesc["name"]),
			unit=str(colDesc["unit"]), format=typecode,
			null=nullValue, array=arr))
		if colDesc["utype"]:
			utypes.append((colInd, str(colDesc["utype"].lower())))

	hdu = pyfits.BinTableHDU.from_columns(
		pyfits.ColDefs(columns))
	for colInd, utype in utypes:
		hdu.header.set("TUTYP%d"%(colInd+1), utype)
	
	cards = hdu.header.cards
	for colInd, desc in enumerate(descriptions):
		cards["TTYPE%d"%(colInd+1)].comment = desc
		hdu.header.set("TCOMM%d"%(colInd+1), desc)

	if not hasattr(serMan.table, "IgnoreTableParams"):
		for param in serMan.table.iterParams():
			if param.value is None:
				continue
		
			key, value, comment = str(param.name), param.value, param.description
			if isinstance(value, str):
				value = utils.defuseFileName(value, False)
			if isinstance(comment, str):
				comment = utils.defuseFileName(comment, False)
			if len(key)>8:
				key = "hierarch "+key

			try:
				hdu.header.set(key, value=value, comment=comment)
			except ValueError as ex:
				# do not fail just because some header couldn't be serialised
				base.ui.notifyWarning(
					"Failed to serialise param %s to a FITS header (%s)"%(
						param.name,
						utils.safe_str(ex)))

	return hdu
	

def _makeFITSTableNOLOCK(dataSet, acquireSamples=True):
	"""returns a hdulist containing extensions for the tables in dataSet.

	You must make sure that this function is only executed once
	since pyfits is not thread-safe.
	"""
	tables = [base.SerManager(table, acquireSamples=acquireSamples)
		for table in list(dataSet.tables.values())]
	extensions = [_makeExtension(table) for table in tables]
	primary = pyfits.PrimaryHDU()
	primary.header.set("DATE", time.strftime("%Y-%m-%d"),
		"Date file was written")
	return pyfits.HDUList([primary]+extensions)


def makeFITSTable(dataSet, acquireSamples=False):
	"""returns a hdulist containing extensions for the tables in dataSet.

	This function may block basically forever.  Never call this from
	the main server, always use threads or separate processes (until
	pyfits is fixed to be thread-safe).

	This will add table parameters as header cards on the resulting FITS
	header.
	"""
	with utils.fitsLock():
		return _makeFITSTableNOLOCK(dataSet, acquireSamples)


def writeFITSTableFile(hdulist):
	"""returns the name of a temporary file containing the FITS data for
	hdulist.
	"""
	handle, pathname = tempfile.mkstemp(".fits", prefix="fitstable",
		dir=base.getConfig("tempDir"))

	# if there's more than the primary HDU, EXTEND=True is mandatory; let's
	# be defensive here
	if len(hdulist)>1:
		hdulist[0].header.set("EXTEND", True, "More exts following")

	with utils.silence():
		hdulist.writeto(pathname, overwrite=True)
	os.close(handle)
	return pathname


def makeFITSTableFile(dataSet, acquireSamples=True):
	"""returns the name of a temporary file containing a fits file
	representing dataSet.

	The caller is responsible to remove the file.
	"""
	hdulist = makeFITSTable(dataSet, acquireSamples)
	return writeFITSTableFile(hdulist)


def writeDataAsFITS(data, outputFile, acquireSamples=False):
	"""a formats.common compliant data writer.

	This will write out table params as header cards.  To serialise
	those yourself (as is required for spectral data model compliant
	tables), set an attribute IgnoreTableParams (with an arbitrary
	value) on the table.
	"""
	data = rsc.wrapTable(data)
	fitsName = makeFITSTableFile(data, acquireSamples)
	try:
		src = open(fitsName, "rb")
		utils.cat(src, outputFile)
		src.close()
	finally:
		os.unlink(fitsName)

common.registerDataWriter("fits", writeDataAsFITS, "application/fits",
	"FITS Binary Table", ".fits")
