"""
"User" defined functions, i.e., ADQL functions defined only on this
system.

See the userFunction docstring on how to use these.
"""

#c Copyright 2008-2020, the GAVO project
#c
#c This program is free software, covered by the GNU GPL.  See the
#c COPYING file in the source distribution.


from gavo import stc
from gavo import utils
from gavo.adql import common
from gavo.adql import fieldinfo
from gavo.adql import morphhelpers
from gavo.adql import nodes
from gavo.adql import tree


UFUNC_REGISTRY = {}


def getStringLiteral(node, description="Argument"):
	"""ensures that node only contains a constant string literal
	and returns its value if so.

	The function raises an UfuncError mentioning description otherwise.
	"""
	if node.type!='characterStringLiteral':
		raise common.UfuncError(
			"%s must be a constant string literal"%description)
	
	return node.value


def userFunction(name, signature, doc, returntype="double precision", 
		unit="", ucd=""):
	"""a decorator adding some metadata to python functions to make
	them suitable as ADQL user defined functions.

	name is the name the function will be visible under in ADQL; signature is a
	signature not including the name of the form '(parName1 type1, parName1
	type2) -> resulttype'; doc is preformatted ASCII documentation.  The
	indentation of the second line will be removed from all lines.

	returntype is the SQL return type, which defaults to double
	precision.  While ADQL 2.0 appears to say that UDFs must be
	numeric, in practice nobody cares; so, return whatever you see fit.
	
	unit and ucd are optional for when you actually have a good guess what's
	coming back from your ufunc.  They can also be callable; in that
	case, they'll be passed the (annotated) arguments, and whatever
	they return will be the unit/ucd.

	The python function receives an array of arguments; this will in
	general be ADQL expression trees.  It must return either 
	
	* a string that will go literally into the eventual serialised SQL 
	  string (so take care to quote; in general, you will use 
	  nodes.flatten(arg) to flatten individual args); 
	* or they may return None, in which case the expression tree
	  remains unchanged.  This is for when the actual implementation is
	  in the database.
	* or they may raise nodes.ReplaceNode(r), where r is a nodes.ADQLNode 
		instance, which then replaces the user defined function in the parse 
		tree and will be annotated as usual.

	If you receive bad arguments or something else goes awry, raise
	a UfuncError.
	"""
	def deco(f):
		f.adqlUDF_name = name
		f.adqlUDF_signature = f.adqlUDF_name+signature.strip()
		f.adqlUDF_doc = utils.fixIndentation(doc, "", 1).strip()
		f.adqlUDF_returntype = returntype
		f.adqlUDF_unit = unit
		f.adqlUDF_ucd = ucd
		UFUNC_REGISTRY[f.adqlUDF_name.upper()] = f
		return f
	return deco


def _makeBooleanizer(funcName, booleanExpr):
	"""makes and registers a booleanizer for funcName.

	booleanExpr is the postgres expression the function should be booleanized
	to.  Refer to the two arguments as %(1)s and %(2)s
	"""
	def _booleanizeThis(node, operator, operand):
		if len(node.args)!=2:
			raise common.UfuncError("%s takes exactly two arguments"%funcName)
		return morphhelpers.addNotToBooleanized(
			booleanExpr%{
				'1': nodes.flatten(node.args[0]), 
				'2': nodes.flatten(node.args[1])},
			operator, operand)

	morphhelpers.registerBooleanizer(funcName.upper(), _booleanizeThis)



@userFunction("gavo_match",
	"(pattern TEXT, string TEXT) -> INTEGER",
	"""
	gavo_match returns 1 if the POSIX regular expression pattern
	matches anything in string, 0 otherwise.
	""",
	"integer")
def _match(args):
	if len(args)!=2:
		raise common.UfuncError("gavo_match takes exactly two arguments")
	return "(CASE WHEN %s ~ %s THEN 1 ELSE 0 END)"%(
		nodes.flatten(args[1]), nodes.flatten(args[0]))


@userFunction("ivo_hasword",
	"(haystack TEXT, needle TEXT) -> INTEGER",
	"""
	gavo_hasword returns 1 if needle shows up in haystack, 0 otherwise.  This
	is for "google-like"-searches in text-like fields.  In word, you can
	actually employ a fairly complex query language; see
	http://www.postgresql.org/docs/8.3/static/textsearch.html
	for details.
	""",
	"integer")
def _hasword(args):
	if len(args)!=2:
		raise common.UfuncError("ivo_hasword takes exactly two arguments")
	return None

_makeBooleanizer("ivo_hasword", 
	"(to_tsvector('english', %(1)s) @@ plainto_tsquery('english', %(2)s))")


@userFunction("ivo_nocasematch",
	"(value TEXT, pattern TEXT) -> INTEGER",
	"""
	ivo_nocasematch returns 1 if pattern matches value, 0 otherwise.
	pattern is defined as for the SQL LIKE operator, but the
	match is performed case-insensitively.  This function in effect
	provides a surrogate for the ILIKE SQL operator that is missing from
	ADQL.

	On this site, this is actually implemented using python's and SQL's
	LOWER, so for everything except ASCII, your milage will vary.
	""",
	"integer")
def _nocasematch(args):
	if len(args)!=2:
		raise common.UfuncError("ivo_nocasematch takes exactly two arguments")
	if args[1].type=='characterStringLiteral':
		args[1].value = args[1].value.lower()
	else:
		args[1] = "LOWER(%s)"%nodes.flatten(args[1])
	return None

_makeBooleanizer("ivo_nocasematch", "(LOWER(%(1)s) like %(2)s)")


@userFunction("ivo_hashlist_has",
	"(hashlist TEXT, item TEXT) -> INTEGER",
	"""
	The function takes two strings; the first is a list of words not
	containing the hash sign (#), concatenated by hash signs, the second is
	a word not containing the hash sign.  It returns 1 if, compared
	case-insensitively, the second argument is in the list of words coded in
	the first argument.  The behaviour in case the the second
	argument contains a hash sign is unspecified.
	""",
	"integer")
def _hashlist_has(args):
	if len(args)!=2:
		raise common.UfuncError("ivo_haslist_has takes exactly two arguments")
	return None

_makeBooleanizer("ivo_hashlist_has", 
	"lower(%(2)s) = ANY(string_to_array(lower(%(1)s), '#'))")

@userFunction("ivo_interval_overlaps",
	"(l1 NUMERIC, h1 NUMERIC, l2 NUMERIC, h2 NUMERIC) -> INTEGER",
	"""
	The function returns 1 if the interval [l1...h1] overlaps with
	the interval [l2...h2].  For the purposes of this function,
	the case l1=h2 or l2=h1 is treated as overlap.  The function
	returns 0 for non-overlapping intervals.
	""")
def _interval_overlaps(args):
	if len(args)!=4:
		raise common.UfuncError(
			"ivo_interval_overlaps takes exactly four arguments")
	l1, h1, l2, h2 = [nodes.flatten(a) for a in args]
	return "((%s)>=(%s) AND (%s)>=(%s) AND (%s)<=(%s) AND (%s)<=(%s))::INTEGER"%(
		h1, l2, h2, l1,
		l1, h1, l2, h2)


@userFunction("ivo_interval_has",
	"(val NUMERIC, iv INTERVAL) -> INTEGER",
	"""
	The function returns 1 if the interval iv contains val, 0 otherwise.
	The lower limit is always included in iv, behaviour on the upper
	limit is column-specific.
	""")
def _interval_has(args):
	if len(args)!=2:
		raise common.UfuncError(
			"ivo_interval_has takes exactly two arguments")
	return None

_makeBooleanizer("ivo_interval_has", "((%(1)s) <@ (%(2)s))")


@userFunction("gavo_to_mjd",
	"(d TIMESTAMP) -> DOUBLE PRECISION",
	"""
	The function converts a postgres timestamp to modified julian date.
	This is naive; no corrections for timezones, let alone time
	scales or the like are done; you can thus not expect this to be
	good to second-precision unless you are careful in the construction
	of the timestamp.
	""")
def _to_mjd(args):
	if len(args)!=1:
		raise common.UfuncError("gavo_to_mjd takes exactly one timestamp argument")
	return "ts_to_mjd(%s)"%nodes.flatten(args[0])


@userFunction("gavo_to_jd",
	"(d TIMESTAMP) -> DOUBLE PRECISION",
	"""
	The function converts a postgres timestamp to julian date.
	This is naive; no corrections for timezones, let alone time
	scales or the like are done; you can thus not expect this to be
	good to second-precision unless you are careful in the construction
	of the timestamp.
	""")
def _to_jd(args):
	if len(args)!=1:
		raise common.UfuncError("gavo_to_jd takes exactly one timestamp argument")
	return "ts_to_jd(%s)"%nodes.flatten(args[0])


def _get_histogram_ucd(args):
	baseUCD = args[0].fieldInfo.ucd
	if baseUCD:
		return "stat.histogram;"+baseUCD
	else:
		return "stat.histogram"


@userFunction("gavo_histogram",
	"(val REAL, lower REAL, upper REAL, nbins INTEGER) -> INTEGER[]",
	"""
	The aggregate function returns a histogram of val with nbins+2 elements.
	Assuming 0-based arrays, result[0] contains the number of underflows (i.e.,
	val<lower), result[nbins+1] the number of overflows.  Elements 1..nbins
	are the counts in nbins bins of width (upper-lower)/nbins.  Clients
	will have to convert back to physical units using some external 
	communication, there currently is no (meta-) data as lower and upper in
	the TAP response.
	""",
	returntype="integer[]",
	ucd=_get_histogram_ucd)
def _gavo_histogram(args):
	if len(args)!=4:
		raise common.UfuncError(
			"gavo_histogram takes exactly four arguments (the column to aggregate,"
			" a lower and upper limit of values to tabulate, and the number"
			" of bins desired).")
	return None


@userFunction("gavo_ipix",
	"(long REAL, lat REAL) -> BIGINT",
	"""
	gavo_ipix returns the q3c ipix for a long/lat pair (it simply wraps
	the 13c_ang2ipix function).

	This is probably only relevant when you play tricks with indices or
	PPMXL ids.
	""",
	returntype="bigint",
	ucd="pos")
def _gavo_ipix(args):
	if len(args)!=2:
		raise common.UfuncError(
			"gavo_ipix takes exactly two arguments.")
	
	int, lat = [nodes.flatten(a) for a in args]
	return "q3c_ang2ipix(%s, %s)"%(int, lat)


class _TransformNode(nodes.FunctionNode):
	"""a node representing a gavo_transform call.

	This has a proper node mainly because of type inference.
	"""
	name = "gavo_transform"
	_a_trans = None

	def flatten(self):
		return "(%s)%s"%(nodes.flatten(self.args[0]), self.trans)
	
	def addFieldInfos(self, context):
		# if our argument is properly annotated, that annotation applies
		# to us as well (though it doesn't quite, with pos.eq vs. pos.galactic;
		# sucks, but sucks with UCDs).  If it's not, we're basically hosed.
		if hasattr(self.args[0], "fieldInfo"):
			self.fieldInfo = self.args[0].fieldInfo

		else:
			# go for an STC-S string
			self.fieldInfo = fieldinfo.FieldInfo("text", "", "", tainted=True)
			self.fieldInfo.properties["xtype"] = "adql:REGION"
				

@userFunction("gavo_transform",
	"(from_sys TEXT, to_sys TEXT, geo GEOMETRY) -> GEOMETRY",
	"""
	The function transforms ADQL geometries between various reference systems.
	geo can be a POINT, a CIRCLE, or a POLYGON, and the function will return a
	geometry of the same type.  In the current implementation, from_sys and
	to_sys must be literal strings (i.e., they cannot be computed through
	expressions or be taken from database columns).

	All transforms are just simple rotations, which is only a rough 
	approximation to the actual relationships between reference systems
	(in particular between FK4 and ICRS-based ones).  Note that, in particular,
	the epoch is not changed (i.e., no proper motions are applied).

	We currently support the following reference frames: ICRS, FK5 (which
	is treated as ICRS), FK4 (for B1950. without epoch-dependent corrections),
	GALACTIC.  Reference frame names are case-sensitive.
	""",
	returntype="GEOMETRY")
def _gavo_transform(args):
	if len(args)!=3:
		raise common.UfuncError(
			"gavo_transform takes exactly three arguments")

	fromSys = getStringLiteral(args[0], "source reference system")
	toSys = getStringLiteral(args[1], "target reference system")

	try:
		trans = stc.getPGSphereTrafo(fromSys, toSys)
	except stc.STCValueError as msg:
		raise common.UfuncError(
			"Cannot compute transformation between %s and %s: %s"%(
				fromSys, toSys, msg))

	if trans is None:
		return args[2]
	else:
		raise nodes.ReplaceNode(_TransformNode(
			funName="gavo_transform", 
			args=(args[2],),
			trans=trans))


@userFunction("gavo_normal_random",
	"(mu REAL, sigma REAL) -> REAL",
	"""The function returns a random number drawn from a normal distribution
	with mean mu and width sigma.

	Implementation note: Right now, the Gaussian is approximated by
	summing up and scaling ten calls to random.  This, hence, is not
	very precise or fast.  It might work for some use cases, and we
	will provide a better implemenation if this proves inadequate.
	""",
	returntype="real")
def _gavo_normal_random(args):
	if len(args)!=2:
		raise common.UfuncError(
			"gavo_normal_random takes mu, sigma arguments.")
	
	return ("(((random()+random()+random()+random()+random()"
		"+random()+random()+random()+random()+random()-5)*(%s)"
		")+(%s))")%(nodes.flatten(args[1]), nodes.flatten(args[0]))


@userFunction("gavo_mocunion",
	"(moc1 MOC, moc2 MOC) -> MOC",
	"""returns the union of two MOCs.
	""",
	returntype="smoc")
def _gavo_mocunion(args):
	if len(args)!=2:
		raise common.UfuncError(
			"gavo_mocunion only has two arguments.")

	return "%s | %s"%(
		nodes.flatten(args[0]),
		nodes.flatten(args[1]))


@userFunction("gavo_mocintersect",
	"(moc1 MOC, moc2 MOC) -> MOC",
	"""returns the intersection of two MOCs.
	""",
	returntype="smoc")
def _gavo_mocintersect(args):
	if len(args)!=2:
		raise common.UfuncError(
			"gavo_mocintersect only has two arguments.")

	return "%s & %s"%(
		nodes.flatten(args[0]),
		nodes.flatten(args[1]))


@userFunction("ivo_string_agg",
	"(expression TEXT, delimiter TEXT) -> TEXT",
	"""
	An aggregate function returning all values of
	expression within a GROUP contcatenated with delimiter
	""",
	"text")
def _string_agg(args):
	if len(args)!=2:
		raise common.UfuncError("ivo_string_agg takes exactly two arguments")
	return "string_agg(%s, %s)"%(
		nodes.flatten(args[0]), nodes.flatten(args[1]))


@userFunction("ivo_apply_pm",
	"(ra DOUBLE PRECISION, dec DOUBLE PRECISION, pmra DOUBLE PRECISION, pmde DOUBLE PRECISON, epdist DOUBLE PRECISION) -> POINT",
	"""Returns a POINT (in the UNDEFINED reference frame) for the position
	an object at ra/dec with proper motion pmra/pmde has after epdist years.

	positions must be in degrees, PMs in should be in julian years (i.e., proper
	motions are expected in degrees/year).  pmra is assumed to contain 
	cos(delta).

	This function goes through the tangential plane.  Since it does not have
	information on distance and radial velocity, it cannot reconstruct
	the true space motion, and hence its results will degrade over time.
	""",
	returntype="spoint")
def _ivo_apply_pm(args):
	if len(args)!=5:
		raise common.UfuncError(
			"ivo_apply_pm requires exactly ra, dec, pmra, pmdec, epdist.")

	# there's no indexing or anything to gain here; drop through to an
	# implementation in the database.
	return None


# the healpix functions are currently available in a pgsphere fork
# on credativ's github area, and there are debian packages in
# GAVO's repo.  We'd hope there'll be a stable backport in late 2020.

@userFunction("ivo_healpix_index",
	"(order INTEGER, ra DOUBLE PRECISION, dec DOUBLE PRECISION) -> BIGINT",
	"""Returns the index of the (nest) healpix with order containing the 
	spherical point (ra, dec).

	An alternative, 2-argument form 
	
	ivo_healpix_index(order INTEGER, p POINT) -> BIGINT

	is also available.
	""",
	returntype="bigint")
def _ivo_healpix_index(args):
	if len(args)==2:
		return "healpix_nest(%s, %s)"%(
			nodes.flatten(args[0]), nodes.flatten(args[1]))
	elif len(args)==3:
		return "healpix_nest(%s, spoint(RADIANS(%s), RADIANS(%s)))"%(
			nodes.flatten(args[0]), nodes.flatten(args[1]), nodes.flatten(args[2]))
	else:
		raise common.UfuncError("ivo_healpix_index takes either (ra, dec, order)"
			" or (point, order) arguments")


@userFunction("ivo_healpix_center",
	"(hpxOrder INTEGER, hpxIndex BIGINT) -> POINT",
	"""returns a POINT corresponding to the center of the healpix with
	the given index at the given order.
	""",
	returntype="spoint")
def _ivo_healpix_center(args):
	if len(args)!=2:
		raise common.UfuncError("ivo_healpix_center only takes (index, order)"
			" arguments")
	return "center_of_healpix_nest(%s, %s)"%(
		nodes.flatten(args[0]), nodes.flatten(args[1]))


@userFunction("gavo_getauthority",
	"(ivoid TEXT) -> TEXT",
	"""returns the authority part of an ivoid (or, more generally a URI).
	So, ivo://org.gavo.dc/foo/bar#baz becomes org.gavo.dc.

	The behaviour for anything that's not a full URI is undefined.
	""",
	returntype="text")
def _gavo_getauthority(args):
	if len(args)!=1:
		raise common.UfuncError("gavo_getauthority only takes an ivoid"
			" argument.")
	return None


@userFunction("gavo_vocmatch",
	"(vocname TEXT, term TEXT, matchagainst TEXT) -> INTEGER",
	"""returns 1 if matchagainst is term or narrower in the IVOA vocabulary
	vocname, 0 otherwise.

	This is intended for semantic querying.  For instance, 
	gavo_vocmatch('datalink/core', 'calibration', semantics) would be 1
	if semantics is any of calibration, bias, dark, or flat.

	For RDF-flavoured vocabularies (strict trees), term will expand to the
	entire branch rooted in term.  For SKOS-flavoured vocabularies (where
	narrower is not transitive), only directly narrower terms will be included.

	Both the term and the vocabulary name must be string literals (i.e.,
	constants).  matchagainst can be any string-valued expression.
	""",
	returntype="integer")
def _gavo_vocmatch(args):
	if len(args)!=3:
		raise common.UfuncError("gavo_getauthority takes three"
			" arguments.")

	from gavo.protocols import vocabularies
	voc = vocabularies.get_vocabulary(
		getStringLiteral(args[0], "vocabulary name"))
	term = getStringLiteral(args[1], "term")
	if term not in voc["terms"]:
		raise common.UfuncError("'%s' is not a term in the vocabulary %s"%(
			term, voc["uri"]))

	expanded = [args[1]]
	for narrower in voc["terms"][term]["narrower"]:
		expanded.append(
			nodes.CharacterStringLiteral(value=narrower))

	args[1] = "(%s)"%(", ".join(nodes.flatten(n) for n in expanded))
	raise nodes.ReplaceNode(nodes.FunctionNode(
		funName='GAVO_VOCMATCH', 
		args=(args[2], args[1])))

_makeBooleanizer("gavo_vocmatch", 
	"%(1)s IN %(2)s")


class UserFunction(nodes.FunctionNode):
	"""A node processing user defined functions.

	See the userFunction docstring for how ADQL user defined functions
	are defined.
	"""
	type = "userDefinedFunction"

	def _getFunc(self):
		try:
			return UFUNC_REGISTRY[self.funName.upper()]
		except:
			raise common.UfuncError("No such function: %s"%self.funName)

	def _polish(self):
		if self.args:
			self.args = list(self.args)
		self.processedExpression = self._getFunc()(self.args)

	def flatten(self):
		if self.processedExpression is None:
			return nodes.FunctionNode.flatten(self)
		else:
			return self.processedExpression

	def addFieldInfo(self, context):
		ufunc = self._getFunc()

		unit = (ufunc.adqlUDF_unit(self.args) if callable(ufunc.adqlUDF_unit)
			else ufunc.adqlUDF_unit)
		ucd = (ufunc.adqlUDF_ucd(self.args) if callable(ufunc.adqlUDF_ucd)
			else ufunc.adqlUDF_ucd)

		self.fieldInfo = fieldinfo.FieldInfo(
			ufunc.adqlUDF_returntype, 
			unit, 
			ucd)


tree.registerNode(UserFunction)
