"""
Tests for slap tables and services
"""

#c Copyright 2008-2023, 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

import unittest
from ctypes.util import find_library

from gavo import api
from gavo.helpers import testtricks
from gavo.helpers import trialhelpers
from gavo.protocols import linetap
from gavo.web import vodal

import tresc

if find_library(linetap._InChIBinding._libname):
	def skipIfNoInChI(obj):
		return obj
else:
	def skipIfNoInChI(obj):
		return unittest.skip("No libinchi")(obj)


class _SLAPTable(tresc.RDDataResource):
	rdName = "data/slaptest"
	dataId = "import"

slapTestTable = _SLAPTable()


class TestWithSLAPTable(testhelpers.VerboseTest):
	resources = [("table", slapTestTable)]

	def runService(self, params):
		res = trialhelpers.runSvcWith(
			api.resolveCrossId("data/slaptest#s"), "slap.xml", params)
		return res


class ImportingTest(TestWithSLAPTable):
	def testValues(self):
		line = list(self.table.connection.queryToDicts(
			"select * from test.slaptest where chemical_element='Mo'"))[0]
		self.assertAlmostEqual(line.pop("initial_level_energy"),
			1.65537152014313e-26)
		self.assertAlmostEqual(line.pop("final_level_energy"),
			4.96611457698311e-18)
		self.assertEqual(line, {
			'id_status': 'identified',
			'initial_name': 'Upper Level',
			'linename': 'Mo 400 A',
			'final_name': 'Lower Level', 'pub': '2003junk.yard.0123X',
			'chemical_element': 'Mo', 'wavelength': 4e-08})
		
	def testSimpleMeta(self):
		self.assertEqual(
			self.table.tableDef.getColumnByName("wavelength").utype,
			"ssldm:Line.wavelength.value")


class _SLAPMetadata(testhelpers.TestResource):
	resources = [("slapTable", slapTestTable)]

	def make(self, dependents):
		request = trialhelpers.FakeRequest("", args={"FORMAt": ["Metadata"]})
		rendered = request.processWithRoot(vodal.SLAPRenderer(request,
			api.resolveCrossId("data/slaptest#s")))
		tree = testhelpers.getXMLTree(rendered, debug=False)
		return tree


class ServiceMetadataTest(testhelpers.VerboseTest):
	resources = [("tree", _SLAPMetadata())]

	def testWavelength(self):
		params = self.tree.xpath("//PARAM[@name='INPUT:wavelength']")
		self.assertEqual(len(params), 1)
		param = params[0]
		self.assertEqual(param.get("utype"), "ssldm:Line.wavelength.value")
		self.assertEqual(param.get("unit"), "m")

	def testServiceKey(self):
		param = self.tree.xpath("//PARAM[@name='INPUT:verb']")[0]
		self.assertEqual(len(param.xpath("VALUES/OPTION")), 3)
	
	def testFieldsDeclared(self):
		param = self.tree.xpath("//FIELD[@ID='wavelength']")[0]
		self.assertEqual(param.get("utype"), "ssldm:Line.wavelength.value")

	def testQueryStatus(self):
		self.assertEqual(
			self.tree.xpath("RESOURCE/INFO[@name='QUERY_STATUS']")[0
				].get("value"), "OK")


class ServiceConstraintsTest(TestWithSLAPTable):
	def testChemicalElement(self):
		res = self.runService({"CHEMICAL_ELEMENT": ["Mo"]}
			).getPrimaryTable().rows
		self.assertEqual(len(res), 1)
		self.assertEqual(res[0]["chemical_element"], "Mo")

	def testChemicalElements(self):
		res = self.runService({"CHEMICAL_ELEMENT": ["Mo,Bi"]}
			).getPrimaryTable().rows
		self.assertEqual(len(res), 2)
		self.assertEqual(set(r["chemical_element"] for r in res),
			set(["Mo", "Bi"]))

	def testWavelength(self):
		res = self.runService({"WAVELENGTH": ["4e-8/1e-7"]}
			).getPrimaryTable().rows
		self.assertEqual(len(res), 1)
		self.assertEqual(res[0]["chemical_element"], "Mo")

	def testInitialEnergyLower(self):
		res = self.runService({"INITIAL_LEVEL_ENERGY": ["1.1e-29/"]}
			).getPrimaryTable().rows
		self.assertEqual(len(res), 2)
		self.assertEqual(set(r["chemical_element"] for r in res),
			set(["Mo", "Bi"]))

	def testInitialEnergyUpper(self):
		res = self.runService({"FINAL_LEVEL_ENERGY": ["/5e-18"]}
			).getPrimaryTable().rows
		self.assertEqual(len(res), 2)
		self.assertEqual(set(r["chemical_element"] for r in res),
			set(["Mo", "H"]))

	def testMAXREC(self):
		res = self.runService({"MAXREC": ["1"]}
			).getPrimaryTable().rows
		self.assertEqual(len(res), 1)


class _SLAPRegistryRecord(testhelpers.TestResource):
	def make(self, dependents):
		from gavo.registry import capabilities
		from gavo.web import vosi
		publication = api.resolveCrossId("data/slaptest#s").publications[0]
		res = vosi.CAP.capabilities[
			capabilities.getCapabilityElement(publication)].render()
		return testhelpers.getXMLTree(res, debug=False), res


class CapabilityTest(testhelpers.VerboseTest, testtricks.XSDTestMixin):
	resources = [("cap", _SLAPRegistryRecord())]

	def testValid(self):
		self.assertValidates(self.cap[1])

	def testType(self):
		self.assertEqual(self.cap[0][0].get(
			"{http://www.w3.org/2001/XMLSchema-instance}type"),
			"slap:SimpleLineAccess")

	def testStandardId(self):
		self.assertEqual(self.cap[0][0].get("standardID"),
			"ivo://ivoa.net/std/SLAP")

	def testComplianceLevel(self):
		self.assertEqual(self.cap[0][0].xpath("complianceLevel")[0].text,
			"full")

	def testDataSource(self):
		self.assertEqual(self.cap[0][0].xpath("dataSource")[0].text,
			"theoretical")

	def testTestQuery(self):
		self.assertEqual(self.cap[0][0].xpath("testQuery/queryDataCmd")[0].text,
			"MAXREC=1")


@skipIfNoInChI
class InchiTest(testhelpers.VerboseTest):
	def testNoLib(self):
		try:
			oldval = linetap._InChIBinding._libname
			linetap._InChIBinding._libname = "nonexistinglib"
			self.assertRaisesWithMsg(api.ReportableError,
				"No libinchi found.",
				linetap._InChIBinding,
				())
		finally:
			linetap._InChIBinding._libname = oldval

	def testInChIOk(self):
		self.assertEqual(
			# this *really* ought to work in strict mode, but alas, the library
			# complains about a FAIL_I2I right now.  Too lazy to investigate
			# right now.
			linetap.getILib().checkInChI("InChI=1S/CH4/h1H4", 0),
			0)

	def testInChINotOk(self):
		self.assertEqual(
			linetap.getILib().checkInChI("InCh=1S/CH44"),
			1)

	def testGetInChIKeyOk(self):
		self.assertEqual(
			linetap.getILib().getInChIKey("InChI=1S/CH4/h1H4"),
			"VNWKTOKETHGBQD-UHFFFAOYSA-N")

	def testGetInChiKeyNotOk(self):
		self.assertRaisesWithMsg(linetap.InChIError,
			"Could not generate InChIKey: 3",
			linetap.getILib().getInChIKey,
			("",))
	
	def testNormalizeInchi(self):
		self.assertEqual(
			linetap.getILib().normalizeInChI("InChI=1S/CH4/h1H4"),
			"InChI=1S/CH4/h1H4")
	
	def testAtomCreation(self):
		a = linetap.getILib().getAtom("Ta", 181)
		self.assertEqual(a.x, 0)
		self.assertEqual(a.elname, b"Ta")
		self.assertEqual(a.isotopic_mass, 181)
		self.assertEqual(a.charge, 0)

	def testAtomCreationWithCharge(self):
		a = linetap.getILib().getAtom("Ta", 181, -3)
		self.assertEqual(a.charge, -3)

	def testInputCreation(self):
		ii = linetap.getILib()
		mol = ii.getInput([
			ii.getAtom("O", 16),
			ii.getAtom("H", 1),
			ii.getAtom("H", 2)])
		mol.connect(0, 1)
		mol.connect(0, 2)

		self.assertEqual(mol.num_atoms, 3)

		self.assertEqual(mol.atom[0].elname, b"O")
		self.assertEqual(mol.atom[0].num_bonds, 2)
		self.assertEqual(mol.atom[0].neighbor[0], 1)
		self.assertEqual(mol.atom[0].neighbor[1], 2)

		self.assertEqual(mol.atom[1].num_bonds, 1)
		self.assertEqual(mol.atom[1].neighbor[0], 0)
		self.assertEqual(mol.atom[2].neighbor[0], 0)

	def testInChICreationIso(self):
		ii = linetap.getILib()
		mol = ii.getInput([
			ii.getAtom("O", 16),
			ii.getAtom("H", 1),
			ii.getAtom("H", 1)])
		mol.connect(0, 1)
		mol.connect(0, 2)

		self.assertEqual(ii.getInChI(mol), "InChI=1S/H2O/h1H2/i1+0/hH2")

	def testInChICreation(self):
		ii = linetap.getILib()
		mol = ii.getInput([
			ii.getAtom("H"),
			ii.getAtom("H"),
			ii.getAtom("C"),
			ii.getAtom("H"),
			ii.getAtom("H")])
		mol.connect(2, 0)
		mol.connect(2, 1)
		mol.connect(2, 3)
		mol.connect(2, 4)

		self.assertEqual(ii.getInChI(mol), "InChI=1S/CH4/h1H4")

	def testFailedCreation(self):
		ii = linetap.getILib()
		mol = ii.getInput([ii.getAtom("Ta")])
		self.assertRaisesWithMsg(linetap.InChIError,
			"Atom indexes must be < 1 here",
			mol.connect,
			(0, 1))


@skipIfNoInChI
class MolCodeTest(testhelpers.VerboseTest,
		metaclass=testhelpers.SamplesBasedAutoTest):
	def testErrorReporting(self):
		self.assertRaisesWithMsg(
			ValueError,
			"Invalid molecule code token: 'totalermurks'",
			linetap.getILib().getInChIFromCode,
			("12C H H totalermurks",))

	def _runTest(self, sample):
		molcode, expInchi, expWeight = sample
		parsed, foundInchi = linetap.getILib().getInChIFromCode(molcode)
		self.assertEqual(expInchi, foundInchi)
		self.assertEqual(expWeight, parsed.getTotalWeight())

	samples = [
		("H", "InChI=1S/H", 1),
		("H 1-2 H", "InChI=1S/H2/h1H", 2),
		("O 1=2 O", "InChI=1S/O2/c1-2", 32),
		("N 1#2 N", "InChI=1S/N2/c1-2", 28),
		("H H O 1-3 2-3", "InChI=1S/H2O/h1H2", 18),
#5
		("O O C 1=3 2=3", "InChI=1S/CO2/c2-1-3", 44),
		("H 1-2 C 2#3 C 3-4 C H H H 4-5 4-6 4-7",
			"InChI=1S/C3H4/c1-3-2/h1H,2H3", 40),
		("H 1-2 N 2=3 C 3-4 H 3-5 H", "InChI=1S/CH3N/c1-2/h2H,1H2", 29),
		("35S 1-2 F 1-3 F 1-4 F 1-5 F 1-6 F 1-7 F",
			"InChI=1S/F6S/c1-7(2,3,4,5)6/i7+3", 149),
		("34S+ 1-2 F 1-3 F 1-4 F 1-5 F 1-6 F",
			"InChI=1S/F5S/c1-6(2,3,4)5/q+1/i6+2", 129),
#10
		("55Fe++++", "InChI=1S/Fe/q+4/i1-1", 55),
		("D 1-2 T", "InChI=1S/H2/h1H/i1+2D", 5),
	]


@skipIfNoInChI
class InChIParsingTest(testhelpers.VerboseTest):
	def testNoInchi(self):
		self.assertRaisesWithMsg(linetap.InChIError,
			"No structural data provided",
			linetap.getILib().parseInChI,
			("keininchi",))

	def testBadLayer(self):
		self.assertRaisesWithMsg(linetap.InChIError,
			"No structural data provided",
			linetap.getILib().parseInChI,
			("InChI=1S/CH4/Breakthis",))

	def testNoMainLayer(self):
		self.assertRaisesWithMsg(linetap.InChIError,
			"No structural data provided",
			linetap.getILib().parseInChI,
			("InChI=1S",))
		self.assertRaisesWithMsg(linetap.InChIError,
			"No structural data provided",
			linetap.getILib().parseInChI,
			("InChI=1/",))

	def testGetChargeNoLayer(self):
		with linetap.getILib().parseInChI(
				"InChI=1S/C2H6O/c1-2-3/h3H,2H2,1H3") as mol:
			self.assertEqual(mol.getTotalCharge(), 0)

	def testGetChargeSimple(self):
		with linetap.getILib().parseInChI("InChI=1/Al/q+1") as mol:
			self.assertEqual(
				mol.getTotalCharge(), 1)


@skipIfNoInChI
class InChIMolweightTest(testhelpers.VerboseTest,
		metaclass=testhelpers.SamplesBasedAutoTest):
	def _runTest(self, sample):
		inchi, expectedMass = sample
		with linetap.getILib().parseInChI(inchi) as mol:
			self.assertEqual(mol.getTotalWeight(), expectedMass)
	
	samples = [
		("InChI=1S/H", 1),
		("InChI=1S/H2/h1H", 2),
		("InChI=1S/H2O/h1H2", 18),
		("InChI=1S/H2O/h1H2/i/hD", 19),
		("InChI=1S/H2O/h1H2/i/hD2", 20),
# 05
		("InChI=1S/H2O/h1H2/i/hTD", 21),
		("InChI=1S/CHBOSSi/c3-1-2-4(1)5/h5H", 100),
		("InChI=1S/CHBOSSi/c3-1-2-4(1)5/h5H/i1+2,2-2,4+3", 103),
		("InChI=1S/F6S/c1-7(2,3,4,5)6", 146),
	]


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