#!/usr/bin/env python
"""
Add WCS and observatory headers to raw Königstuhl scans.

Well, there are some plates that were guided on objects, had objective
prisms or are otherwise messy.  For those, we don't add WCS but instead
X-NOCALB ("not calibrated") headers with a short notice why it doesn't
calibrate.
"""



import csv
import math
import os
import resource
import shutil
import sys

from astropy import coordinates, units as u, time
import numpy

from gavo import api
from gavo import utils
from gavo.base import coords
from gavo.helpers import anet
from gavo.helpers import processing
from gavo.utils import fitstools
from gavo.utils import pyfits


def _getHorizonpars(ra, dec, lon, lat, height, jd):
	"""returns airmass, zenith distance and hour angle of the point
	at ra and dec for an observer at lon, lat and height at julian date jd.
	
	All arguments are plain numbers in degrees, except for height, which is
	in meters, and jd, which is in days.
	"""
	def sec(x):
		return 1/math.cos(x)

	obs = coordinates.EarthLocation(lat=lat, lon=lon, height=height*u.m)
	obj = coordinates.SkyCoord(ra, dec, unit="deg")
	transformed = obj.transform_to(
		coordinates.AltAz(obstime=time.Time(jd, format="jd"), location=obs))

	z = math.pi/2-transformed.alt.to(u.deg).value
	airmass = sec(z)*(1-0.0012*(sec(z)**2-1))
	return airmass, z/math.pi*180


def computeHourangle(starTime, alpha):
	"""returns the hour angle in a time-like string.

	starTime is a string containing the sexagesimal star time, alpha is deg.
	"""
	return utils.degToHms(
		utils.hmsToDeg(starTime, sepChar=":")-alpha,
		sepChar=":", secondFracs=0)


_utfMap = [
	('\xc3\xb6', "oe"),
	('\xc3\xa4', "ae"),
	('\xc3\xbc', "ue"),
	('\xc3\xbc', "ue"),
	('\xc3\x84', "Ae"),
]

def clean8Bit(val):
	"""returns the utf-8 byte string val as a plain ASCII str.
	"""
	for src, dst in _utfMap:
		val = val.replace(src, dst)
	return val.encode("ascii").decode("ascii")


PLATE_META_BY_FIRST_CHAR = {
# This dictionary gives plate metadata by the plate character.
# The data was typically given by the scanner project.
# (cf. http://vo.ari.uni-heidelberg.de/arivo/HDAP)
	'A': [
		('SITELAT', "37:13:27"),
		('SITELONG', "-02:32:15", "Positive long is east"),
		('SITEELEV', 2168, "Observatory Elevation [m]"),
		('EMULSION', lambda hdr, rec: rec["Emulsion"]),
		("OBSERVAT", "Calar Alto (493)"),
		("TELESCOP", "Schmidt"),
		("DETECTOR", "Photographic Plate"),],
	'B': [
		('SITELAT', "49:23:55"),
		('SITELONG', "8:43:15", "Positive long is east"),
		('SITEELEV', 560, "Observatory Elevation [m]"),
		('EMULSION', "blue sensitive"),
		('OBSERVAT', "Heidelberg Koenigstuhl (24)"),
		('TELESCOP', "Bruce Astrograph"),
		('DETECTOR', "Photographic Plate"),
		('INSTRUME', lambda hdr, rec: "Camera "+rec["Camera"]),],
	'C': [
		('SITELAT', "37:13:27"),
		('SITELONG', "-02:32:15", "Positive long is east"),
		('SITEELEV', 2168, "Observatory Elevation [m]"),
		('EMULSION', lambda hdr, rec: rec["Emulsion"]),
		("OBSERVAT", "Calar Alto (493)"),
		("TELESCOP", "1.23m in Cassegrain focus w/corrector"),
		("FOCALLEN", 9857, "Focal length [mm]"),
		("DETECTOR", "Photographic Plate"),
		("PLTSCL", 20.9, "Plate scale (arcsec/mm)"),
		],
	'D': [
		('SITELAT', "49:23:55"),
		('SITELONG', "8:43:15", "Positive long is east"),
		('SITEELEV', 560, "Observatory Elevation [m]"),
		('EMULSION', lambda hdr, rec: rec["Emulsion"]),
		("OBSERVAT", "Heidelberg Koenigstuhl (24)"),
		("TELESCOP", "72cm Walz Reflektor"),
		("FOCALLEN", 2815, "Focal length [mm]"),
		("DETECTOR", "Photographic Plate"),
		("PLTSCL", 73.3, "Plate scale (arcsec/mm)"),],
	'E': [
		('SITELAT', "37:13:27"),
		('SITELONG', "-02:32:15", "Positive long is east"),
		('SITEELEV', 2168, "Observatory Elevation [m]"),
		('EMULSION', lambda hdr, rec: rec["Emulsion"]),
		("OBSERVAT", "Calar Alto (493)"),
		("TELESCOP", "2.2m in Cassegrain focus w/corrector"),
		("FOCALLEN", 17037, "Focal length [mm]"),
		("DETECTOR", "Photographic Plate"),
		("PLTSCL", 12.1, "Plate scale (arcsec/mm)"),],
	'F1': [
		('SITELAT', "37:13:27"),
		('SITELONG', "-02:32:15", "Positive long is east"),
		('SITEELEV', 2168, "Observatory Elevation [m]"),
		('EMULSION', lambda hdr, rec: rec["Emulsion"]),
		("OBSERVAT", "Calar Alto (493)"),
		("TELESCOP", "3.5m in primary focus w/3-lens corrector"),
		("FOCALLEN", 17037, "Focal length [mm]"),
		("DETECTOR", "Photographic Plate"),
		("PLTSCL", 15., "Plate scale (arcsec/mm)"),],
	'F2': [
		('SITELAT', "37:13:27"),
		('SITELONG', "-02:32:15", "Positive long is east"),
		('SITEELEV', 2168, "Observatory Elevation [m]"),
		('EMULSION', lambda hdr, rec: rec["Emulsion"]),
		("OBSERVAT", "Calar Alto (493)"),
		("TELESCOP", "3.5m in primary focus w/2-lens corrector"),
		("FOCALLEN", 12195, "Focal length [mm]"),
		("DETECTOR", "Photographic Plate"),
		("PLTSCL", 16.9, "Plate scale (arcsec/mm)"),],
	'G1': [
		('SITELAT', "49:36:36"),
		('SITELONG', "8:41:38", "Positive long is east"),
		('SITEELEV', 110, "Observatory Elevation [m]"),
		('EMULSION', lambda hdr, rec: rec["Emulsion"]),
		('OBSERVAT', "Max Wolf's residence in Heidelberg, Maerzgasse"),
		('TELESCOP', "Wolf's Doppelastrograph"),
		('DETECTOR', "Photographic Plate"),],
	'G2': [
		('SITELAT', "49:23:55"),
		('SITELONG', "8:43:15", "Positive long is east"),
		('SITEELEV', 560, "Observatory Elevation [m]"),
		('EMULSION', lambda hdr, rec: rec["Emulsion"]),
		('OBSERVAT', "Heidelberg Koenigstuhl (24)"),
		('TELESCOP', "Wolf's Doppelastrograph"),
		('DETECTOR', "Photographic Plate"),],
	'H': [
		('SITELAT', "-29:15:15.43"),
		('SITELONG', "-70:44:04.54", "Positive long is east"),
		('SITEELEV', 2335, "Observatory Elevation [m]"),
		('EMULSION', lambda hdr, rec: rec["Emulsion"]),
		('OBSERVAT', "ESO La Silla"),
		('TELESCOP', "2.2m MPG telescope"),
		('DETECTOR', "Photographic Plate"),
		('FOCALLEN',  "Instrument data lost")],
}


def addMetaForObs(plateNumber, header, record):
	"""returns canned metadata for plateNumber.

	This adds whatever items are necessary to fitsHeader; in particular,
	SITELAT, SITELONG, and SITEELEV are in fitsHeader after the function
	has run.
	"""
	key = plateNumber[0]

	if key=='G':
		# Wolf's telescope: on K"onigstuhl from 1202 onward
		if int(plateNumber[1:])<1202:
			key = "G1"
		else:
			key = "G2"
	elif key=="F":
		# ESO MPG 3.5m; 2-lens vs. 3-lens via size
		if header["NAXIS1"]>18000:  # 3-lens is large
			key = "F1"
		else:
			key = "F2"

	for rec in PLATE_META_BY_FIRST_CHAR[key]:
		try:
			key, value = rec
			comment = None
		except ValueError:
			key, value, comment = rec
		if callable(value):
			header.set(key, value(header, record), comment)
		else:
			header.set(key, value, comment)


def addObsHeaders(fName, header, record):
	"""amends the pyfits header with information from the plate cat entry record.
	"""
	if "X-NOCALB" in header:
		centerOfPlate = None
	else:
		centerOfPlate = coords.getCenterFromWCSFields(header)

	header.set("ORIGIN", "GAVO data center")
	header.set("DATE", record["Scandate"], "Date plate was scanned")
	header.set("OBJECT", clean8Bit(record["Object"]),
		"Special object on plate")
	header.set("DATE-OBS", record["Date"], "Date of observation")
	header.set("UT", record["UT"], "UT at mean epoch")
	header.set("ST", record["ST"], "ST at mean epoch")
	header.set("JD", float(record["JD"]), "JD at mean epoch")
	header.set("TM_START", record["TM_Start"], "UT at start of observation")
	header.set("TM_END", record["TM_End"], "UT at end of observation")
	header.set("EXPTIME", float(record["Exptime"]),
		"Effective exposure time [seconds]")
	header.set("COLOR", record["Colour"])

	addMetaForObs(record["Plate"], header, record)

	if centerOfPlate is not None:
		alpha, delta = centerOfPlate
		airmass, zd = _getHorizonpars(
			centerOfPlate[0], centerOfPlate[1],
			header["SITELAT"], header["SITELONG"], float(header["SITEELEV"]),
			float(record["JD"]))

		header.set("AIRMASS", "%.4f"%airmass, "Airmass at mean epoch")
		header.set("HA", computeHourangle(record["ST"], alpha),
			"Hour angle at mean epoch")
		header.set("ZD", zd, "Zenith distance at mean epoch")
		header.set("RA", alpha, "approximate center of plate J2000.0")
		header.set("DEC", delta, "approximate center of plate J2000.0")

	header.set("FILTER", record["Filter"])
	header.set("OBSERVER", clean8Bit(record["Observer"]))
	header.set("PLATE-ID", record["Plate"], "in observation catalogue")
	header.set("IMTITLE", str(os.path.basename(fName)), "source file name")
	header.set("REMARKS", record["Remark"])
	header.add_comment('If you use this image, please acknowledge')
	header.add_comment('"This work made use of the HDAP which was produced at')
	header.add_comment('Landessternwarte Heidelberg-K&ouml;nigstuhl under')
	header.add_comment('grant No. 00.071.2005 of the Klaus-Tschira-Foundation."')
	header.add_comment('Thanks.')
	header.add_history('Calibration arguments: %s'%" ".join(sys.argv[1:]))


def loadNocalibReasons(rd):
	"""returns a dictionary mapping product keys to reasons for not calibrating
	them.

	See comments in res/noposition.txt
	"""
	reasons = {}
	curReason = "unspecified"
	with open(os.path.join(rd.resdir, "res", "noposition.txt")) as f:
		for num, ln in utils.iterSimpleText(f):
			if ln.startswith(">"):
				curReason = ln[1:].strip()
			else:
				reasons[ln] = curReason
	return reasons


def loadWolfPalisa(rd):
	"""returns a mapping from B-plate ids to Wolf-Palisa ids.

	See comments in res/wolf-palisa_mapping.txt
	"""
	res = {}
	with open(os.path.join(rd.resdir, "res", "wolf-palisa_mapping.txt")) as f:
		for num, ln in utils.iterSimpleText(f):
			parts = ln.split()
			res[parts[0]] = "B%s%s"%(parts[9], parts[10])
	return res


class HDAPProcessor(processing.AnetHeaderProcessor):
	sourceExtractorControl_in = """
DETECT_TYPE      CCD
DETECT_MINAREA   %f
DETECT_THRESH    %f
SEEING_FWHM      1.2
"""

	sp_indices = ["index-4113.fits", "index-4112.fits", "index-4111.fits",
		"index-4110.fits", "index-4109.fits", "index-4108.fits"]
	sp_lower_pix = 0.6
	sp_upper_pix = 1.1
	sp_endob = 60
	sp_tweak = True
	sp_plots = False
	sp_total_timelimit = 80
	sp_downsample = None

	def commentFilter(self, value):
		# some versions of pyfits added an "=" sign after comment.
		# That messed up our comments, and the startswith things are attempts
		# to clean that up
		return (processing.AnetHeaderProcessor.commentFilter(self, value)
			or value.startswith('"This work')
			or value.startswith("If you use")
			or value.startswith("Landessternwarte")
			or value.startswith("grant No.")
			or value.startswith("Thanks."))
	
	def historyFilter(self, value):
		return (processing.AnetHeaderProcessor.historyFilter(self, value)
			or "Calibration" in value)

	def _isProcessed(self, srcName):
		hdr = self.getPrimaryHeader(srcName)
		return "CD1_1" in hdr or "X-NOCALB" in hdr

	def classify(self, srcName):
		stem = os.path.splitext(os.path.basename(srcName))[0]
		key = self.getProductKey(srcName)
		if stem not in self.plateCat:
			return "Not in plate cat"
		elif key in self.nocalibReasons:
			return "Special plates"
		elif os.path.exists(self._makeCacheName(srcName)):
			return "Solved"
		return "Not solved"

	def makeObjectFilter(self, badBorder=0.1):
		"""returns a function to throw out funny-looking objects and objects
		near the border.
		"""
		def filterObjects(inName):
			if self.opts.suppressFilter:
				return
			shutil.copy(inName, "origData.fits")
			hdulist = pyfits.open(inName)
			data = hdulist[1].data
			width = max(data.field("X_IMAGE")) # XXX TODO: get from image
			height = max(data.field("Y_IMAGE"))
			data = data[data.field("ELONGATION")<self.opts.maxelong]

			if self.opts.minelong is not None:
				data = data[data.field("ELONGATION")>self.opts.minelong]
			data = data[data.field("X_IMAGE")>width*badBorder]
			data = data[data.field("X_IMAGE")<width-width*badBorder]
			data = data[data.field("Y_IMAGE")>height*badBorder]
			data = data[data.field("Y_IMAGE")<height-height*badBorder]

			if self.opts.dropBrightest:
				data.sort(order=["FLUX_AUTO"])
				data = data[self.opts.dropBrightest:]

			hdu = pyfits.BinTableHDU(numpy.array(data))
			hdu.writeto("foo.xyls")
			hdulist.close()
			os.rename("foo.xyls", inName)

		return filterObjects

	def _solveAnet(self, srcName):
		# If we know the plate cannot be solved, give up immediately
		# with a header field giving the reason
		key = self.getProductKey(srcName)
		if  key in self.nocalibReasons:
			hdr = pyfits.Header()
			hdr.set("X-NOCALB", self.nocalibReasons[key])
			return hdr.cards
		else:
			return processing.AnetHeaderProcessor._solveAnet(self, srcName)

	def _mungeHeader(self, srcName, origHeader):
		stem = os.path.splitext(os.path.basename(srcName))[0]
		if not stem in self.plateCat:
			raise KeyError("Not in plate cat: %s"%stem)
		addObsHeaders(srcName, origHeader, self.plateCat[stem])

		if stem in self.wolfPalisaPlates:
			wpRem = "Wolf-Palisa plate %s"%self.wolfPalisaPlates[stem]
			rems = origHeader["REMARKS"]
			if rems:
				origHeader.set("REMARKS", rems+"; "+wpRem)
			else:
				origHeader.set("REMARKS", wpRem)

		return origHeader

	def _createAuxillaries(self, dd):
		self.resdir = dd.rd.resdir

		rdr = csv.reader(
			open(os.path.join(
				dd.rd.resdir,
				"res/plattenkatalog.csv"),
				encoding="utf-8"),
			delimiter=";")

		fieldNames = next(rdr)
		self.plateCat = dict([(d["Plate"], d) for d in
			[dict(list(zip(fieldNames, rec))) for rec in rdr]])
		self.manualPath = os.path.join(dd.rd.resdir, "res", "manual")

		self.nocalibReasons = loadNocalibReasons(dd.rd)

		self.wolfPalisaPlates = loadWolfPalisa(dd.rd)

	_specialIndices = ["/data/gavo/inputs/lswscans/customindex/"+p
		for p in ["custom-*.idx"]]

	def _runAnet(self, srcName):
		manualName = os.path.join(self.manualPath,
			os.path.basename(srcName)+".manual")
		if os.path.exists(manualName):
			with open(manualName) as f:
				return fitstools.parseCards(f.read())

		sp = self.solverParameters
		sp["endob"] = self.opts.endob
		sourceExtractorControl = self.sourceExtractorControl_in%(
			self.opts.minarea, self.opts.detectThresh)
		objectFilter = None

		# Now patch anet control based on plate characteristics
		idChar = os.path.basename(srcName)[0]
		if idChar in ['A', 'B']:
			sp["indices"] = ["/usr/share/astrometry/index-tycho2-*"]
			sourceExtractorControl = (self.sourceExtractorControl_in+"BACKPHOTO_TYPE LOCAL\n")%(
				self.opts.minarea, self.opts.detectThresh)
			sp["total_timelimit"] = 600

		elif idChar=='C':
			sp["lower_pix"] = 0.18
			sp["upper_pix"] = 0.22
			sourceExtractorControl = self.sourceExtractorControl_in%(
				self.opts.minarea*3, self.opts.detectThresh)

		elif idChar=='D':
			sp["lower_pix"] = 0.65
			sp["upper_pix"] = 0.75
			sourceExtractorControl = "BACKPHOTO_TYPE LOCAL"
			sp["indices"] = ["index-4112.fits", "index-4111.fits",
				"index-4110.fits", "index-4109.fits"]
			objectFilter = self.makeObjectFilter()

		elif idChar=='E':
			sp["lower_pix"] = 0.1
			sp["upper_pix"] = 0.14
			sp["indices"] = ["index-4204-*.fits", "index-4203-*.fits",
				"index-4202-*.fits"]
			objectFilter = self.makeObjectFilter()
			sourceExtractorControl = self.sourceExtractorControl_in%(
				self.opts.minarea*5, self.opts.detectThresh)

		elif idChar=='F':
			sp["lower_pix"] = 0.02
			sp["upper_pix"] = 0.19
			sp["indices"] = ["index-4108.fits", "index-4207-*.fits", "index-4206-*.fits", "index-4205-*.fits", "index-4204-*.fits", "index-4203-*.fits", "index-4202-*.fits"]

		elif idChar=='G':
			sp["lower_pix"] = 1
			sp["upper_pix"] = 20
			sp["total_timelimit"] = 120
			objectFilter = self.makeObjectFilter(0.15)
			sp["indices"] = self._specialIndices
			sp["endob"] = 100
			sourceExtractorControl = ("BACKPHOTO_TYPE LOCAL\nDETECT_MINAREA %f\n"
				"DETECT_THRESH %f\n")%(
					self.opts.minarea, self.opts.detectThresh)

		elif idChar=='H':
			sp["lower_pix"] = 0.1
			sp["upper_pix"] = 0.14
			sourceExtractorControl = self.sourceExtractorControl_in%(
				self.opts.minarea*5, self.opts.detectThresh)
			sp["indices"] = self._specialIndices+sp["indices"]

		else:
			sp["lower_pix"] = 0.1
			sp["upper_pix"] = 3
			print("No params for %s plates yet"%idChar)

		return anet.getWCSFieldsFor(srcName,
			sp,
			sourceExtractorControl,
			objectFilter,
			self.opts.copyTo,
			self.indexPath,
			self.opts.beVerbose)


	@staticmethod
	def addOptions(optParser):
		processing.AnetHeaderProcessor.addOptions(optParser)
		optParser.add_option("--minarea", help="sextractor minarea parameter;"
			" you probably should try at least 50 in addition to the defaults",
			action="store", type="int", dest="minarea", default=300)
		optParser.add_option("--maxelong", help="maximal elongation allowed"
			" for extracted objects", action="store", type="float", dest="maxelong",
			default=2)
		optParser.add_option("--minelong", help="minimal elongation allowed"
			" for extracted objects (default None)", action="store",
			type="float", dest="minelong", default=None)
		optParser.add_option("--detectthreshold", help="detection threshold"
			" for sextractor", action="store", type="float", dest="detectThresh",
			default=4)
		optParser.add_option("--no-filter", help="do not filter sextractor"
			" result", action="store_true", dest="suppressFilter", default=False)
		optParser.add_option("--endob", help="maximal object index for blind"
			" solver", action="store", dest="endob", type="int", default=50)
		optParser.add_option("--dropBrightest", help="drop N brightest objects"
			" from sextractor", action="store", dest="dropBrightest",
			type="int", default=None)


if __name__=="__main__":
# Quick hack to guard against sextractor running wild
	resource.setrlimit(resource.RLIMIT_AS, (10000000000, -1))
	processing.procmain(HDAPProcessor, "lswscans/res/positions.rd", "reimport")