# -*- coding: utf-8 -*-
"""
A core wrapping Corrado's space spectral matching code.
"""

# to support more input parameters:
# (1) define in the core's inputTableXML
# (2) in Core.run, say how to translate from input val to what's in the control
#     file
# (3) amend makeParFile
# (4) if necessary, add a sensible default in runSpace, progArgs


import os
import shutil
import subprocess
import tempfile

from gavo import api
from gavo import rsc
from gavo import svcs
from gavo import utils

# see README
DUMP_OUTPUT = False

SPACE_BINARIES = {
#	Binaries of old versions available; this maps the keys from
# version (in inputTableXML) to excutable names.  So, if
# you change things here, you'll need to change them down there, too.
	"paper": "SPACE-paper",
	"current": "SPACE",
	None: "SPACE",
}


class Core(svcs.Core):
	inputTableXML = """<inputTable>
			<inputKey name="spectrum" type="file" required="True"
				multiplicity="single"
				tablehead="Spectrum"
				description="ASCII file with two columns: wavelength (in Angstrom) and
					continuum normalized flux. The spectrum must be radial velocity
					corrected (wavelengths in rest frame). The spectral resolution power
					should be between 2000 and 20000. SP_Ace handles spectra in the
					stellar parameters intervals Teff=[3600,7400]K, logg=[0.2,5.0],
					[M/H]=[-2.4,0.4]dex.">
				<property name="adaptToRenderer">True</property>
			</inputKey>
			<inputKey name="fwhm" multiplicity="single" required="True"
				unit="Angstrom" ucd="spect.line.width;instr"
				tablehead="Instr. FWHM"
				description="Starting value for estimation of FWHM of the
					 instrument line profile."
				>0.4</inputKey>
			<inputKey name="wave_lims" type="text"
				multiplicity="single"
				tablehead = "Wave intervals"
				description = "Give up to five wavelength intervals you want to
					analyze, starting from the lowest.  Intervals not covered by the
					library will be ignored.  The default setting is the range of
					wavelenghts currently processed by the software."
					>4800 6860</inputKey>
			<inputKey name="tforce" multiplicity="single"
				ucd="stat.param;phys;temperature.effective" unit="K"
				tablehead="Fixed Teff"
				description="Force solver to assume this temperature.
					Leave empty to let SP_Ace estimate this parameter."/>
			<inputKey name="gforce" multiplicity="single"
				ucd="stat.param;phys.gravity"
				tablehead="Fixed gravity"
				description="Force  solver to assume this gravity.
					Leave empty to let SP_Ace estimate this parameter."/>
			<inputKey name="snrforce" multiplicity="single"
				ucd="stat.param;stat.snr"
				tablehead="Fixed SNR"
				description="Force solver to assume this signal to noise
				ratio on every pixel. Suggestion: Leave empty to let SP_Ace
				to estimate this parameter pixel-by-pixel."/>
			<inputKey name="rv_ini" multiplicity="single"
				ucd="spect.dopplerVeloc" unit="km/s"
				tablehead="Initial RV estimate"
				description="Provide an estimation of the object's radial
					velocity in case the spectrum is not RV-corrected."/>
			<inputKey name="compute_errors" type="boolean" multiplicity="single"
				tablehead="Compute Errors?"
				description="Make SP_Ace estimate errors (this increases runtime
					significantly)."/>
			<inputKey name="alpha" type="boolean" multiplicity="single"
				tablehead="Aggregate alpha?"
				description="If selected, the output will only contain two
					abundances: 'alphas' and 'metals', which are the estimation
					of the abundances of alpha-process elements and non-alpha-process
					elements as if they were one element.  Hint: This option is mostly
					useful with low-resolution (R ~ 2000) spectra.  See the tutorial
					for more details."/>
			<inputKey name="no_norm" type="boolean" multiplicity="single"
				tablehead="Keep original normalization?"
				description="Process the spectrum with the normalization provided
					by the user? This option switches off the internal
					re-normalization done by SP_Ace. Suggestion: Do not set unless
					you are absolutely sure that your normalization is right."/>
			<inputKey name="norm_rad" multiplicity="single"
				tablehead="Re-normalization parameter"
				description="This parameter determines the flexibility of
					the curve used to fit the continuum during the re-normalization
					of the spectrum.  The higher the number the more rigid the curve
					is (see discussion in the tutorial and in Sec.7.4 of the
					paper). Suggestion: leave the default value unless you know
					very well what you are doing.">30</inputKey>
			<inputKey name="abd_loop" type="boolean" multiplicity="single"
				tablehead="Use iterative loop?"
				description="Switch on the loop that iteratively derives
					stellar parameters (Teff, log g, [M/H]) and chemical
					abundances. Suggestion: do not set (see Sec. 8.5.1 of the paper).
					If your spectrum has low resolution (R ~ 2000), see discussion
					in the tutorial."/>
		</inputTable>"""

	def run(self, service, inputTable, queryMeta):
		return runSpace(inputTable.getParam("spectrum")[1].read(),
			inputTable.getParamDict(),
			service.rd)


def makeParFile(args, resdir):
	"""turns args (as defined by the core's input table) to the contents of
	a control file for SP_ACE.
	"""
	parFile = ["obs_sp_file 'spectrum.asc'",
		"GCOGlib '%s'"%os.path.join(resdir, "data/spectral-library-v1.4/"),
		"null_value null",
		"Salaris_MH"
		]

	# booleans
	for ourKey, theirKey  in [
			("compute_errors", "error_est"),
			("alpha", "alpha"),
			("no_norm", "no_norm"),
			("abd_loop", "ABD_loop")]:
		if args.get(ourKey, False):
			parFile.append(theirKey)

	if args.get("wave_lims"):
		try:
			pars = [int(w) for w in args["wave_lims"].split()]
			if len(pars)%2:
				raise ValueError("no interval")
		except ValueError:
			raise api.ValidationError("limits must be pairs of wavelengths in"
				" Angstrom (usefully between 4800 and 6860)", "wave_lims")
		parFile.append("wave_lims "+" ".join(str(n) for n in pars))

	# simple floats
	for ourKey, theirKey in [
			("fwhm", "fwhm"),
			("tforce", "T_force"),
			("gforce", "G_force"),
			("rv_ini", "RV_ini"),
			("snrforce", "sn_ratio"),
			("norm_rad", "norm_rad")]:
		if args.get(ourKey) is not None:
			parFile.append("%s %f"%(theirKey, args[ourKey]))

	return "\n".join(parFile)+"\n"


def runSpace(spectrum, inArgs, rd):
	"""runs Corrado's code on the data in spectrum.
	rd must be the s_pace/q RD; inArgs is a dictionary containing other
	program arguments as defined by Core's input table.
	"""
	parFile = makeParFile(inArgs, rd.resdir)
	binaryName = SPACE_BINARIES[inArgs.get("version")]

	wd = tempfile.mkdtemp("sandbox", dir=api.getConfig("tempDir"))
	try:
		with open(os.path.join(wd, "spectrum.asc"), "wb") as f:
			f.write(spectrum)
		with open(os.path.join(wd, "space.par"), "wb") as f:
			f.write(parFile.encode("utf-8"))

		if spectrum.startswith(b"This is not a spectrum.  Please just test."):
			with open(os.path.join(wd, "space_TGM_ABD.dat"), "wb") as f:
				f.write(rd.openRes("test_data/response_scaffold.dat").read())
			with open(os.path.join(wd, "space_model.dat"), "wb") as f:
				f.write(rd.openRes("test_data/space_model_scaffold.dat").read())

		else:
			try:
				subprocess.check_output([os.path.join(rd.resdir, "bin", binaryName),
					"space.par"], cwd=wd)
			except subprocess.CalledProcessError as ex:
				# in general, the last line of the output contains the error message
				try:
					errMsg = ex.output.split(b"\n")[-2].strip()
				except IndexError:
					errMsg = b"No useful error message from SP_ACE"
				raise api.ValidationError("Computation failed: %s"
						%errMsg.decode("utf-8", "ignore"),
					"spectrum")

		if DUMP_OUTPUT:
			with open(os.path.expanduser("~/response_scaffold.dat"), "wb") as f:
				f.write(open(os.path.join(wd, "space_TGM_ABD.dat"), "rb").read())
			with open(os.path.expanduser("~/space_model_scaffold.dat"), "wb") as f:
				f.write(open(os.path.join(wd, "space_model.dat"), "rb").read())

		res = api.makeData(rd.getById("import_programResult"),
			forceSource=os.path.join(wd, "space_TGM_ABD.dat"))
		res.getPrimaryTable().addMeta("info",
			parFile, infoName="parameter_file",
			infoValue="Parameters used for computation", format="literal")

		supplement = api.makeData(rd.getById("parse_resultspectrum"),
			forceSource=os.path.join(wd, "space_model.dat"))
		
		return rsc.makeCombinedData(res.dd, {
			"primary": res.getPrimaryTable(),
			"spectrum": supplement.getPrimaryTable()})
			

	finally:
		shutil.rmtree(wd)


if __name__=="__main__":
	from gavo import base
	base.DEBUG = True
	DUMP_OUTPUT = True
	api.LoggingUI(api.ui)
	runSpace(open("test_data/fast_spectrum.txt", "rb").read(),
		{"fwhm": 0.4, "wave_lims": "4800 6860", "compute_errors": True},
		api.getRD("sp_ace/q"))