"""
IVOA cone search: Helper functions, a core, and misc.
"""

#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 import base
from gavo import svcs
from gavo.protocols import simbadinterface  #noflake: for registration
from gavo.svcs import outputdef


def getRadialCondition(td, ra, dec, sr,
		raColName=None, decColName=None):
	"""returns a sql literal for building a spatial query over the tableDef
	td.

	ra, dec, and sr are as in SCS.  While this function does some mild
	last-resort checking in case things break, you must not pass in
	untrusted content in there.  It's ok to pass in columns or expressions,
	though.

	raCol and decCol are determined through UCDs if not given.  This
	will preferentially use q3c and fall back to pgsphere if no q3c indices
	can be discerned on the columns.

	In case no SCS UCDs are present, the function will also accept
	spoints with pos.eq;meta.main.  This, of course, will not work
	with SCS itself (out of the box).
	"""
	if (isinstance(ra, str) and ";" in ra
		) or (isinstance(dec, str) and ";" in dec
		) or (isinstance(sr, str) and ";" in sr):
		raise base.ReportableError(
			"getRadialCondition's last resort alarm triggered.")

	try:
		if raColName is None:
			raCol = td.getColumnByUCD("pos.eq.ra;meta.main")
		else:
			raCol = td.getColumnByName(raColName)

		if decColName is None:
			decCol = td.getColumnByUCD("pos.eq.dec;meta.main")
		else:
			decCol = td.getColumnByName(decColName)
	except base.StructureError:
		# presumably a UCD location failure
		pointCols = td.getColumnsByUCD("pos.eq")
		if pointCols:
			return ("{ptname} <@ scircle("
				"spoint(radians({constra}), radians({constdec})),"
				" radians({sr}))").format(ptname=pointCols[0].name,
						constra=ra, constdec=dec, sr=sr)
		else:
			raise

	if "q3c" in (raCol.isIndexed() or []):
		pattern = ("q3c_radial_query({varra}, {vardec},"
			" {constra}, {constdec}, {sr})")
	else:
		pattern = ("spoint(radians({varra}), radians({vardec})) "
			" <@ scircle("
			"spoint(radians({constra}), radians({constdec})),"
			" radians({sr}))")
	
	return pattern.format(
		varra=raCol.name, vardec=decCol.name,
		constra=ra, constdec=dec, sr=sr)


def findNClosest(alpha, delta, tableDef, n, fields, searchRadius=5):
	"""returns the n objects closest around alpha, delta in table.

	n is the number of items returned, with the closest ones at the
	top, fields is a sequence of desired field names, searchRadius
	is a radius for the initial q3c search and will need to be
	lowered for dense catalogues and possibly raised for sparse ones.

	The last item of each row is the distance of the object from
	the query center in degrees.
	"""
	with base.getTableConn() as conn:
		raField = tableDef.getColumnByUCDs("pos.eq.ra;meta.main",
			"POS_EQ_RA_MAIN").name
		decField = tableDef.getColumnByUCDs("pos.eq.dec;meta.main",
			"POS_EQ_RA_MAIN").name
		res = list(conn.query("SELECT %s,"
				" (spoint(radians(%s), radians(%s)) <->"
				" spoint(radians(%%(alpha)s), radians(%%(delta)s))) as dist_"
				" FROM %s WHERE"
				" q3c_radial_query(%s, %s, %%(alpha)s, %%(delta)s,"
				" %%(searchRadius)s)"
				" ORDER BY dist_ LIMIT %%(n)s"%
					(",".join(fields), raField, decField, tableDef.getQName(),
						raField, decField),
			locals()))
		return res


def parseHumanSpoint(cooSpec, colName=None):
	"""tries to interpret cooSpec as some sort of cone center.

	Attempted interpretations include various forms of coordinate pairs
	and simbad objects; hence, this will in general cause network traffic.

	If no sense can be made, a ValidationError on colName is raised.
	"""
	try:
		cooPair = base.parseCooPair(cooSpec)
	except ValueError:
		simbadData = base.caches.getSesame("web").query(cooSpec)
		if not simbadData:
			raise base.ValidationError("%s is neither a RA,DEC"
				" pair nor a simbad resolvable object."%cooSpec, colName)
		cooPair = simbadData["RA"], simbadData["dec"]
	return cooPair


def getConeColumns(td):
	"""returns the columns the cone search will use as positions in a
	tableDef.

	This will raise an error if these are not present or not unique.
	Both new-style and old-style UCDs are accepted.
	"""
	raColumn = td.getColumnByUCDs(
		"pos.eq.ra;meta.main", "POS_EQ_RA_MAIN")
	decColumn = td.getColumnByUCDs(
		"pos.eq.dec;meta.main", "POS_EQ_DEC_MAIN")
	return raColumn, decColumn


class SCSCore(svcs.DBCore):
	"""A core performing cone searches.

	This will, if it finds input parameters it can make out a position from,
	add a _r column giving the distance between the match center and
	the columns that a cone search will match against.

	If any of the conditions for adding _r aren't met, this will silently
	degrade to a plain DBCore.

	You will almost certainly want a::

		<FEED source="//scs#coreDescs"/>
	
	in the body of this (in addition to whatever other custom conditions
	you may have).
	"""
	name_ = "scsCore"

	def onElementComplete(self):
		super().onElementComplete()
		# raColumn and decColumn must be from the queriedTable (rather than
		# the outputTable, as it would be preferable), since we're using
		# them to build database queries.
		self.raColumn, self.decColumn = getConeColumns(self.queriedTable)
		try:
			self.idColumn = self.outputTable.getColumnByUCDs(
				"meta.id;meta.main", "ID_MAIN")
		except base.StructureError:
			base.ui.notifyWarning("SCS core at %s: Output table has no unique"
				" meta.id;meta.main column.  This service will be invalid."%
				self.getSourcePosition())

		self.distCol = base.resolveCrossId("//scs#distCol")
		# let me indulge in this ugly hack -- outputTable belongs to
		# us anyway, and saving a copy is totally not worth it
		self.outputTable.columns[0:0] = [self.distCol]
		self.outputTable.columns.redoIndex()

		if not self.hasProperty("defaultSortKey"):
			self.setProperty("defaultSortKey", self.distCol.name)

	def _guessDestPos(self, inputTable):
		"""returns RA and Dec for a cone search possibly contained in inputTable.

		If no positional query is discernible, this returns None.
		"""
		pars = inputTable.getParamDict()
		if pars.get("RA") is not None and pars.get("DEC") is not None:
			return pars["RA"], pars["DEC"]
		elif pars.get("hscs_pos") is not None:
			try:
				return parseHumanSpoint(pars["hscs_pos"], "hscs_pos")
			except (ValueError, base.ValidationError):
				# We do not want to fail for this fairly unimportant thing.
				# If the core actually needs the position, it should fail itself.
				return None
		else:
			return None

	def _getDistColumn(self, destPos):
		"""returns an outputField selecting the distance of the match
		object to the cone center.
		"""
		if destPos is None:
			select = "NULL"
		else:
			select = "degrees(spoint(radians(%s), radians(%s)) <-> %s)"%(
				self.raColumn.name, self.decColumn.name,
				"spoint '(%fd,%fd)'"%destPos)

		return self.distCol.change(select=select)

	def _fixupQueryColumns(self, destPos, baseColumns):
		"""returns the output columns from baseColumns for a query
		centered at destPos.

		In particular, the _r column is primed so it yields the right result
		if destPos is given.
		"""
		res = []
		for col in baseColumns:
			if col.name=="_r":
				res.append(self._getDistColumn(destPos))
			else:
				res.append(col)
		return res

	def _makeResultTableDef(self, service, inputTable, queryMeta):
		destPos = self._guessDestPos(inputTable)

		outCols = self._fixupQueryColumns(destPos,
			self.getQueryCols(service, queryMeta))

		return base.makeStruct(outputdef.OutputTableDef,
			parent_=self.queriedTable.parent,
			id="result",
			onDisk=False,
			columns=outCols,
			params=self.queriedTable.params)
