"""
Tests for conversion between utype sequences and STC ASTs.
"""

#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.


from gavo.helpers import testhelpers

from gavo import stc
from gavo.stc import dm


class CoosysGenerTest(testhelpers.VerboseTest):
	"""tests for generation of coosys utype dicts from STC ASTs.
	"""
	def _assertSetmatch(self, literal, expected):
		ast = stc.parseSTCS(literal)
		self.assertEqual(dict(stc.getUtypes(ast)),
			expected)

	def testRef(self):
		tree = dm.STCSpec(astroSystem=stc.getLibrarySystem("TT-ICRS-OPT-BARY-TOPO"))
		self.assertEqual(stc.getUtypes(tree),
			[('stc:AstroCoordSystem.href',
				'ivo://STClib/CoordSys#TT-ICRS-OPT-BARY-TOPO')])

	def testSimpleSystems(self):
		self._assertSetmatch("Time TT Position ICRS GEOCENTER Spectral BARYCENTER"
			" Redshift BARYCENTER VELOCITY OPTICAL", {
				'stc:AstroCoordSystem.SpaceFrame.CoordRefFrame': 'ICRS',
				'stc:AstroCoordSystem.TimeFrame.TimeScale': 'TT',
				'stc:AstroCoordSystem.RedshiftFrame.DopplerDefinition': 'OPTICAL',
				'stc:AstroCoordSystem.RedshiftFrame.value_type': 'VELOCITY',
				'stc:AstroCoordSystem.RedshiftFrame.ReferencePosition': 'BARYCENTER',
				'stc:AstroCoordSystem.SpaceFrame.CoordFlavor': 'SPHERICAL',
				'stc:AstroCoordSystem.SpectralFrame.ReferencePosition': 'BARYCENTER',
				'stc:AstroCoordSystem.SpaceFrame.ReferencePosition': 'GEOCENTER'})
	
	def testWithEquinox(self):
		self._assertSetmatch("Position FK4 J1975.0", {
			'stc:AstroCoordSystem.SpaceFrame.CoordRefFrame': 'FK4',
			'stc:AstroCoordSystem.SpaceFrame.CoordRefFrame.Equinox': 'J1975.0',
			'stc:AstroCoordSystem.SpaceFrame.CoordFlavor': 'SPHERICAL',})
	
	def testWithDistance(self):
		self._assertSetmatch("Position ICRS BARYCENTER SPHER3 unit deg deg pc",{
			'stc:AstroCoordSystem.SpaceFrame.CoordFlavor.coord_naxes': '3',
			'stc:AstroCoordSystem.SpaceFrame.CoordRefFrame': 'ICRS',
			'stc:AstroCoordSystem.SpaceFrame.CoordFlavor': 'SPHERICAL',
			'stc:AstroCoordSystem.SpaceFrame.ReferencePosition': 'BARYCENTER'})

	def testWithPleph(self):
		self._assertSetmatch("Position ICRS BARYCENTER JPL-DE405", {
			'stc:AstroCoordSystem.SpaceFrame.CoordFlavor': 'SPHERICAL',
			'stc:AstroCoordSystem.SpaceFrame.CoordRefFrame': 'ICRS',
			'stc:AstroCoordSystem.SpaceFrame.ReferencePosition.PlanetaryEphem':
				'JPL-DE405',
			'stc:AstroCoordSystem.SpaceFrame.ReferencePosition': 'BARYCENTER'})


CR = stc.ColRef


class CooGenerTest(testhelpers.VerboseTest):
	"""tests for generation of column utype dicts from STC ASTs.
	"""
	def _assertAssmatch(self, literal, truth):
		ast = stc.parseQSTCS(literal)
		gr = [p for p in stc.getUtypes(ast)
			if not p[0].startswith("stc:AstroCoordSystem")]
		self.assertEqual(dict(gr), truth)
	
	def testTrivialPos(self):
		self._assertAssmatch('Position ICRS "ra" "dec"', {
			'stc:AstroCoords.Position2D.Value2.C1': CR("ra"),
			'stc:AstroCoords.Position2D.Value2.C2': CR("dec"),
		})

	def testMixedPos(self):
		self._assertAssmatch('Position ICRS "ra" 20.0', {
			'stc:AstroCoords.Position2D.Value2.C1': CR("ra"),
			'stc:AstroCoords.Position2D.Value2.C2': '20.0',
		})

	def testVecPos(self):
		self._assertAssmatch('Position ICRS [point]',
			{'stc:AstroCoords.Position2D.Value2': CR("point")})
	
	def testVecPos3(self):
		self._assertAssmatch('Position ICRS SPHER3 [point]',
			{'stc:AstroCoords.Position3D.Value3': CR('point')})

	def testVecEpoch(self):
		self._assertAssmatch('Position ICRS Epoch J2010.5', {
			'stc:AstroCoords.Position2D.Epoch': '2010.5',
			'stc:AstroCoords.Position2D.Epoch.yearDef': 'J'})

	def testBesselEpoch(self):
		self._assertAssmatch('Position ICRS Epoch B2010.5', {
			'stc:AstroCoords.Position2D.Epoch': '2010.5',
			'stc:AstroCoords.Position2D.Epoch.yearDef': 'B'})

	def testVeloc(self):
		self._assertAssmatch(
			'Position ICRS VelocityInterval Velocity "pmra" "pmde"', {
				'stc:AstroCoords.Velocity2D.Value2.C2': CR('pmde'),
				'stc:AstroCoords.Velocity2D.Value2.C1': CR('pmra')})

	def testErrorRadius(self):
		self._assertAssmatch(
			'Position ICRS Error "ep" "ep"', {
				'stc:AstroCoords.Position2D.Error2Radius': CR('ep')})

	def testTime(self):
		self._assertAssmatch('Time TT "obsDate"', {
			'stc:AstroCoords.Time.TimeInstant': CR("obsDate"),})

	def testTimeA(self):
		self._assertAssmatch(
			'TimeInterval TT "start" "end"', {
				'stc:AstroCoordArea.TimeInterval.StartTime': CR('start'),
				'stc:AstroCoordArea.TimeInterval.StopTime': CR('end')})
	
	def testGeoRef(self):
		self._assertAssmatch('Circle FK5 J1000.0 [circle]', {
			'stc:AstroCoordArea.Circle': CR('circle')})
		self._assertAssmatch('Box FK5 J1000.0 [bbox]', {
			'stc:AstroCoordArea.Box': CR('bbox')})
		self._assertAssmatch('Polygon FK5 J1000.0 [poly]', {
			'stc:AstroCoordArea.Polygon': CR('poly')})

	def testGeoSplit(self):
		self._assertAssmatch('Circle FK5 J1000.0 "cx" "cy" "radius"', {
			'stc:AstroCoordArea.Circle.Center.C2': CR('cy'),
			'stc:AstroCoordArea.Circle.Center.C1': CR('cx'),
			'stc:AstroCoordArea.Circle.Radius': CR('radius')})

	def testError(self):
		self._assertAssmatch('Position ICRS Error "e_ra" "e_dec"', {
			'stc:AstroCoords.Position2D.Error2.C1': CR('e_ra'),
			'stc:AstroCoords.Position2D.Error2.C2': CR('e_dec')})

	def testEpoch(self):
		self._assertAssmatch('Position ICRS Epoch "epum" 10 20', {
			'stc:AstroCoords.Position2D.Epoch': CR('epum'),
			'stc:AstroCoords.Position2D.Epoch.yearDef': 'J',
			'stc:AstroCoords.Position2D.Value2.C1': '10.0',
			'stc:AstroCoords.Position2D.Value2.C2': '20.0'})

	def testRedshift(self):
		self._assertAssmatch('Redshift "z" Error "zErr"', {
			'stc:AstroCoords.Redshift.Value': CR('z'),
			'stc:AstroCoords.Redshift.Error': CR('zErr')
		})

	def testCombined(self):
		self._assertAssmatch('Time TT "dateObs" Error "e_date"'
			' Circle ICRS "cra" "cdec" "crad" Position "ra" "dec"'
			'   Error "e_ra" "e_dec" Size "s_ra" "s_dec"'
			' SpectralInterval "bandLow" "bandHigh"'
			' Redshift "z" Error "zErr"', {
					'stc:AstroCoordArea.SpectralInterval.HiLimit': CR('bandHigh'),
					'stc:AstroCoordArea.SpectralInterval.LoLimit': CR('bandLow'),
					'stc:AstroCoordArea.Circle.Center.C2': CR('cdec'),
					'stc:AstroCoordArea.Circle.Center.C1': CR('cra'),
					'stc:AstroCoordArea.Circle.Radius': CR('crad'),
					'stc:AstroCoords.Time.TimeInstant': CR('dateObs'),
					'stc:AstroCoords.Position2D.Value2.C2': CR('dec'),
					'stc:AstroCoords.Position2D.Error2.C2': CR('e_dec'),
					'stc:AstroCoords.Position2D.Error2.C1': CR('e_ra'),
					'stc:AstroCoords.Time.Error': CR('e_date'),
					'stc:AstroCoords.Position2D.Value2.C1': CR('ra'),
					'stc:AstroCoords.Position2D.Size2.C2': CR('s_dec'),
					'stc:AstroCoords.Position2D.Size2.C1': CR('s_ra'),
					'stc:AstroCoords.Redshift.Value': CR('z'),
					'stc:AstroCoords.Redshift.Error': CR('zErr'),
				})


class UtypeASTTest(testhelpers.VerboseTest):
	"""tests for building STC ASTs out of utype sequences.
	"""
	def _getASTFromSTCS(self, stcs):
		ast = stc.parseQSTCS(stcs)
		utypes = stc.getUtypes(ast)
		return stc.parseFromUtypes(utypes)

	def testSimplePos(self):
		ast = self._getASTFromSTCS('Position GALACTIC "long" "lat"')
		self.assertEqual(ast.astroSystem.spaceFrame.refFrame, "GALACTIC_II")
		self.assertEqual(ast.place.value[0].dest, "long")
		self.assertEqual(ast.place.frame.refFrame, "GALACTIC_II")

	def testWithError(self):
		ast = self._getASTFromSTCS('Position GALACTIC Error "e1" "e1"')
		self.assertEqual(ast.place.error.radii[0].dest, "e1")

	def testWithEquinox(self):
		ast = self._getASTFromSTCS("Position FK4 J1975.0 Epoch B2000.0")
		self.assertEqual(ast.astroSystem.spaceFrame.equinox, "J1975.0")
		self.assertEqual(ast.place.yearDef, "B")
		self.assertEqual(ast.place.epoch, 2000.)

	def testTime(self):
		ast = self._getASTFromSTCS(
			'Time TT TOPOCENTER "dateObs" Error "clockdamage"')
		self.assertEqual(ast.time.value.dest, "dateObs")
		self.assertEqual(ast.time.error.values[0].dest, "clockdamage")
		self.assertEqual(ast.time.frame.refPos.standardOrigin, "TOPOCENTER")
		self.assertEqual(ast.time.frame.timeScale, "TT")

	def testGeoComp(self):
		ast = self._getASTFromSTCS('Circle ICRS [errc]')
		self.assertEqual(ast.areas[0].geoColRef.dest, "errc")

	def testTimeTDBScale(self):
		ast = self._getASTFromSTCS('Time TDB "time"')
		self.assertEqual(ast.time.frame.timeScale, 'TDB')

	def testPleph(self):
		ast = self._getASTFromSTCS('Time TDB TOPOCENTER JPL-DE405')
		self.assertEqual(ast.astroSystem.timeFrame.refPos.planetaryEphemeris,
			"JPL-DE405")


class UtypeRoundtripTest(testhelpers.VerboseTest, metaclass=testhelpers.SamplesBasedAutoTest):
	"""tests for working roundtrip of utype de-/serialization.
	"""

	def _assertSublist(self, sub, full):
		sub, full = dict(sub), dict(full)
		for k, v in sub.items():
			if k not in full:
				raise AssertionError("%s not in full"%k)
			if sub[k]!=full[k]:
				raise AssertionError("%s has wrong value, %s instead of %s"%(
					k, full[k], sub[k]))

	def _runTest(self, args):
		inTypes = args
		ast = stc.parseFromUtypes(inTypes)
		outTypes = stc.getUtypes(ast, includeDMURI=True)
		# allow additional keys in output (non-removed defaults)
		self._assertSublist(inTypes, outTypes)

	samples = [
		[],
		[('stc:AstroCoordSystem.SpaceFrame.CoordRefFrame', 'ICRS')],
		[('stc:AstroCoordSystem.SpaceFrame.CoordRefFrame', 'ICRS'),
			('stc:AstroCoords.Position2D.Epoch', '2002.0'),
			('stc:DataModel.URI', 'http://www.ivoa.net/xml/STC/stc-v1.30.xsd'),
		],
		[('stc:AstroCoords.Position2D.Value2.C1', stc.ColRef('ra')),
			('stc:AstroCoords.Position2D.Value2.C2', stc.ColRef('dec'))],
		[('stc:AstroCoordSystem.SpaceFrame.CoordRefFrame', 'ICRS'),
			('stc:AstroCoordSystem.TimeFrame.TimeScale', 'TT'),
			('stc:AstroCoordSystem.RedshiftFrame.DopplerDefinition', 'OPTICAL'),
			('stc:AstroCoordSystem.RedshiftFrame.value_type', 'VELOCITY'),
			('stc:AstroCoordSystem.RedshiftFrame.ReferencePosition', 'BARYCENTER'),
			('stc:AstroCoordSystem.SpaceFrame.CoordFlavor', 'SPHERICAL'),
			('stc:AstroCoordSystem.SpectralFrame.ReferencePosition', 'BARYCENTER'),
			('stc:AstroCoordSystem.SpaceFrame.ReferencePosition', 'GEOCENTER'),
		 	('stc:AstroCoords.Position2D.Value2.C1', stc.ColRef('ra')),
		  ('stc:AstroCoords.Position2D.Value2.C2', stc.ColRef('dec'))],
		[('stc:AstroCoordSystem.SpaceFrame.CoordRefFrame', 'ICRS'),
			('stc:AstroCoordSystem.SpaceFrame.CoordFlavor', 'CARTESIAN'),
			('stc:AstroCoordSystem.SpaceFrame.CoordFlavor.coord_naxes', '3')],
		[
			('stc:AstroCoords.Time.TimeInstant', '2000-01-01T00:00:00'),
			('stc:AstroCoordArea.Circle', stc.ColRef('errc'))],
		[
			('stc:AstroCoordSystem.href',
				'ivo://STClib/CoordSys#TT-ICRS-OPT-BARY-TOPO')],
		[
			('stc:AstroCoordSystem.SpaceFrame.CoordFlavor', 'SPHERICAL'),
			('stc:AstroCoordSystem.SpaceFrame.CoordFlavor.coord_naxes', '3'),
			('stc:AstroCoordSystem.SpaceFrame.CoordRefFrame', 'ICRS'),
			('stc:AstroCoords.Position3D.Value3.C1', stc.ColRef('raj2000')),
			('stc:AstroCoords.Position3D.Value3.C2', stc.ColRef('dej2000')),
			('stc:AstroCoords.Position3D.Value3.C3', stc.ColRef('dist_binney'))
		],
		[
			('stc:AstroCoordSystem.SpaceFrame.CoordFlavor', 'SPHERICAL'),
			('stc:AstroCoordSystem.SpaceFrame.CoordFlavor.coord_naxes', '3'),
			('stc:AstroCoordSystem.SpaceFrame.CoordRefFrame', 'ICRS'),
			('stc:AstroCoords.Position3D.Value3.C2', stc.ColRef('dej2000')),
		],

	]


if __name__=="__main__":
	testhelpers.main(CooGenerTest)