#!/usr/bin/env python
# PURPOSE:	qc_imadisp.py : displays a single 2D image (complete frame and details)
# AUTHOR:	Burkhard Wolff, ESO-DMO
# VERSION:	1.0	-- May 2009, edited from qc_mstdisp.py
#		1.1	-- options --plotsize and --compress added (2009-05-27)
#		1.2	-- symbol overplotting added (2009-07-24)
#		1.3	-- additional information can be printed in figure header (2009-08-07)
#		1.4	-- upgrade to numpy V1.3 and matplotlib V0.99 (2009-10-30/2010-02-05)
#		1.4.1	-- options --xcenter and --ycenter added (2010-02-23)
#		1.4.2	-- overplot symbols also on second subplot (2010-02-26)
#		1.4.3	-- symbol colour als function parameter (2010-03-30)
#		1.5	-- functions get_rawimg and get_proimg introduced; handle NAN in product frames (2014-01-17)
#		1.5.1	-- bug fix with outpng name (2014-01-24)
#		1.5.2	-- remove usage of has_key dictionary attribute (2014-03-12)
#		1.5.3	-- usage of argparse module (2014-04-07)
#		1.5.4	-- function parameters xcentrange/ycentrange overwrite values in configuration (2014-12-17)
#
# PARAMETERS:	-a <AB> [ -e <EXT> ] [ -t <TYPE> ] [ -s <SHAPE> ] 
#		[ --lcut=<LOWER_CUT> ] [ --hcut=<HIGHER_CUT> ] 
#		[ --xsym=<XVALUES> ] [ --ysym=<YVALUES> ]
#		[ --xcenter=<XCENTRANGE> ] [ --ycenter=<YCENTRANGE> ]
#		[ --plotindex<IDX> ] [ --plottag=<TAG> ]
#		[ --plotsize=a4|a5|a6 ] [ --compress=<COMP_FACTOR> ]
#		[ --addinfo=<CHAR_STRING> ]
#
# COMMENTS:	The script can be used standalone. It needs qclib.py but does not 
#		require any further configuration files.
#		The script can also be imported into other Python scripts; then,
#		the function display_image() can be called directly.

_version_ = "1.5.4"

# =====================================================================================
# 0. initialization: import modules
# =====================================================================================

# QC configuration and library
from config_cf import *				# names of configuration files
from qclib import *				# qclib classes, functions, etc.

# general modules
import string					# string handling
import time					# to get plot date
import logging					# for writing info, warning, and error messages

# =====================================================================================
# 0.1 helper functions

def get_rawimg(HDU, ext=0, prescan=0, overscan=0):
	"""returns a raw image array, chopped from pre and overscan"""

	image = HDU[ext].data
	# chop prescan and overscan
	if overscan > 0:
		image = image[:,prescan:-overscan]
	else:
		image = image[:,prescan:]

	return image

def get_proimg(HDU, ext=0):
	"""returns a product image array, NAN replaced by 0"""

	image = HDU[ext].data
	image[numpy.isnan(image)] = 0.0

	return image

# =====================================================================================
# 1. function for displaying 2D images
# =====================================================================================

def display_image(AB, proHDUs, rawHDUs, ext=0, plot_config=[ ('raw', 0, (1,2), (1,), (), (), pylab.cm.hot) ], 
		xcentrange=(), ycentrange=(), symcol='b',
		aspect='equal', add_info='', plot_tag='S', plot_index=0, imcomp=1, fig_size='a4'):

	# =============================================================================
	# 1.1 import general configuration dependend on RAW_TYPE and RAW_MATCH_KEYs in AB

	module_name = get_confname(config_files, AB)
	if module_name == '':
		logging.info('configuration file could not be found. Using default.')
		# define some variables
		overscan = 0
		prescan = 0
		#plot_type = raw_type
		parameters = []
	else:
		# save xcentrange/ycentrange before import
		xcr = xcentrange
		ycr = ycentrange
		# import
		exec 'from ' + module_name + ' import * '
		logging.info('configuration imported from ' + module_name)
		# copy back
		if len(xcr) == 2:
			xcentrange = xcr
		if len(ycr) == 2:
			ycentrange = ycr

	# =============================================================================
	# 1.2 define some variables

	ins_name = os.environ.get('DFO_INSTRUMENT')
	masterfile = AB.content['PROD_ROOT_NAME'] + '_0000.fits'

	# =============================================================================
	# 1.3 create plots

	# create figure
	fs_title, fs_sub, fs_foot = get_fontsizes(fig_size)
	fig = pylab.figure(10, figsize=paper_size(fig_size))

	# loop over plot definitions
	for pdata in plot_config:
	
		fig.clear()
		# header of figure
		# first lines: AB, raw type, date
		plot_type = AB.content['RAW_TYPE']
		fig.text(0.01, 0.99, ins_name + ': ' + plot_type, horizontalalignment='left', 
				verticalalignment='top', fontsize=fs_title)
		fig.text(0.5, 0.99, AB.content['AB_NAME'], horizontalalignment='center', verticalalignment='top', fontsize=fs_title)
		fig.text(0.99, 0.99, AB.content['DATE'], horizontalalignment='right', verticalalignment='top', fontsize=fs_title)

		# parameters from AB or as configured
		paramstr = ''
		if len(parameters) == 0:
			if 'RAW_MATCH_KEY' in AB.content:
				for raw_match_key in AB.content['RAW_MATCH_KEY']:
					paramstr = paramstr + ' ' + string.split(string.replace(raw_match_key[0], '=', ' '))[-1]
		else:
			for param in parameters:
				if param[1] in proHDUs[catg_master][0].header:
					paramstr = paramstr + ' ' + param[0] + ' = ' + str(proHDUs[catg_master][0].header[param[1]])
		
		# second line: additional information and parameters
		if add_info == '':
			if fig_size == 'a6':
				fig.text(0.5, 0.955, paramstr, horizontalalignment='center', verticalalignment='top', 
						fontsize=fs_sub)
			else:
				fig.text(0.5, 0.96, paramstr, horizontalalignment='center', verticalalignment='top',
						fontsize=fs_sub)
		else:
			if fig_size == 'a6':
				fig.text(0.99, 0.955, paramstr, horizontalalignment='right', verticalalignment='top', 
						fontsize=fs_sub)
				fig.text(0.01, 0.955, add_info, horizontalalignment='left', verticalalignment='top',
						fontsize=fs_sub)
			else:
				fig.text(0.99, 0.96, paramstr, horizontalalignment='right', verticalalignment='top', 
						fontsize=fs_sub)
				fig.text(0.01, 0.96, add_info, horizontalalignment='left', verticalalignment='top', 
						fontsize=fs_sub)

	
		# bottom: file name
		if pdata[0] == 'pro':
			filename = proHDUs[pdata[1]][0].header['PIPEFILE']
		else:
			filename = string.split(string.replace(AB.content['RAWFILE'][pdata[1]][0], '/', ' '))[-1]
		fig.text(0.01, 0.01, 'File: '+filename+'[ext='+str(ext)+']', 
				horizontalalignment='left', verticalalignment='bottom', fontsize=fs_foot)
		
		# bottom: figure date
		year, month, day, hour, min, sec = time.localtime()[0:6]
		today = str('%04i-%02i-%02i %02i:%02i:%02i' %(year, month, day, hour, min, sec))
		fig.text(0.99, 0.01, 'created ' + today, horizontalalignment='right', verticalalignment='bottom', fontsize=fs_foot)

		# plotting
		logging.info('plotting '+pdata[0]+' '+str(pdata[1]))
		mainplot = ImagePlot(image_map=pdata[2])

		# first subplot
		if pdata[0] == 'raw':
			image = get_rawimg(rawHDUs[pdata[1]], ext, prescan, overscan)
		else:
			image = get_proimg(proHDUs[pdata[1]], ext)
		nx = image.shape[1]
		ny = image.shape[0]
		if len(xcentrange) != 2:
			xcentrange = (4*nx/10, 6*nx/10)
		if len(ycentrange) != 2:
			ycentrange = (4*ny/10, 6*ny/10)
		if pdata[0] == 'raw':
			title_str = 'raw file #'+str(pdata[1])+' '
			title_str = title_str + rawHDUs[pdata[1]][0].header['HIERARCH ESO DPR CATG'] + ' '
			title_str = title_str + rawHDUs[pdata[1]][0].header['HIERARCH ESO DPR TYPE'] + ' '
			title_str = title_str + rawHDUs[pdata[1]][0].header['HIERARCH ESO DPR TECH']
		else:
			title_str = pdata[1]
		if len(pdata) > 6:
			mainplot.add_image(image, 1, title_str, cuts=pdata[3], compress=imcomp, cmap=pdata[6])
		else:
			mainplot.add_image(image, 1, title_str, cuts=pdata[3], compress=imcomp)
		
		# add symbols
		if len(pdata) >= 6:
			if len(pdata[4]) > 0 and len(pdata[5]) > 0 and len(pdata[4]) < 1000:
				if fig_size == 'a4':
					mainplot.add_symbols(pdata[4], pdata[5], 1, symsize=8, edgecolor=symcol, edgewidth=0.75)
				elif fig_size == 'a5':
					mainplot.add_symbols(pdata[4], pdata[5], 1, symsize=6, edgecolor=symcol, edgewidth=0.5)
				else:
					mainplot.add_symbols(pdata[4], pdata[5], 1, symsize=4, edgecolor=symcol, edgewidth=0.4)

		# second subplot
		if pdata[2][0] == 2 or pdata[2][1] == 2:
			if len(pdata) > 6:
				mainplot.add_image(image, 2, title_str, 
						cuts=pdata[3], xrange=xcentrange, yrange=ycentrange, cmap=pdata[6])
			else:
				mainplot.add_image(image, 2, title_str, 
						cuts=pdata[3], xrange=xcentrange, yrange=ycentrange)
			# add symbols
			if len(pdata) >= 6:
				if len(pdata[4]) > 0 and len(pdata[5]) > 0:
					if fig_size == 'a4':
						mainplot.add_symbols(pdata[4], pdata[5], 2, symsize=8, edgecolor=symcol, edgewidth=0.75)
					elif fig_size == 'a5':
						mainplot.add_symbols(pdata[4], pdata[5], 2, symsize=6, edgecolor=symcol, edgewidth=0.5)
					else:
						mainplot.add_symbols(pdata[4], pdata[5], 2, symsize=4, edgecolor=symcol, edgewidth=0.4)

			mainplot.draw(addsquare=(xcentrange[0], ycentrange[0], xcentrange[1], ycentrange[1]), 
					aspect=aspect, figsize=fig_size)
		else:
			mainplot.draw(aspect=aspect, figsize=fig_size)

		# save figures
		#outpng = '%s_%s%02i.png' % (string.rstrip(masterfile,'.fits'), plot_tag, plot_index)
		#outpng = '%s_%s%02i.png' % (string.rstrip(masterfile,'_0123456789.fits'), plot_tag, plot_index)
		outpng = '%s%s%02i.png' % (string.rstrip(masterfile,'0123456789.fits'), plot_tag, plot_index)
		pylab.savefig(outpng, dpi=150, orientation='portrait')
		logging.info('plot saved as ' + outpng)
		plot_index = plot_index + 1

	return plot_index

# =====================================================================================
# 2. function main(), only called in standalone usage
# =====================================================================================

def main():
	# command line parser
	parser = argparse.ArgumentParser(parents=[basic_parser],
			description='Displays a single 2D image.')

	# additional options
	parser.add_argument('--version', action='version', version='%(prog)s ' + _version_)
	parser.add_argument('-e', '--ext', metavar='EXT', dest='ext', type=int, default=0, help='fits extension')
	parser.add_argument('-t', '--type', metavar='TYPE', dest='type', default='0',
			help='''file to be plotted; 0, 1, 2, -1, etc. for first, second, third, last, etc. raw frame; 
			or PRO.CATG of product [default=0]''')
	parser.add_argument('-s', '--shape', metavar='SHAPE', dest='shape', default='12',
			help='plot matrix: 11, 12, or 21; default is 12')
	parser.add_argument('--lcut', metavar='LCUT', dest='lcut', default='0', help='lower cut value for plots')
	parser.add_argument('--hcut', metavar='HCUT', dest='hcut', default='0', help='higher cut value for plots')
	parser.add_argument('--addinfo', metavar='STRING', dest='info', default='',
			help='additional information string to be printed on figure')
	parser.add_argument('--aspect', metavar='ASPECT', dest='aspect', default='equal',
			help='aspect ratio of plots: equal|auto')
	parser.add_argument('--xsym', metavar='XVALS', dest='xsymbols', default='[]',
			help='''xvalues for symbols to be overplotted;
			specified as Python list, e.g. --xsym="[5,6,2]"''')
	parser.add_argument('--ysym', metavar='YVALS', dest='ysymbols', default='[]',
			help='''yvalues for symbols that are overplotted;
			specified as Python list, e.g. --ysym="[1,-4,7]"''')
	parser.add_argument('--xcenter', metavar='XCENTRANGE', dest='xcenter', default='()',
			help='''range of x values for central plot;
			specified as Python list or tuple, e.g. --xcenter="(500,600)"''')
	parser.add_argument('--ycenter', metavar='YCENTRANGE', dest='ycenter', default='()',
			help='''range of y values for central plot;
			specified as Python list or tuple, e.g. --ycenter="(200,300)"''')
	parser.add_argument('--plotindex', metavar='IDX', dest='plot_index', type=int, default=0,
			help='''plot index used for png output; 
			output is written to <PROD_ROOT_NAME>_OOOO_<TAG><IDX>.png''')
	parser.add_argument('--plottag', metavar='TAG', dest='plot_tag', default='S',
			help='''plot tag used for png output; 
			output is written to <PROD_ROOT_NAME>_OOOO_<TAG><IDX>.png''')
	parser.add_argument('--plotsize', metavar='SIZE', dest='plot_size', default='a4',
			help='plot size, either a4, a5, or a6 [optional]')
	parser.add_argument('--compress', metavar='COMP', dest='compress', type=int, default=1,
			help='compression factor for images, useful to save execution time for large images [default=1]')

	# parse arguments/options
	args = parser.parse_args()

	# convert symbol lists (that are char strings) into Python lists
	exec 'xsymbols = ' + args.xsymbols
	exec 'ysymbols = ' + args.ysymbols

	# convert x and y center range into Python lists/tuples
	exec 'xcentrange = ' + args.xcenter
	exec 'ycentrange = ' + args.ycenter

	# configure plot
	nx = int(args.shape[0])
	ny = int(args.shape[1])
	if args.lcut == args.hcut:
		cuts = (1,) # automatic determination of cut values
	else:
		cuts = (float(args.lcut), float(args.hcut))
	if args.type[0] in ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '-']:
		plot_config = [ ('raw', int(args.type), (nx,ny), cuts, xsymbols, ysymbols) ]
	else:
		plot_config = [ ('pro', args.type, (nx,ny), cuts, xsymbols, ysymbols) ]

	compress = int(args.compress)
	if compress < 1:
		compress = 1
	
	if args.aspect == 'auto':
		aspect='auto'
	else:
		aspect='equal'
	
	# set logging level
	set_logging()

	logging.info('started')

	# parse AB
	# AB.content : dictionary with string content of AB

	AB = AssociationBlock(args.ab)
	logging.info(args.ab + ' parsed')

	# reading fits files
	# get list with HDUs of raw files

	logging.info('opening raw frames')
	rawHDUs = AB.get_raw()

	# get dictionary with HDUs of all product files

	logging.info('opening product frames')
	proHDUs = AB.get_pro()

	# draw plots
	display_image(AB, proHDUs, rawHDUs, ext=args.ext, plot_config=plot_config, add_info=str(args.info),
			xcentrange=xcentrange, ycentrange=ycentrange,
			aspect=aspect, plot_index=args.plot_index, plot_tag=str(args.plot_tag), 
			imcomp=compress, fig_size=args.plot_size)

	logging.info('finished')

# =====================================================================================
# 3. if standalone call procedure main()
# =====================================================================================

if __name__ == '__main__':
	main()
	sys.exit(0)