#!/usr/bin/env python

# create scatter plot of two keys in an LDAC table

# HISTORY INFORMATION
# ===================
#
# 09.09.2010:
# I made the script more robust against calling it
# with a too old Python version

# standard-library imports:
import sys
import string

# check for Python 2.X with X >= 5:
version = string.split(string.replace(string.split(sys.version)[0],
                                      '+', ''), ".")
# well, Python version 3 just gives us a syntax error at the
# first print statement :-)
if map(int, version) >= [3, 0, 0] or map(int, version) < [2, 5, 0]:
    sys.stderr.write("This script needs Python 2.Y.X (with Y >= 5)\n\n")
    sys.stderr.write("You have Python V%s.%s.%s\n" \
                      % (version[0], version[1], version[2]))
    sys.exit(1)

# non standard-library imports:
try:
    import numpy as np
except ImportError:
    sys.stderr.write("This script needs the Python module 'numpy'!\n")   
    sys.stderr.write("see http://numpy.scipy.org/\n")                        
    sys.exit(1)                                                                 

try:
    import argparse
except ImportError:                                                             
    sys.stderr.write("This script needs the Python module 'argparse'!\n")   
    sys.stderr.write("see http://code.google.com/p/argparse/\n")
    sys.exit(1)                                                                 

try:
    import ldac
except ImportError:                                                             
    sys.stderr.write("This script needs LDAC Python binding 'ldac.py'!\n")   
    sys.stderr.write("see http://sites.google.com/site/applegatearchive/software\n")
    sys.exit(1)                                                                 


# Handle command line arguments
parser = argparse.ArgumentParser(
    formatter_class=argparse.RawDescriptionHelpFormatter,
    epilog="""
DESCRIPTION:
The program creates a scatter plot from two keys (-k) in a Table (-t) 
within an LDAC catalogue (-i). If no output file (-o) is specified the
plot goes interacticely to the 'matplotlib' X11 device. Conditions on
the objects to be plotted can be passed as lower and upper limits on any 
existing key in the LDAC table with (-c).

EXAMPLES:
- ldacplot.py -i 962743p_1C.cat -t LDAC_OBJECTS -k X_IMAGE Y_IMAGE 

  Create a scatter plot of positions; show the plot interactively
  on the screen

- ldacplot.py -i 962743p_1C.cat -t LDAC_OBJECTS -k 'MAG_APER(1)' 'MAG_APER(2)'

  Plot first and second aperture magnitude; show the plot interactively
  on the screen
  
- ldacplot.py -i 962743p_1C.cat -t LDAC_OBJECTS -k X_IMAGE Y_IMAGE \\
   -c MAG_AUTO 17.0 22.0 FLUX_RADIUS 2.0 3.0 -o plot.png

  Create a scatter plot of positions from objects with 17.0 < MAG_AUTO < 22.0
  and 2.0 < FLUX_RADIUS < 3.0; save the plot in file 'plot.png',  

AUTHOR:
Thomas Erben (terben@astro.uni-bonn.de

ldacplot.py is based on the LDAC Python binding by Douglas Applegate
(see http://sites.google.com/site/applegatearchive/software)
"""
)
parser.add_argument('-i', '--input_file', nargs=1,
                    required=True, help='input file name')
parser.add_argument('-t', '--table', nargs=1,
                    required=True, help='input LDAC table name')
parser.add_argument('-k', '--keys', nargs=2,
                    required=True, help='keys to plot')
parser.add_argument('-c', '--condition', nargs='*',
                    help='condition on the keys to plot')
parser.add_argument('-ti', '--title', nargs=1,
                    help='Plot Title (OPTIONAL)')
parser.add_argument('-o', '--output_file', nargs=1,
                    help='output file (OPTIONAL)')
parser.add_argument('-l', '--limits', nargs=4,
                    type=float, help='plot limits (OPTONAL)')
parser.add_argument('-s', '--marker_style', nargs=1, default=['k.'],
                    help='pyplot marker style (default=\'k.\')')
parser.add_argument('-m', '--marker_size', nargs=1, default=[5],
                    type=int, help='point symbol size (default=5)')

args = parser.parse_args()

# give meaningful variablenames to the command line
# arguments:
catname = args.input_file[0]
tablename = args.table[0]
firstkeyname = args.keys[0]
secondkeyname = args.keys[1]
markerstyle = args.marker_style[0]
markersize = args.marker_size[0]

# handle the conditions:
condkeys = None

if args.condition != None:
    if len(args.condition) % 3 == 0:
        condkeys = args.condition[0::3]
        condmin = args.condition[1::3]
        condmax = args.condition[2::3]
    else:
        sys.stderr.write("Error in condition keys!\n")
        sys.exit(1)

# read input catalogue:
cat = ldac.openObjectFile(catname, table=tablename)

if cat == None:
    sys.stderr.write("ERROR: Cannot open catalog %s with table %s\n" % 
                     (catname, tablename))
    sys.exit(1)

# get input keys
try:
    firstkey = cat[firstkeyname]
    secondkey = cat[secondkeyname]
except KeyError, NameError:
    sys.stderr.write("ERROR: Cannot read key %s and/or key %s\n" % 
                     (firstkeyname, secondkeyname))
    sys.exit(1)

# determine the objects to plot:
condition = (np.isfinite(firstkey)) & (np.isfinite(secondkey))

if condkeys != None:
    for i in xrange(len(condkeys)):
        condition = condition & (cat[condkeys[i]] >= float(condmin[i])) & \
            (cat[condkeys[i]] <= float(condmax[i]))

# do the plotting with matplotlib:

# we need to deal with the import of matplotlib
# AFTER we know whether we need to save the output
# in a file. If we do so we need to disable the
# output device which might not be available in remote
# processing.
try:
    import matplotlib
except ImportError:
    sys.stderr.write("This script needs the modules 'matplotlib'!\n")   
    sys.stderr.write("see http://matplotlib.sourceforge.net/\n")
    sys.exit(1)                                                                 
    
if args.output_file != None:
    matplotlib.use('Agg')
import matplotlib.pyplot as plt

print "Plotting %d objects\n" % (len(firstkey[condition]))
fig = plt.figure()
ax = fig.add_subplot(111)
ax.plot(firstkey[condition], secondkey[condition], markerstyle, 
        markersize=markersize)
ax.set_xlabel(firstkeyname)
ax.set_ylabel(secondkeyname)

if args.limits != None:
    ax.set_xlim(args.limits[0], args.limits[1])
    ax.set_ylim(args.limits[2], args.limits[3])

if args.title != None:
    ax.set_title(args.title[0])

if args.output_file != None:
    print "Saving plot to %s\n" % (args.output_file[0])
    fig.savefig(args.output_file[0])
else:
    plt.show()