"""
Commands for obtaining information about various things in the data center.
"""

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


# XXX: TODO: we should throw the AnnotationMaker stuff away and
# move the core of annotateDBTable to rsc.dbtable, and rsc.Limits.
# There's now inMemoryTable.getLimits that does a similar thing.


import datetime

from argparse import ArgumentParser

from gavo import api
from gavo import base
from gavo import stc
from gavo import svcs
from gavo import utils
from gavo.protocols import scs


class Interval(object):
	"""an interval class.

	This should probably move to utils as soon as it is really used
	outside of this module; I suppose we should have a native implementation
	of xtype="interval".
	"""
	def __init__(self, min, max):
		self.min, self.max = min, max
	
	def __str__(self):
		return "%g %g"%(self.min, self.max)


class AnnotationMaker(object):
	"""A class for producing column annotations.
	
	An annotation simply is a dictionary with some well-known keys.  They
	are generated from DB queries.  It is this class' responsibility
	to collect the DB query result columns pertaining to a column and
	produce the annotation dictionary from them.

	To make this happen, it is constructed with the column; then, for
	each property queried, addPropertyKey is called.  Finally, addAnnotation
	is called with the DB result row (see annotateDBTable) to actually
	make and attach the dictionary.
	"""
	def __init__(self, column):
		self.column = column
		if not hasattr(self.column, "annotations"):
			self.column.annotations = {}
		self.propDests = {}

	def doesWork(self):
		return self.propDests

	def getOutputFieldFor(self, propName, propFunc, nameMaker):
		"""returns an OutputField that will generate a propName annotation
		from the propFunc function.

		propFunc for now has a %(name)s where the column name must be
		inserted.

		nameMaker is something like a base.VOTNameMaker.
		"""
		destCol = nameMaker.makeName(propName+"_"+self.column.name)
		self.propDests[destCol] = propName
		return api.makeStruct(svcs.OutputField,
			name=destCol, 
			select=propFunc%{"name": self.column.name}, 
			type=self.column.type)

	def annotate(self, resultRow):
		"""builds an annotation of the column form resultRow.

		resultRow is a dictionary containing values for all keys registred
		through addPropertyKey.

		If the column already has an annotation, only the new keys will be
		overwritten.
		"""
		for srcKey, destKey in self.propDests.items():
			self.column.annotations[destKey] = resultRow[srcKey]


def annotateDBTable(td, 
		extended=True, 
		requireValues=False, 
		samplePercent=None,
		acquireColumnMeta=True):
	"""returns the TableDef td with domain annotations for its columns.

	td must be an existing on-Disk table.

	The annotations come in a dictionary-valued attribute on each
	column annotated.  Possible keys include "min", "max", "avg", and "hasnulls".
	Only columns with the appropriate types (i.e., orderable, and, 
	for avg, numeric) are annotated.

	Without extended, only min and max are annotated.  With
	requireValues, only numeric columns that already have a values child
	are annotated.

	samplePercent uses TABLESAMPLE SYSTEM to only look at about this percentage
	of the rows.
	"""
	outputFields, annotators = [], []
	nameMaker = base.VOTNameMaker()

	for col in td:
		annotator = AnnotationMaker(col)

		if col.type in base.ORDERED_TYPES or col.type.startswith("char"):
			if requireValues:
				if (col.type not in base.NUMERIC_TYPES 
						or not col.values
						or col.values.min is None):
					continue

			outputFields.append(annotator.getOutputFieldFor("max",
				"MAX(%(name)s)", nameMaker))
			outputFields.append(annotator.getOutputFieldFor("min",
				"MIN(%(name)s)", nameMaker))
		
		if extended:
			if col.type in base.NUMERIC_TYPES:
				outputFields.append(annotator.getOutputFieldFor("avg",
					"AVG(%(name)s)", nameMaker))
			

			outputFields.append(annotator.getOutputFieldFor("hasnulls",
					"BOOL_OR(%(name)s IS NULL)", nameMaker))

		if annotator.doesWork():
			annotators.append(annotator)

	with base.getTableConn() as conn:
		td.nrows = estimateTableSize(td.getQName(), conn)
		dbtable = api.TableForDef(td, connection=conn)

		if acquireColumnMeta:
			if not hasattr(dbtable, "getTableForQuery"):
				raise api.ReportableError("Table %s cannot be queried."%td.getQName(),
					hint="This is probably because it is an in-memory table.  Add"
					" onDisk='True' to make tables reside in the database.")

			limitsTable = dbtable.getTableForQuery(
				outputFields, "", samplePercent=samplePercent)
			resultRow = limitsTable.rows[0]

		else:
			annotators = []

	for annotator in annotators:
		annotator.annotate(resultRow)


def getSCSCoverageQuery(td, order):
	"""returns a database query for getting a MOC for a table suitable
	for cone search.

	This will return None if no such query can be built.
	"""
	try:
		raCol, decCol = scs.getConeColumns(td)
	except (ValueError, base.NotFoundError):
		return None

	fragments = [
		"SELECT smoc('%d/' || string_agg(format('%%%%s', hpx), ','))"%order,
		"FROM (",
		"  SELECT DISTINCT healpix_nest(%d,"
			" spoint(RADIANS(%s), RADIANS(%s))) AS hpx "%(
				order, str(raCol.name), str(decCol.name)),
		"FROM %s"%td.getQName(),
		"WHERE %s IS NOT NULL AND %s IS NOT NULL"%(
			str(raCol.name), str(decCol.name)),
		"GROUP BY hpx",
		") as q"]
	return "\n".join(fragments)


def getSSAPCoverageQuery(td, order):
	"""returns a database query for getting a MOC for a table using
	one of our standard SSAP mixins.

	This will return None if no such query can be built.
	"""
	mixinsHandled = ["//ssap#hcd", "//ssap#mixc", "//ssap#view"]
	for mixin in mixinsHandled:
		if td.mixesIn(mixin):
			break
	else:
		return None

	fragments = [
		"SELECT smoc('%d/' || string_agg(format('%%%%s', hpx), ','))"%order,
		"FROM (",
		"  SELECT DISTINCT healpix_nest(%d, ssa_location) AS hpx"%order,
		"FROM %s WHERE ssa_location IS NOT NULL GROUP BY hpx"%td.getQName(),
		") as q"]
	return "\n".join(fragments)


def getSIAPCoverageQuery(td, order):
	"""returns a database query for getting a MOC for a table using
	//siap#pgs (i.e., SIAv1)

	TODO pgsphere is missing features to make this really work properly,
	so so far we're just using centers.  Make sure you're using small
	orders here.

	This will return None if no such query can be built.

	For SIAv2, no such thing is available yet; the assumption is that
	for the time being, for all image collections there's always a
	SIAv1 service, too.
	"""
	if not td.mixesIn("//siap#pgs"):
		return None

	fragments = [
		"SELECT smoc('%d/' || string_agg(format('%%%%s', hpx), ','))"%order,
		"FROM (",
		"  SELECT DISTINCT healpix_nest(%d, "
			" spoint(RADIANS(centerAlpha), RADIANS(centerDelta))) AS hpx "%(order),
		"FROM %s WHERE centerAlpha*centerDelta IS NOT NULL"%td.getQName(),
		"GROUP BY hpx",
		") as q"]
	return "\n".join(fragments)


def getObscoreCoverageQuery(td, order):
	"""returns a database query for getting a MOC for tables with obscore
	columns

	TODO pgsphere is missing features to make this really work properly,
	so so far we're just using s_ra and s_dec; we should move to s_region.

	This will return None if no such query can be built.
	"""
	if not ("s_ra" in td and "s_dec" in td):
		return None

	fragments = [
		"SELECT smoc('%d/' || string_agg(format('%%%%s', hpx), ','))"%order,
		"FROM (",
		"  SELECT DISTINCT healpix_nest(%d, "
			" spoint(RADIANS(s_ra), RADIANS(s_dec))) AS hpx "%(order),
		"FROM %s WHERE s_ra*s_dec IS NOT NULL"%td.getQName(),
		"GROUP BY hpx",
		") as q"]
	return "\n".join(fragments)


def getMOCQuery(td, order):
	"""returns a MOC-generating query for a tableDef with standard
	columns.

	(this is a helper for getMOCForStdTable)
	"""
	for generatingFunction in [
			getSIAPCoverageQuery,
			getSSAPCoverageQuery,
			getSCSCoverageQuery,
			getObscoreCoverageQuery,
		]:
		mocQuery = generatingFunction(td, order)
		if mocQuery is not None:
			return mocQuery
	else:
		raise base.ReportableError("Table %s does not have columns DaCHS knows"
			" how to get a coverage from."%td.getFullId())


def getMOCForStdTable(td, order=6):
	"""returns a MOC for a tableDef with one of the standard protocol mixins.

	The function knows about SCS and SSAP for now; protocols are tested
	for in this order.
	"""
	with base.getTableConn() as conn:
		moc = list(conn.query(getMOCQuery(td, order)))[0][0]
	return moc


def _getTimeTransformer(col):
	"""returns a function turning values in col to julian years.

	This is very much a matter of heuristics; we build upon what's happening
	in utils.serializers.
	"""
	if col.type in ["timestamp", "date"]:
		return lambda val: stc.dateTimeToMJD(val)
		
	elif ((col.ucd and "MJD" in col.ucd.upper())
			or col.xtype=="mjd"
			or "mjd" in col.name):
		return utils.identity

	elif col.unit=="yr" or col.unit=="a":
		return lambda val: stc.dateTimeToMJD(stc.jYearToDateTime(val))

	elif col.unit=="d":
		return lambda val: val-stc.JD_MJD

	elif col.unit=="s":
		return lambda val: stc.dateTimeToMJD(datetime.utcfromtimestamp(val))
	
	else:
		raise base.NotImplementedError("Cannot figure out how to get an MJD"
			" from column %s"%col.name)


def _min(s):
	return "MIN(%s)"%s
def _max(s):
	return "MAX(%s)"%s


def getTimeLimitsExprs(td):
	"""returns the names of columns hopefully containing minimal and
	maximal time coverage of each row of a tabe defined by td.

	As required by getScalarLimits, this will also return a function
	that (hopefully) turns the detected columns to julian years,

	This tries a couple of known criteria for columns containing times
	in some order, and the first one matching wins.

	This will raise a NotFoundError if none of our heuristics work.
	"""
	# obscore and friends
	try:
		return (_min(td.getColumnByName("t_min").name), 
			_max(td.getColumnByName("t_max").name),
			_getTimeTransformer(td.getColumnByName("t_min")))
	except base.NotFoundError:
		pass
	
	# SSAP
	try:
		col = td.columns.getColumnByUtype(
			"ssa:Char.TimeAxis.Coverage.Location.Value"
			)
		return _min(col.name), _max(col.name), utils.identity
	except base.NotFoundError:
		pass
	
	# our SIAP mixins
	try:
		col = td.getColumnByName("dateObs"
			)
		return _min(col.name), _max(col.name), utils.identity
	except base.NotFoundError:
		pass
	
	# Any table with appropriate, sufficiently unique UCDs
	try:
		col = td.getColumnByUCD("time.start")
		# we assume time.start and time.end work the same way and one
		# transformer is enough.
		return (_min(col.name),
			_max(td.getColumnByUCD("time.end").name)
			), _getTimeTransformer(col)
	except ValueError:
		pass
	
	for obsUCD in ["time.epoch", "time.epoch;obs"]:
		try:
			col = td.getColumnByUCD(obsUCD)
			return _min(col.name), _max(col.name), _getTimeTransformer(col)
		except ValueError:
			pass
	
	raise base.NotFoundError("temporal coverage", "Columns to figure out", 
		"table "+td.getFullId())


def getSpectralTransformer(sourceUnit):
	"""returns a function transforming a spectral value to joules, as
	requried by VODataService.
	"""
	if sourceUnit=="m":
		return lambda val: base.PLANCK_H*base.LIGHT_C/val
	raise NotImplementedError("Cannot transform from %s to Joule"%sourceUnit)


def getSpectralLimitsExprs(td):
	"""returns the name of columns hopefully containing minimal and
	maximal spectral coverage.

	As transformer function, we currently return the identity, as we're
	only using IVOA standard columns anyway.  Based on unit and ucd,
	we could pretty certainly do better.

	If this doesn't find any, it raise a NotFoundError.
	"""
	# obscore and friends
	try:
		return (_max(td.getColumnByName("em_max").name), 
			_min(td.getColumnByName("em_min").name),
			getSpectralTransformer("m"))
	except base.NotFoundError:
		pass
	
	# SSAP
	try:
		return (_max(td.getColumnByName("ssa_specend").name), 
			_min(td.getColumnByName("ssa_specstart").name),
			getSpectralTransformer("m"))
	except base.NotFoundError:
		pass

	# SIAv1
	try:
		return (_max(td.getColumnByName("bandpassHi").name), 
			_min(td.getColumnByName("bandpassLo").name),
			getSpectralTransformer("m"))
	except base.NotFoundError:
		pass

	raise base.NotFoundError("spectral coverage", "Columns to figure out", 
		"table "+td.getFullId())


def iterScalarLimits(td, columnsDeterminer):
	"""yields Interal instances for time or spectral coverage.

	ColumnsDeterminer is a function td -> (mincol, maxcol, transformer) expected
	to raise a NotFoundError if no appropriate columns can be found.  This is
	either getTimeLimitsExprs or getSpectralLimitsExprs at this
	point.  transformer here is a function val -> val turning what's coming
	back from the database to what's expected by the coverage machinery
	(e.g., MJD -> jYear).

	It's conceivable that at some point we'll give multiple intervals, 
	and hence this is an iterator (that can very well yield nothing for
	a variety of reasons).
	"""
	try:
		minExpr, maxExpr, transformer = columnsDeterminer(td)
	except base.NotFoundError:
		return

	query = "SELECT %s, %s FROM %s"%(
		str(minExpr), str(maxExpr), td.getQName())
	with base.getTableConn() as conn:
		for res in conn.query(query):
			if res[0] is not None and res[1] is not None:
				yield Interval(transformer(res[0]), transformer(res[1]))


_PROP_SEQ = ("min", "avg", "max", "hasnulls")

def printTableInfo(td, samplePercent=None):
	"""tries to obtain various information on the properties of the
	database table described by td.
	"""
	annotateDBTable(td, samplePercent=samplePercent)
	propTable = [("col",)+_PROP_SEQ]
	for col in td:
		row = [col.name]
		for prop in _PROP_SEQ:
			if prop in col.annotations:
				row.append(utils.makeEllipsis(
					utils.safe_str(col.annotations[prop]), 30))
			else:
				row.append("-")
		propTable.append(tuple(row))
	print(utils.formatSimpleTable(propTable))


def estimateTableSize(tableName, connection):
	"""returns an estimate for the size of the table tableName.

	This is precise for tables postgres guesses are small (currently, 1e6
	rows); for larger tables, we round up postgres' estimate.

	The function will return None for non-existing tables.
	"""
	try:
		estsize = base.UnmanagedQuerier(connection).getRowEstimate(tableName)
	except KeyError:
		return None

	if estsize<1e6:
		res = list(connection.query("SELECT COUNT(*) FROM %s"%tableName))
		if not res:
			return None
		return res[0][0]
	
	else:
		return utils.roundO2M(estsize)


def parseCmdline():
	parser = ArgumentParser(
		description="Displays various stats about the table referred to in"
			" the argument.")
	parser.add_argument("tableId", help="Table id (of the form rdId#tableId)")
	parser.add_argument("-m", "--moc-order", type=int, default=None,
		dest="mocOrder",
		help="Also print a MOC giving the coverage at MOC_ORDER (use MOC_ORDER=6"
			" for about 1 deg resolution).",
		metavar="MOC_ORDER")
	parser.add_argument("-s", "--sample-percent", type=float, default=None,
		dest="samplePercent", metavar="P",
		help="Only look at P percent of the table to determine min/max/mean.")
	return parser.parse_args()


def main():
	args = parseCmdline()
	td = api.getReferencedElement(args.tableId, api.TableDef)
	printTableInfo(td, args.samplePercent)

	if args.mocOrder:
		print("Computing MOC at order %s..."%args.mocOrder)
		print(getMOCForStdTable(td, args.mocOrder))
