"""
Some plausibility testing for our siap infrastructure.
"""

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

from gavo.helpers import testhelpers

from gavo import base
from gavo import rsc
from gavo import rscdesc  #noflake: for registration
from gavo.base import coords
from gavo.helpers import trialhelpers
from gavo.formats import votablewrite
from gavo.protocols import siap
from gavo.utils import fitstools

import tresc


class TestWCSTrafos(unittest.TestCase):
	"""Tests for transformations between WCS and pixel coordinates.
	"""
	wcs = {
		"CRVAL1": 0,   "CRVAL2": 0,
		"CRPIX1": 50,  "CRPIX2": 50,
		"CD1_1": 0.01, "CD1_2": 0,
		"CD2_1": 0,    "CD2_2": 0.01,
		"NAXIS1": 100, "NAXIS2": 100,
		"CUNIT1": "deg", "CUNIT2": "deg",
		"CTYPE1": 'RA---TAN-SIP', "CTYPE2": 'DEC--TAN-SIP',
		"LONPOLE": 180.,
		"NAXIS": 2,
	}

	def _testInvertibilityReal(self):
		for crvals in [(0,0), (40,60), (125,-60), (238,80)]:
			self.wcs["CRVAL1"], self.wcs["CRVAL2"] = crvals
			for crpixs in [(0,0), (50,50), (100,100), (150,0)]:
				self.wcs["CRPIX1"], self.wcs["CRPIX2"] = crpixs
				for pixpos in [(0,0), (50, 50), (0, 100), (100, 0)]:
					fwTrafo = coords.getWCSTrafo(self.wcs)
					bwTrafo = coords.getInvWCSTrafo(self.wcs)
					x, y = bwTrafo(*fwTrafo(*pixpos))
					self.assertAlmostEqual(x, pixpos[0], 5)
					self.assertAlmostEqual(y, pixpos[1], 5)

	def testInvertibility(self):
		self._testInvertibilityReal()
		self.wcs["CD2_1"] = 0.001
		self._testInvertibilityReal()
		self.wcs["CD1_2"] = -0.001
		self._testInvertibilityReal()


class PixelScaleTest(testhelpers.VerboseTest):
	wcs = {
		"CRPIX1": 50, "CRPIX2": 50,
		"CD1_1": 1E-04,  "CD1_2": -1E-06,
		"CD2_1": -1E-06, "CD2_2": -1E-04,
		"CTYPE1": 'RA---TAN', "CTYPE2": 'DEC--TAN',
		"CRVAL1": None, "CRVAL2": None,
		"NAXIS1": 100, "NAXIS2": 100,
		"CUNIT1": "deg", "CUNIT2": "deg",}

	def testEquator(self):
		wcs = self.wcs.copy()
		wcs["CRVAL1"], wcs["CRVAL2"] = 0, 0
		self.assertAlmostEqualVector(coords.getPixelSizeDeg(wcs), (1e-4, 1e-4))

	def testStitch(self):
		wcs = self.wcs.copy()
		wcs["CRVAL1"], wcs["CRVAL2"] = -1e-3, 0
		self.assertAlmostEqualVector(coords.getPixelSizeDeg(wcs), (1e-4, 1e-4))
	
	def testNorth(self):
		wcs = self.wcs.copy()
		wcs["CRVAL1"], wcs["CRVAL2"] = 359.99, 70
		self.assertAlmostEqualVector(coords.getPixelSizeDeg(wcs), (1e-4, 1e-4))

	def testPole(self):
		wcs = self.wcs.copy()
		wcs["CRVAL1"], wcs["CRVAL2"] = 10, -90
		rasz, decsz = coords.getPixelSizeDeg(wcs)
		self.assertAlmostEqual(decsz, 1e-4, 5)
		self.assertFalse(rasz>1)


class CooQueryTestBase(testhelpers.VerboseTest):
	"""base class for functional testing of the SIAP code.
	"""
	data = None

	# queries with expected numbers of returned items
	_testcases = [
		("0,0", "1", (1, 0, 1, 1)),
		("45,-45.6", "1", (0, 0, 0, 1)),
		("1,45", "3,1", (1, 0, 3, 3)),
		("1,46", "1.1", (0, 0, 0, 3)),
		("161,45", "3,1", (1, 0, 3, 3)),
		("161,46", "1.1", (0, 0, 0, 3)),
		("0,90", "360,2", (0, 1, 1, 1)),
# XXX TODO: do some more here
	]
	_intersectionIndex = {
		"COVERS": 0,
		"ENCLOSED": 1,
		"CENTER": 2,
		"OVERLAPS": 3,
	}

	def _runTests(self, type):
		if self.data is None:  # only do something if given data (see below)
			return
		table = self.data.getPrimaryTable()
		for center, size, expected in self._testcases:
			pars = {}
			fragment = siap.getQuery(table.tableDef, {
				"POS": center,
				"SIZE": size,
				"INTERSECT": type}, pars)
			res = list(table.connection.query(
				"SELECT * FROM %s WHERE %s"%(table.tableName, fragment),
				pars))
			self.assertEqual(len(res), expected[self._intersectionIndex[type]],
				"%d instead of %d matched when queriying for %s %s"%(len(res),
				expected[self._intersectionIndex[type]], type, (center, size)))

	# queries with expected numbers of returned items
	_testcases = [
		("0,0", "1", (1, 0, 1, 1)),
		("45,-45.6", "1", (0, 0, 0, 1)),
		("1,45", "3,1", (1, 0, 3, 3)),
		("1,46", "1.1", (0, 0, 0, 3)),
		("161,45", "3,1", (1, 0, 3, 3)),
		("161,46", "1.1", (0, 0, 0, 3)),
		("0,90", "360,2", (0, 1, 1, 1)),
# XXX TODO: do some more here
	]

	def testCOVERS(self):
		self._runTests("COVERS")
	
	def testENCLOSED(self):
		self._runTests("ENCLOSED")
	
	def testCENTER(self):
		self._runTests("CENTER")
	
	def testOVERLAPS(self):
		self._runTests("OVERLAPS")


class PgSphereQueriesTest(CooQueryTestBase):
	"""tests for actual queries on the sphere with trivial WCS data.
	"""
	resources = [("data", tresc.siapTestTable)]


class CutoutQueriesTest(testhelpers.VerboseTest):
	resources = [("data", tresc.siapTestTable)]

	def _runQuery(self, params):
		svc = testhelpers.getTestRD().getById("siapcutout")
		return trialhelpers.runSvcWith(svc, "siap.xml", params
			).getPrimaryTable().rows

	def testWCSFixedSimple(self):
		row = self._runQuery({"POS": ["1,2"], "SIZE": ["0.1,0.2"]})[0]
		self.assertAlmostEqual(row["wcs_cdmatrix"][0], 0.01)
		self.assertAlmostEqual(row["wcs_cdmatrix"][1], 0)
		self.assertEqual(row["pixelSize"], [11, 21])
		self.assertEqual(row["coverage"].asSTCS("Unknown"),
			"Polygon Unknown 0.9499129572 1.8990431715 0.9499129572 2.0987721868 1.0498824794 2.0987083322 1.0498824794 1.8989853842")

	def testWCSFixedCrop(self):
		row = self._runQuery({"POS": ["5,5"], "SIZE": ["1,2"]})[0]
		self.assertEqual(row["pixelSize"], [50, 99])
		self.assertEqual(row["coverage"].asSTCS("Unknown"),
			"Polygon Unknown 4.5007198987 4.0010872016 4.5007198987 4.9720632011 4.9873652888 4.9685772926 4.9873652888 3.9982770726")


class ImportTest(testhelpers.VerboseTest):
	resources = [("conn", tresc.dbConnection)]
	setupCost = 0.5
	noWCSRec = {
			"imageTitle": "uncalibrated image",
			"NAXIS1": 1000,
			"NAXIS2": 100,
			"NAXIS": 2,
			"dateObs": None,
			"accref": "uu",
			"accsize": None,
			"embargo": None,
			"owner": None,}

	def testPlainPGS(self):
		dd = testhelpers.getTestRD().getById("pgs_siaptest")
		data = rsc.makeData(dd, connection=self.conn, forceSource=[
			testhelpers.computeWCSKeys((34, 67), (0.3, 0.4))])
		try:
			table = data.tables["pgs_siaptable"]
			res = list(table.iterQuery([table.tableDef.getColumnByName(n) for
				n in ("centerAlpha", "centerDelta")], ""))[0]
			self.assertEqual(int(res["centerDelta"]), 67)
		finally:
			data.dropTables(rsc.parseNonValidating)
			self.conn.commit()
	
	def testRaisingOnNull(self):
		dd = testhelpers.getTestRD().getById("pgs_siaptest")
		self.assertRaises(base.ValidationError,
			rsc.makeData,
			dd, connection=self.conn, forceSource=[self.noWCSRec])

	def testNullIncorporation(self):
		dd = testhelpers.getTestRD().getById("pgs_siapnulltest")
		data = rsc.makeData(dd, connection=self.conn, forceSource=[
			self.noWCSRec, testhelpers.computeWCSKeys((34, 67), (0.25, 0.5))])
		try:
			table = data.tables["pgs_siaptable"]
			res = sorted(table.iterQuery(
				[table.tableDef.getColumnByName("accref")], ""),
				key=lambda d: d["accref"])
			self.assertEqual(res,
				[{'accref': 'image/(34, 67)/(0.25, 0.5)'}, {'accref': 'uu'}])
		finally:
			data.dropTables(rsc.parseNonValidating)
			self.conn.commit()


class SIAPTestResponse(testhelpers.TestResource):
	resources = [("siapTable", tresc.siapTestTable)]

	def make(self, deps):
		svc = testhelpers.getTestRD().getById("pgsiapsvc")
		data = trialhelpers.runSvcWith(svc, "siap.xml",
			{"POS": "0, 0", "SIZE": "180,180"})
		vot = votablewrite.getAsVOTable(data,
			tablecoding="td", suppressNamespace=True)
		return vot, testhelpers.getXMLTree(vot, debug=False)

_siapTestResponse = SIAPTestResponse()


class SIAPResponseTest(testhelpers.VerboseTest):

	resources = [("siapResp", _siapTestResponse)]

	def _getSTCGroup(self):
		try:
			return _siapTestResponse._stcGroup
		except AttributeError:
			_siapTestResponse._stcGroup = self.siapResp[1].find(
			".//GROUP[@utype='stc:CatalogEntryLocation']")
		return _siapTestResponse._stcGroup

	def _getFieldForUtype(self, utype):
		ref = self._getSTCGroup().find(
			"FIELDref[@utype='%s']"%utype).get("ref")
		return self.siapResp[1].find(".//FIELD[@ID='%s']"%ref)

	def testSTCDefined(self):
		self.assertTrue(len(self._getSTCGroup()))

	def testCoverageDefined(self):
		self.assertEqual(
			self._getFieldForUtype("stc:AstroCoordArea.Polygon").get("name"),
			"coverage")

	def testSpectralDefined(self):
		self.assertEqual(
			self._getFieldForUtype("stc:AstroCoordArea.SpectralInterval.HiLimit"
				).get("name"),
			"bandpassHi")

	def testPositionDefined(self):
		self.assertEqual(
			self._getFieldForUtype("stc:AstroCoords.Position2D.Value2.C1",
				).get("name"),
			"centerAlpha")

	def testRefsysDefined(self):
		self.assertEqual(self._getSTCGroup().find(
			"PARAM[@utype='stc:AstroCoordSystem.SpaceFrame.CoordRefFrame']").get(
				"value"), "ICRS")

	def testTimesys(self):
		timesysEls = self.siapResp[1].xpath("RESOURCE[@type='results']/TIMESYS")
		timesysEl = testhelpers.pickSingle(timesysEls)
		self.assertEqual(timesysEl.get("refposition"), "UNKNOWN")
		self.assertEqual(timesysEl.get("timeorigin"), "MJD-origin")
		self.assertEqual(timesysEl.get("timescale"), "TT")


class ScaleHeaderTest(testhelpers.VerboseTest):
# This is a test for fitstricks.shrinkWCSHeader.
# It's here just because we already deal with wcs in this module
	def _assertCoordsMatch(self, pair1, pair2):
		self.assertAlmostEqual(pair1[0], pair2[0])
		self.assertAlmostEqual(pair1[1], pair2[1])

	def testSimple(self):
		fullHdr = fitstools.headerFromDict(
			testhelpers.computeWCSKeys((23, 27), (1, 2), cutCrap=True))
		halfHdr = fitstools.shrinkWCSHeader(fullHdr, 2)

		self.assertEqual(halfHdr["IMSHRINK"], 'Image scaled down 2-fold by DaCHS')
		self.assertEqual(fullHdr["NAXIS2"]/2, halfHdr["NAXIS1"])

		toPhysOld = coords.getWCSTrafo(fullHdr)
		toPhysNew = coords.getWCSTrafo(halfHdr)

		self._assertCoordsMatch(
			toPhysOld(fullHdr["CRPIX1"],fullHdr["CRPIX2"]),
			toPhysNew(halfHdr["CRPIX1"],halfHdr["CRPIX2"]))

		self._assertCoordsMatch(
			toPhysOld(1, 1), toPhysNew(1, 1))

		self._assertCoordsMatch(
			toPhysOld(fullHdr["NAXIS1"]+1,fullHdr["NAXIS2"]+1),
			toPhysNew(halfHdr["NAXIS1"]+1,halfHdr["NAXIS2"]+1))

	def testTypesFixed(self):
		fullHdr = fitstools.headerFromDict(
			testhelpers.computeWCSKeys((23, 27), (1, 2), cutCrap=True))
		fullHdr.set("BZERO", 32768)
		fullHdr.set("BSCALE", 1)
		fullHdr.set("BITPIX", 8)
		halfHdr = fitstools.shrinkWCSHeader(fullHdr, 2)
		self.assertFalse("BZERO" in halfHdr)
		self.assertFalse("BSCALE" in halfHdr)
		self.assertEqual(halfHdr["BITPIX"], -32)


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