#! /usr/bin/env python
import pyfits
import glob
import string
import os
import sys
import getopt
from pylab import *

__author__  = "Alan Patterson"
__version__ = "05/15/08D"

def run (*args):
    """Finds all FITS files in the given directory
    Creates plot for each FITS file if a file extension is also given
    
    Usage:
        do fits_plot  [-v] [-s] [-y <ylabel>] [-t star_type_file]
			[-o radec_override_file]  dirname   [plotext]
	
	-v                   - causes informational messages
	-s                   - causes RA and Dec values to appear in
			       sexagesimal representation (hh:mm:ss.sss)
	ylabel               - is the Y axis label for the plots
			       (default is  Flux)
	star_type_file	     - file containing CSV values for (starname,type)
	radec_override_file  - file containing CSV values for override ra
			       and dec position values  (starname, ra, dec)
	dirname              - directory in which FITS files appear
	plotext              - extension for plot files e.g. (ps, png, pdf)
    """
    opts, params = getopt.getopt(args,'t:y:vso:')
    if len(params) not in [1,2]:
        print __doc__
	return sys.exit(1)
    dirname = params[0]
    if len(params) == 2:
        plotext = params[1]
	savflag = 1
        if plotext.lower() not in ['ps','pdf','png']:
            print 'Unrecognized plot file type'
	    return sys.exit(0)
    else:
        plotext = ''
	savflag = 0
    
    verbose = 0
    sexa = 0
    csvfile = None
    ylab = "Flux"
    typfile = None
    
    for opt in opts:
        if opt[0] == '-v':
	    verbose = 1
	if opt[0] == '-s':
	    sexa = 1
	if opt[0] == '-o':
	    csvfile = opt[1]
	if opt[0] == '-y':
	    ylab = opt[1]
	if opt[0] == '-t':
	    typfile = opt[1]
	    
    if verbose:
        if sexa:
	    print 'Output legend has RA and Dec in sexagesimal form'
	else:
	    print 'Output legend has RA and Dec in degrees'
	    
    if not os.path.exists(dirname):
        print 'Directory does not exist!'
	return sys.exit(0)
	
    filetemplate = '*.fits'
    file2 = os.path.join(dirname,filetemplate)
    filelist = glob.glob(file2)
    
    # position correction file
    if csvfile is not None:
        if verbose:
	    print 'Overwriting header RA and Dec data with data from'
	    print '   %s' % csvfile
        csvdata = get_csv_data(csvfile)
        newposdict = ra_dec_table(csvdata)
    else:
        newposdict = {}
	
    # star type file
    if typfile is not None:
        if verbose:
	    print 'Readin star type data from %s' % typfile
        typdata = get_csv_data(typfile)
	typdict = {}
	for typ in typdata:
	    typdict[typ[0]] = typ[1]
    else:
        typdict = {}
    
    for fil in filelist:
        plotfile = string.split(fil,'.')[0] + '.' + plotext
        if not os.path.exists(plotfile):
            if verbose == '1':
                print fil
            x,y,h = get_plot_data(fil)
	    # get pos from fits header
	    posdict = get_header_data(h)
	    # replace RA DEC data 
	    starname = get_starname(fil)
	    if newposdict.has_key(starname):
	        posdict['ra'] = newposdict[starname]['ra']
		posdict['dec'] = newposdict[starname]['dec']
            # setup legend
	    if typdict.has_key(starname):
	        startype = typdict[starname]
	    else:
	        startype = ''
	    leg = starname + ' ' + get_legend(posdict,sexa,startype)
	    # setup titles (xaxis, yaxis, top)
	    titles = [r'$\lambda (\angstrom)$', \
	    	      ylab, \
		      '%s' % starname]
	    plotroutine(x,y,plotfile,savflag,leg,titles)
	else:
	    print 'Output plot file %s already exists' % plotfile
    return

def get_plot_data(fil):
    """Produces the plot data and header from the FITS file
    	fil         - FITS file
	
    Returns
    	x           - Wavelength
	y           - Intensity
	h           - FITS header
    """
    f = pyfits.open(fil)
    h = pyfits.getheader(fil)
    if h.has_key('CDELT1'):
        xscale = h['CDELT1']
        y = f[0].data
    elif h.has_key('CD1_1'):
        xscale = h['CD1_1']
        y = f[0].data[0][0]
    if h.has_key('CRPIX1'):
        xoffset = xscale*(1. - h['CRPIX1'])
    else:
        xoffset = 0
    x0 = h['CRVAL1'] + xoffset
    x = []
    a = x0
    for i in range(len(y)):
        x.append(a)
        a = a + xscale
    f.close()
    return x,y,h

def get_header_data(h):
    """Extract RA, DEC and DATE-OBS data from the FITS header
    Returns a dictionary with keys 'ra', 'dec', 'date'
    Values are  None  if the key does not occur in the FITS header
    """
    #res ={'ra': None, 'dec': None, 'date': None}
    res ={}
    if h.has_key('RA'):
        v = h['RA']
        res['ra'] = get_degree_from_string(v,15.0)
    elif h.has_key('POSTN-RA'):
        v = h['POSTN-RA']
        res['ra'] = get_degree_from_string(v,15.0)
    else:
        print 'No RA or POSTN-RA keyword found'
    if h.has_key('DEC'):
        v = h['DEC']
        res['dec'] = get_degree_from_string(v,1.0)
    elif h.has_key('POSTN-DE'):
        v = h['POSTN-DE']
        res['dec'] = get_degree_from_string(v,1.0)
    else:
        print 'No DEC or POSTN-DE keyword found'
    if h.has_key('DATE-OBS'):
        v = h['DATE-OBS']
        res['date'] = v
    return res
	
def get_degree_from_string(s,c):
    """Converts input string to degree value.  Input may be a float
     already or a sexagesimal string ('dd:mm:ss.sss').  The sexagesimal
     separator may also be a space i.e. ' '
     c is the conversion from hms to float degrees
    """
    try:
	val = float(s)
    except:
        if string.find(s,':') != -1:
	    sep = ':'
	elif string.find(s,' ') != -1:
	    sep = ' '
	else:
	    return None
	z = string.split(s,sep)
	val = float(z[0])
	div = 60.0
	for i in range(1,min(3,len(z))):
	    if val < 0.0:
	        val = val - float(z[i])/div
	    else:
	        val = val + float(z[i])/div
	    div = div*60.0
	val = val*c
    return val
    
def get_starname(fil):
    """Determines star name from file name
    """
    fbase = os.path.basename(fil)
    starname = string.split(fbase,'.')[0].upper()
    return starname

def sexa_string(v):
    """Converts the degree value to sexagesimal representation
    as a string.
    """
    if v < 0.0: 
    	v = abs(v)
	n = -1
    else: 
        n = 1
    hh = int(v) 
    v2 = (v-hh)*60.0
    mm = int(v2)
    ss = (v2-mm)*60.0
    if n == -1:
        hh = -hh
    s = '%02i:%02i:%06.3f' % (hh,mm,ss)
    return s
        
def get_legend(posdict,sexa,startype):
    """
    """
    if posdict.has_key('ra'):
        if sexa:
	    v = posdict['ra']/15.0
	    ra = sexa_string(v)
	else:
            ra = '%8.5f' % posdict['ra']
    else:
        if posdict.has_key('date'):
	    return startype + ' ' + posdict['date']
	else:
	    return startype
    if posdict.has_key('dec'):
        if sexa:
	    dec = sexa_string(posdict['dec'])[:-1]
	else:
            dec = '%8.5f' % posdict['dec']
    else:
        dec = ''
    leg = ra + ' ' + dec + ' (J2000)'
    if posdict.has_key('date'):
        leg = leg + ' ' + startype + '  ' + posdict['date']
    return leg
    
def plotroutine(x,y,fil,savflag,leg,titles):
    """
    """
    figure(1,(11.0,8.5))
    plot(x,y)
    #xlabel(r'$\lambda (\angstrom)$', size=13)
    #ylabel('Normalized flux')
    if titles[0] is not None:
        xlabel(titles[0])
    if titles[1] is not None:
        ylabel(titles[1])
    if titles[2] is not None:
        title(titles[2])
    legend([leg], loc=3)
    grid(True)
    xlim((x[0],x[-1]))
    if savflag:
        savefig(fil,orientation = 'landscape')
        clf()
    return

def get_csv_data(csvfile):
    """Reads data in csv file and splits it on the commas.
    Returns a list (for each record) of lists (each data point in the
     record).
    """
    specdata = []
    f = open(csvfile,'r')
    line = f.readline()
    while len(line)>0:
        z = string.split(line,',')
        specdata.append((z))
        line = f.readline()
    f.close()
    return specdata
    
def ra_dec_table(data):
    dict={}
    for value in data:
        v1 = get_degree_from_string(string.strip(value[1]),15.0)
	v2 = get_degree_from_string(string.strip(value[2]),1.0)
        dict[string.strip(value[0])]={'ra':v1,'dec':v2}
    return dict
    
# to run on commandline
if __name__ == '__main__':
    input1 = None
    if len(sys.argv) > 1:
        apply(run, tuple(sys.argv[1:]))
    else:
        run(None)