#!/usr/bin/env python

import os, os.path, sys, time, getopt, decimal, re
import sxpar

""" 
updatePluggingStatus:

Interface with plateDb to calculate if a plugging is done with respect to (s/n)^2.  It
will optionally load the status values into plateDb.

The details of the algorithm can be found in the BOSS note "PlateDbDone-v3" which is 
located in the "doc" directory of idlspec2d.

This code could be so much cleaner, if I had more than two days to spend on it.

Written by Gary Kushner (LBL).  March 2010.

"""

try:
	from platedb.APODatabaseConnection import db
	from platedb.db_connection import engine, metadata, Session
	from platedb.ModelClasses import *
	from sqlalchemy import between
	import sqlalchemy
except:
	print '\nMake sure to "setup platedb" before running !!!'
	print '\n\n'
	raise


####
def usage():
	"""Display usage and exit"""
	
	usageCMD = os.path.basename(sys.argv[0])

	print "DEPRECATED -- Use $PLATEDB_DIR/bin/updatePluggingStatus instead\n"
	print "usage:"
	print "\t%s -t MJD_start-MJD_end -p plugging -d" % usageCMD
	print " "
	print "  -t : Update all the pluggings that had exposures between the give dates."
	print "       -t 55000-55100"
	print "  -p : Update the given plugging (e.g. 3874-55278-01)."
	print "  -d : update plateDb with the new status."
	print "  -v : verbose output"

	sys.exit(1)
	

####
class Config:
	"""Config Info"""
	
	def __init__(self):
		self.cameras  = ['b1', 'b2', 'r1', 'r2']
		self.updateDb = False
		self.plugging = None
		self.mjdStart = None
		self.mjdEnd   = None
		self.verbose  = False
		
	def __str__(self):
		return ("cameras   " + str(self.cameras) + "\n" +
				"UpdateDb: " + str(self.updateDb) + "\n" +
		        "plugging: " + str(self.plugging) + "\n" +
				"mjdStart: " + str(self.mjdStart) + "\n" +
				"mjdEnd:   " + str(self.mjdEnd) + "\n" +
				"self.verbose: " + str(self.verbose));


####
def parseCmdLine(args):
	"""Parse command line arguments and return a Config"""

	# parse with options
	try:
		opts, pargs = getopt.gnu_getopt(sys.argv[1:], "t:p:dv")
	except:
		print "Illegal option.\n"
		usage()

	if len(pargs) != 0:
		print "All parameters need to be given using an option switch.\n"
		usage()

	cfg = Config()

	#	Fill in the config
	for (opt, value) in opts:
		if opt == "-t":
			cfg.mjdStart = value.split("-")[0]
			cfg.mjdEnd   = value.split("-")[1]
			if re.match("^5[56]\d\d\d$", cfg.mjdStart) == None:
				print cfg.mjdStart + " is not a valid mjd [55000, 56999].\n"
				usage()
			if re.match("^5[56]\d\d\d$", cfg.mjdEnd) == None:
				print cfg.mjdEnd + " is not a valid mjd [55000, 56999].\n"
				usage()
		if opt == "-p":					# e.g.  3874-55278-01
			cfg.plugging = value
			if re.match("^\d{4}-\d{5}-\d\d$", cfg.plugging) == None:
				print cfg.plugging + " is not a valid plugging."
				print "They look like this 3874-55278-01.\n"
				usage()
		if opt == "-d":
			cfg.updateDb = True
		if opt == "-v":
			cfg.verbose = True

	#	Display config values if verbose
	print "Config values: \n" + str(cfg)
		
	return cfg


####
def findPluggings(session, cfg):
	"""Find all the plugging objects to deal with"""
	
	pluggings = []

	#	if we have an mjd range, find all the pluggings with exposures between the range
	if cfg.mjdStart != None:
		try:
			tai_start = int(cfg.mjdStart) * 86400
			tai_end   = int(cfg.mjdEnd)   * 86400
			if cfg.verbose:
				print "TAI_START:TAI_END: " + str(tai_start) + ":" + str(tai_end)
			pluggings = session.query(Plugging).with_lockmode('update').join(Observation, Exposure).filter(between(Exposure.start_time, tai_start, tai_end)).all()
		except sqlalchemy.orm.exc.NoResultFound:
			print "WARNING:  No pluggings found in given MJD range !"
			pass
		except:
			print "Error looking for pluggings in MJD Range!"
			raise
	
	#	if we have a specific plugging, add that (plPlugMapM-3874-55278-01.par)
	if cfg.plugging != None:
		try:
			p = session.query(Plugging).with_lockmode('update').join(PlPlugMapM).filter(PlPlugMapM.filename=="plPlugMapM-"+cfg.plugging+".par").one()
			if not p in pluggings:
				pluggings.append(p)
		except sqlalchemy.orm.exc.NoResultFound:
			print "Warning:  Could not find plugging: " + cfg.plugging
			pass
		except:
			print "Error looking for plugging " + cfg.plugging
			raise

	return pluggings


####
def getExposures(session, cfg, plugging):
	"""Find all the Exposure objects for the given plugging"""
	
	exposures = []
	
	try:
		exposures = session.query(Exposure).join(Observation).filter(Observation.plugging==plugging).all()
	except sqlalchemy.orm.exc.NoResultFound:
		print "Could not find any Exposures (CameraFrames) for the pluggings!"
		raise
	except:
		print "Error trying to get Exposures (CameraFrames) for the pluggings!"
		raise
			
	return exposures

####
def calculateFlag(session, cfg, plugging, exposures):
	"""Calculate plugging flag using sum(sn2) and other criteria"""
	
	#	Define some constants
	exposureExcellent = 1
	exposureBad       = 2
	exposureTest      = 3
	exposureText      = ["", "Excellent", "Bad", "Test"]
	
	#	Load information that applies to all camera (The values are hard coded to pk)
	try:
		flagAuto           = session.query(PluggingStatus).filter_by(pk=0).one()
		flagGood           = session.query(PluggingStatus).filter_by(pk=1).one()
		flagIncomplete     = session.query(PluggingStatus).filter_by(pk=2).one()
		flagOverGood       = session.query(PluggingStatus).filter_by(pk=3).one()
		flagOverIncomplete = session.query(PluggingStatus).filter_by(pk=4).one()
	except sqlalchemy.orm.exc.MultipleResultsFound:
		print "More than one plugging status value found.  Expecting only one! \n\n"
		raise
	except sqlalchemy.orm.exc.NoResultFound:
		print "Could not load a plugging status flag."
		raise
	except KeyError:	# Why do I sometimes get a key error on record not found?  Ask SQLAlchemy!!
		print "Could not load a plugging status flag. (KeyError)"
		raise
	except:
		print "Problem loading a plugging status flag \n\n"
		raise

	if cfg.verbose:
		print "Initial plugging status is " + plugging.status.label
		
	#	If the plugging is in an overridden state, then there is nothing for us to do
	if plugging.status == flagOverGood or plugging.status == flagOverIncomplete:
		return plugging.status

	#	Loop over all the cameras
	for camName in cfg.cameras:
		if cfg.verbose:
			print "starting camera " + camName
		
		#	First get some variables from the database
		try:
			camera    = session.query(Camera).filter_by(label=camName).one()
			sn2Thresh = session.query(BossSN2Threshold).filter_by(camera=camera).one()
		
			#	Calculate the SN2
			sumsn2        = 0
			goodExposures = 0
			for exposure in exposures:
				if exposure.status.pk != exposureExcellent:
					if cfg.verbose:
						print "exposure " + str(exposure.exposure_no) + " is " + exposureText[exposure.status.pk]
					continue
				else:
					goodExposures += 1

				if cfg.verbose:
					print "exposure " + str(exposure.exposure_no) + " " + camera.label +" sn2 is ",

				try:
					cframe = session.query(CameraFrame).filter_by(exposure=exposure).filter_by(camera=camera).one()
	#				cframe = [x for x in exposure.cameraFrames if x.camera == camera][0]
					sn2 = cframe.sn2
					if sn2 > sn2Thresh.sn2_min:
						sumsn2 += sn2
					if cfg.verbose:
						print str(sn2) + "; sum is " + str(sumsn2)
				except sqlalchemy.orm.exc.MultipleResultsFound:
					print " "
					print "More than one CameraFrame found.  Expecting only one! \n\n"
					raise
				# Why do I sometimes get a key error on record not found?  Ask SQLAlchemy!!
				except (sqlalchemy.orm.exc.NoResultFound, KeyError):
					print "!WARNING:  Could not get sn2 from platedb"
					pass
				except:
					print " "
					print "Problem loading CameraFrame \n\n"
					raise				
					
			#	If we don't have enough sn2, the plugging is incomplete
			if cfg.verbose:
				print "final sumsn2 is " + str(sumsn2)
			if sumsn2 < sn2Thresh.sn2_threshold:
				return flagIncomplete
				
			#	If we don't have enough exposures, the plugging is incomplete
			if cfg.verbose:
				print "number of exposures: " + str(len(exposures))
				print "number of good exposures: " + str(goodExposures)
			if goodExposures < sn2Thresh.min_exposures:
				return flagIncomplete
		except:
			print "Problem calculating sumsn2.\n"
			raise
			
	#	If we get here, then the plugging is done in all cameras
	return flagGood
				

####
def updateFlag(session, cfg, plugging, flag):
	"""Update plateDb with the flag"""

	try:
		plugging.status = flag
	except:
		print "Error updating plugging flag"
		raise

	
####
def processFlags(session, cfg):
	"""Main Script"""
	
	pluggings = findPluggings(session, cfg)
	print " "
	print "Pluggings Found in plateDb:"
	for p in pluggings:
		print str(p.plate.plate_id) + "-" + str(p.fscan_mjd) + "-" + str(p.fscan_id)
	
	for plugging in pluggings:
		if cfg.verbose:
			print "Starting plugging " + str(plugging.plate.plate_id) + "-" + str(plugging.fscan_mjd) + "-" + str(plugging.fscan_id)
		exposures = getExposures(session, cfg, plugging)
		if cfg.verbose:
			print "Found " + str(len(exposures)) + " exposures"
		
		flag = calculateFlag(session, cfg, plugging, exposures)
		print "Plugging " + str(plugging.plate.plate_id) + "-" + str(plugging.fscan_mjd) + "-" + str(plugging.fscan_id) + " is " + flag.label
		
		if cfg.updateDb:
			print "Updating plateDb with flag"
			updateFlag(session, cfg, plugging, flag)
			
	print "Goodbye!"
		


####
def main(argv):
	"""Handle database connection and call main script"""
	
	session = None
	config  = None
	
	try:
		config  = parseCmdLine(sys.argv[1:])
		session = Session()
		processFlags(session, config)
		session.commit()
	except:
		if session != None:
			session.rollback()
		
		raise
	finally:
		if session != None:
			session.close()
			engine.dispose()
		


#### Start of script

if __name__=='__main__':
	main(sys.argv[1:])