#!/usr/bin/python
import sys, os, numpy, math, pylab, commands
from optparse import OptionParser

usage = "usage: %prog -f <file> [options]"
parser = OptionParser(usage)
parser.add_option("-f", "--filename", dest="filename", default="",
                  help="filename to run pmpar on")
parser.add_option("-a", "--axistype", dest="axistype", default="rd",
                  help="Plot xy: RA=r, Dec=d, Time=t. Default=rd")
parser.add_option("-m", "--nopm", dest="subpm", action="store_true", 
                  default=False, help="Subtract proper motion")
parser.add_option("-p", "--nopx", dest="subpx", action="store_true",
                  default=False, help="Subtract parallax")
parser.add_option("-o", "--outprefix", dest="outprefix", default="",
                  help="Prefix for output file - default is input filename")
parser.add_option("-k", "--keylocation", dest="keylocation", 
                  default="right top", help="Legend location on graph")
parser.add_option("--annotate", dest="annotate", action="store_true",
                  default=False, help="Annotate px and pm on plot")
parser.add_option("-g", "--grayscale", dest="dogray", action="store_true",
                  default=False, help="Make the plot grayscale")
parser.add_option("-s", "--fontsize", dest="fontsize", default="18",
                  help="Font size for axis labels/title - default is 18")

(options, junk) = parser.parse_args()
filename = options.filename
axistype = options.axistype
subpm    = options.subpm
subpx    = options.subpx
outfile  = options.outprefix
keyloc   = options.keylocation
annotate = options.annotate
dogray   = options.dogray
fontsize = options.fontsize
pmparexe = '/nfs/cluster/ska/adeller/v190/bootstrap_pmpar/bootstrap_pmpar/pmpar'
if filename == "" or not os.path.exists(filename):
    print "You must enter a valid filename! Aborting."
    sys.exit(1)
if outfile == "":
    outfile = filename
outfile = outfile + '.plot.ps'
if dogray:
    outfile = outfile + ".gray"

# Run pmpar, extract information
suffix = ""
if subpm:
    if subpx:
        suffix = " -omp"
    else:
        suffix = " -om"
elif subpx:
    suffix = " -op"

pmpar_msgs = commands.getoutput(pmparexe + ' ' + filename)
if not suffix == "":
    os.system(pmparexe + ' ' + filename + suffix)
centrera = "" 
centredec = ""
pulsar = ""
parallax = ""
pmra = ""
pmdec = ""
lines = pmpar_msgs.split('\n')
for line in lines:
    if line == "":
        continue
    keyword = line.split()[0]
    if keyword == 'Name':
       pulsar = line.split()[2]
    elif keyword == 'RA':
       centrera = line.split()[2]
    elif keyword == 'Dec':
       centredec = line.split()[2]
    elif keyword == 'pi':
       parallax = line.split()[2:]
    elif keyword == 'mu_a':
       pmra = line.split()[2:]
    elif keyword == 'mu_d':
       pmdec = line.split()[2:]
if centrera == "" or centredec == "" or pulsar == "" or parallax == "" or \
pmra == "" or pmdec == "":
    print "Couldn't get all details from pmpar output - aborting!!!"
    sys.exit(1)

# Produce all the labels
title = "Motion of pulsar " + pulsar
ralabel = "Relative RA (mas) from " + centrera
declabel = "Relative Declination (mas) from " + centredec
timelabel = "Time (MJD)"
parallaxlabel = "{/Symbol p} = %4.3f +- %4.3f mas" % (float(parallax[0]), 
                                                   float(parallax[2]))
pmralabel="{/Symbol m}_{/Symbol a} = %6.3f +- %4.3f mas/yr"%(float(pmra[0]),
                                                             float(pmra[2]))
pmdeclabel="{/Symbol m}_{/Symbol d} = %6.3f +- %4.3f mas/yr"%(float(pmdec[0]),
                                                              float(pmdec[2]))
axislabels = ["", ""]
tcols = ["0","0"]
ocols = ["0","0"]
oecols = ["",""]
ecols = ["x","y"]
for i in range(2):
    if axistype[i] == 'r':
        axislabels[i] = ralabel
        tcols[i] = "2"
        ocols[i] = "2"
        oecols[i] = ":3"
    elif axistype[i] == 'd':
        axislabels[i] = declabel
        tcols[i] = "3"
        ocols[i] = "4"
        oecols[i] = ":5"
    elif axistype[i] == 't':
        axislabels[i] = timelabel
        tcols[i] = "1"
        ocols[i] = "1"
        ecols[i] = ""
    else:
        print "axistype label " + axistype + " is not valid - aborting!!!"
        sys.exit(1)

# Create a gnuplot script to make the graph, and run it
terminalstr = "postscript eps enhanced color\n"
l1style = "lt 1 lw 3 pt 0"
l2style = "lt 2 lw 3 pt 0"
l3style = "lt 3 lw 3"
if dogray:
    terminalstr = "postscript eps enhanced monochrome\n"
    l1style = "lt 2 lw 2 pt 0"
    l2style = "lt 1 lw 3 pt 0"
    l3style = "lt 3 lw 3"
gnuplotout = open('temp.gnuplot.script', 'w')
gnuplotout.write("set term " + terminalstr)
gnuplotout.write("set style line 1 " + l1style + "\n")
gnuplotout.write("set style line 2 " + l2style + "\n")
gnuplotout.write("set style line 3 " + l3style + "\n")
gnuplotout.write("set output \"" + outfile + "\"\n")
gnuplotout.write("set title \"" + title + "\" font \"Helvetica," \
                  + fontsize + "\"\n")
if annotate:
    if keyloc.split()[0] == 'right':
        pos = 0.75
    else:
        pos = 0.02
    gnuplotout.write("set label \"" + parallaxlabel + \
                     "\" at graph %f, graph 0.56\n" % (pos))
    gnuplotout.write("set label \"" + pmralabel + \
                     "\" at graph %f, graph 0.5\n" % (pos))
    gnuplotout.write("set label \"" + pmdeclabel + \
                     "\" at graph %f, graph 0.44\n" % (pos))
gnuplotout.write("set xlabel \"" + axislabels[0] + "\" font \"Helvetica," \
                  + fontsize + "\"\n")
gnuplotout.write("set ylabel \"" + axislabels[1] + "\" font \"Helvetica," \
                  + fontsize + "\"\n")
gnuplotout.write("set key " + keyloc + "\n")
plotline = "plot \"pmpar_t\" u " + tcols[0] + ":" + tcols[1] + \
           " w l ls 1 title \"Predicted path\""
if axistype == 'rd':
    plotline = plotline + ", \"pmpar_e\" u 6:7 w p ls 3 " + \
               "title \"Predicted epoch point\""
plotline = plotline + ", \"pmpar_e\" u " + ocols[0] + ":" + ocols[1] + \
                      oecols[0] + oecols[1] + " w " + ecols[0] + \
                      ecols[1] + "e ls 2 title \"Position measurements\""
gnuplotout.write(plotline + "\n")
gnuplotout.close()
os.system("gnuplot < temp.gnuplot.script")