#!/usr/bin/env python2.5
# $Id: ordPlot 1114 2008-03-11 12:52:53Z ocibu $
#
# A routine to plot the output of the ord apps.
# The above magic line works under osx 10.5 and ubuntu 7.10. Dunno if it is good
# any further...

import sys, string, time, datetime, numpy, matplotlib, pylab, math

def main():
    from optparse import OptionParser
    parser = OptionParser()
    parser.add_option("-d", "--debug", help="Increase the debugLevel",
                      default=0, dest="debugLevel", action="count")

    parser.add_option("-i", help="Input data file, defaults to stdin.",
                      dest="inputFile", type="string", action="store")

    parser.add_option("-t", dest="title", type="string", action="store",
                      help="Specify a title for the plot. \
                      Defaults to the name of the input stream.")

    parser.add_option("-l", "--legend", dest="legend", action="count",
                      help="Include a legend.")

    parser.add_option("-p", "--prn-highlight", dest="prnHighlight", action="append",
                     help="Highlight the indicated prn. Specify all for all prns.")

    parser.add_option("-o", "--ords-only", help="Only plot the ords (types 0 & 1).",
                      dest="ordsOnly", default=0, action="count")

    parser.add_option("-c", "--clocks-only", help="Only plot the clocks.",
                      dest="clocksOnly", default=0, action="count")

    parser.add_option("--clock-delta", help="Plot clock delta instead of clock residual.",
                      dest="clockDelta", default=0, action="count")

    parser.add_option("-s", dest="saveFig", action="store", type="string",
                      help="Save the figure to the indicated file")

    parser.add_option("-y", dest="yRange", action="store", type="string",
                      help="Fix the y range on the ords to be +- this value or include\
                      the percentage data indicated.")
                      
    parser.add_option("--start-time", dest="tStart", action="store",
                      help="Start time. Format as \"YYYY DOY HH:MM:SS.S\" (Note\
                      the trailing decimal place).") 

    parser.add_option("--smoothing", dest="smoothing", action="store", type="float",
                      default=0,
                      help="Smooth hilighted prns and clock data with specified \
                      window length. Window length is specied in seconds.")

    parser.add_option("--end-time", dest="tEnd", action="store",
                      help="End time. Format as \"YYYY DOY HH:MM:SS.S\" (Note\
                      the trailing decimal place).") 

    parser.add_option("-w", "--warts", dest="wartiness", action="count",
                      help="Increase the importants of warts on the plot.\
                      Zero (the default) means don't even plot them. One\
                      means plot them but don't autoscale to show them all\
                      (just show all the ords). Two means autoscale to show\
                      all the warts. Three means only show the warts and\
                      don't show any ords. Only zero and 2 have been\
                      implimented.")                    

    (options, args) = parser.parse_args()

    if (len(args) and options.inputFile == None):
        options.inputFile = args[0]

    inputFile = sys.stdin
    if (options.inputFile):
        inputFile = open(options.inputFile)

    if (options.title == None):
        options.title = inputFile.name

    prns = range(1,33)
    highlightPrns = []
    if options.prnHighlight:
        if options.prnHighlight[0] == 'all':
            highlightPrns = range(1,33)
        else:
            highlightPrns = [int(p) for p in options.prnHighlight]

    if (options.debugLevel):
        print "Processing: %s" % inputFile.name
        print "Debug level: %d" % options.debugLevel
        print "Title: %s" % options.title
        print "Warts: %s"% options.wartiness
        print "Smoothing: %d seconds" % options.smoothing
        print "Y axis:",
        if options.yRange:
            print options.yRange
        else:
            print "auto ranged"
        if options.debugLevel:
            print "highlighting prns:",highlightPrns

    # ------------------------------------------------------------------
    # Here we start reading in the ord file
    ordList=([],[],[],[])      # time, prn, ord, elevation
    wartList=([],[],[],[])     # time, prn, ord, elevation
    clockList=([],[])          # time, offset
    ocdList=([],[])

    rleClockList=[]

    for line in inputFile:
        line = line.strip()
        if options.debugLevel>1:
            print line
        if len(line)==0: continue

        if line[0] == "#": continue
        if line[0] == '>':
            if line[1] == "c":
                words=line.split()
                if len(words) < 9:
                    print "bad rle line"
                else:
                    t0 = parse_time(words[1:4])
                    t1 = parse_time(words[4:7])
                    offset = float(words[7])
                    slope = float(words[9])
                    abdev = float(words[10])
                    rleClockList.append( (t0, t1, offset, slope, abdev) )
            continue

        words=line.split()
        t = parse_time(words[0:3])

        ordType = int(words[3])
        if ordType == 0:
            if len(words) < 7:
                print "bad ord line"
                continue
            
            prn = int(words[4])
            ord = float(words[7])
            elev = float(words[5])
            wart = int(words[8],16)
            if wart==0:
                ordList[0].append(t)
                ordList[1].append(prn)
                ordList[2].append(ord)
                ordList[3].append(elev)
            else:
                wartList[0].append(t)
                wartList[1].append(prn)
                wartList[2].append(ord)
                wartList[3].append(elev)
        elif ordType == 1:
            if len(words) < 2: print "bad clock residual line"
            ocdList[0].append(t)
            ocdList[1].append(float(words[4]))
        elif ordType == 50:
            if len(words) < 5: print "bad clk line"
            clockList[0].append(t)
            clockList[1].append(float(words[4])) #offset

        if options.debugLevel>2 and len(clockList[0]) >= 200: break

    ords = numpy.array(ordList)
    warts = numpy.array(wartList)
    clocks = numpy.array(clockList)
    ocds = numpy.array(ocdList)

    # Since these are now in numpy arrays, delete the source to save some memory
    del ordList, clockList, wartList, ocdList
    # done reading in the ord file
    # ------------------------------------------------------------------
    
    # Now figure out how many axes we need to use
    plotOrds = True
    plotClocks = True

    if clocks.size==0 or options.ordsOnly: plotClocks = False
    if ((clocks.size==0 and options.clockDelta) and ords.size == 0 and warts.size == 0 and ocds.size == 0) or options.clocksOnly: plotOrds = False
    axesCount=0;
    if plotOrds: axesCount+=1
    if plotClocks: axesCount+=1

    if options.debugLevel:
        print "Read %d ords, %d clocks, %d ocds %d warts %d rle" %\
              (len(ords[0]), len(clocks[0]), len(ocds[0]), len(warts[0]),
               len(rleClockList))

    if axesCount == 0:
        print "No data to plot. Exiting"
        sys.exit()

    # A key handler for matplotlib
    def press(event):
        if event.key=='q' or event.key==' ':
            pylab.close()

    # Here we start generating the plots
    fig = pylab.figure()
    pylab.connect('key_press_event', press)
    yprops = dict(rotation=90,
                  horizontalalignment='right',
                  verticalalignment='center',
                  family='monospace',
                  x=-0.01)

    scale_props = dict(horizontalalignment="right",
                       verticalalignment="bottom",
                       size=8, family="sans-serif")

    xMajorFmt=pylab.DateFormatter("%H:%M\n%j")
    xMinorFmt=pylab.NullFormatter()
    xMajorLoc=matplotlib.dates.DayLocator()
    xMinorLoc=matplotlib.dates.HourLocator()

    rExtent=0.89
    if options.legend:
        rExtent=0.82

    if axesCount == 2:
        ax1 = fig.add_axes([0.08, 0.52, rExtent, 0.42])
    elif axesCount == 1:
        ax1 = fig.add_axes([0.08, 0.10, rExtent, 0.82])
    ax1.set_ylim((0,0))

    if plotOrds:
        label="All prns"
        for prn in prns:
            onePrn = ords[:,ords[1]==prn]
            oneWart = warts[:,warts[1]==prn]
            if prn in highlightPrns:
                smoothPlot(ax1, onePrn[0], onePrn[2], ',', "prn %2d"%prn,
                           options.smoothing)
                if options.wartiness:
                    smoothPlot(ax1, oneWart[0], oneWart[2], ',', "prn %2dW"%prn, 0)
            else:
                smoothPlot(ax1, ords[0], ords[2], 'g,', label, 0)
                if options.wartiness:
                    smoothPlot(ax1, oneWart[0], oneWart[2], 'r,', label, 0)
                label="none"

        setLimit(ax1, ords[2], options.yRange)

        if options.clockDelta:
            # Plot the first deritive of the clock...
            ax3=fig.add_axes(ax1.get_position())
            time_diff = 86400 * (clocks[0,1:] - clocks[0,0:-1])
            clock_diff = 1e9 * 3.3356e-9 * (clocks[1,1:] - clocks[1,0:-1])
            dc = 1e3 * clock_diff/time_diff
            smoothPlot(ax1, clocks[0,:-1], dc, 'b-', "clk dif"%prn,
                       options.smoothing, options.yRange)
            setLimit(ax1, dc, options.yRange)
        elif len(ocds[0]):
            # otherwise plot the clock deviations
            ax3=fig.add_axes(ax1.get_position())
            smoothPlot(ax3, ocds[0], ocds[1], 'b-', "clk res"%prn,
                       options.smoothing, options.yRange)
            setLimit(ax3, ocds[1], options.yRange)
            
        # If there are rle clocks, draw a vertical line where each new model
        # starts
        if not options.clockDelta:
            for t0, t1, y0, m, d in rleClockList:
                ax1.axvline(t0, label='_nolegend_')

        if options.legend:
            ax1.legend(numpoints=2, pad=0.1, labelsep = 0, handlelen=0.005,
                       handletextsep=0.01, axespad=0.0, loc=(1,0))
            leg = pylab.gca().get_legend()
            ltext = leg.get_texts()
            llines = leg.get_lines()
            lframe = leg.get_frame()
            lframe.set_facecolor('0.4')
            pylab.setp(ltext, size=8, family="sans-serif")
            pylab.setp(llines, linewidth=2)
            leg.draw_frame(False)
        if not options.clockDelta:
            ax1.set_ylabel('ord (meters)', **yprops)
        else:
            ax1.set_ylabel('ord (meters)\nclock diff (10^-12)', **yprops)
    
        ax1.grid(True)
        pylab.figtext(rExtent+.08, 0.95, "y range %s"%options.yRange,
                      **scale_props)

    # This allows the creation of futher axes that will share the x axis
    # with the first plot.
    axprops = dict()
    axprops['sharex'] = ax1

    if axesCount == 2:
        ax2 = fig.add_axes([0.08, 0.10, rExtent, 0.38], **axprops)
    elif axesCount == 1:
        ax2 = ax1

    if plotClocks:
        if options.debugLevel:
            print "Debiasing clocks"
        clockOffset = int(numpy.average(clocks[1]))
        clocks[1] = clocks[1] - clockOffset

        pylab.figtext(rExtent+0.08, 0.49, 
                      "Clock offset removed: %d m"%clockOffset,
                      **scale_props)

        smoothPlot(ax2, clocks[0], clocks[1], 'g,', "clock offset",
                   options.smoothing, options.yRange)
        ax2.grid(True)
        ax2.set_ylabel('clock (meters)', **yprops)

        # Only plot the linear clock estimate if there is data for it...
        for t0, t1, y0, m, d in rleClockList:
            y1 = y0 + m * (t1 - t0)
            t = numpy.array([t0, t1])
            y = numpy.array([y0, y1])
            ax2.plot_date(t, y, 'b-', linewidth=1, label='_nolegend_')
            yu = y + d
            yl = y - d
            yy = pylab.concatenate( (yu, yl[::-1]) )
            tt = pylab.concatenate( (t, t[::-1]) )
            ax2.fill(tt, yy, facecolor='b', alpha=0.4, label='_nolegend_')
            
        if options.legend:
            ax2.legend(numpoints=2, pad=0.1, labelsep = 0, handlelen=0.005,
                       handletextsep=0.01, axespad=0.0, loc=(1,0))
            leg = pylab.gca().get_legend()
            pylab.setp(leg.get_texts(), size=8, family="sans-serif")
            leg.draw_frame(False)


    ax2.xaxis.set_major_formatter(xMajorFmt)
    xlabels=ax2.get_xticklabels()
    ylabels=ax2.get_yticklabels()
    pylab.setp(xlabels, fontsize=10, family='sans-serif')
    pylab.setp(ylabels, fontsize=10, family='sans-serif')

    ax2.xaxis.set_minor_formatter(xMinorFmt)
    
    # set x axis range
    if options.tStart:
        tMin = parse_time(options.tStart.split()[0:3])
    else:
        if len(ords[0]):
            tMin = min(ords[0])
        else:
            tMin = min(clocks[0])
            
    if options.tEnd:
        tMax = parse_time(options.tEnd.split()[0:3])
    else:
        if len(ords[0]):
            tMax = max(ords[0])
        else:
            tMax = max(clocks[0])
        
    ax2.set_xlim(xmin=tMin, xmax=tMax)

    # Axis labels on the upper plot would be bad since they would be
    # drawn over the upper part of the lower plot
    if axesCount > 1:
        pylab.setp(ax1.get_xticklabels(), visible=False)
    ax1.set_title(options.title)

    if (options.saveFig == None):
        pylab.show()
    else:
       pylab.savefig(options.saveFig)
# end of main


def parse_time(words):
    fsec = float(words[2][8:10])
    ydhms =  words[0]+" "+words[1]+" "+words[2][0:8]
    utime = time.strptime(ydhms, "%Y %j %H:%M:%S")
    dtime = datetime.datetime(utime[0], utime[1], utime[2],
                              utime[3], utime[4], utime[5], int(fsec*1e6))
    t0 = matplotlib.dates.date2num(dtime)
    return t0
# end of parse_time()


def setLimit(ax, y, limit=None):
    if not limit:
        lim = (min(y), max(y))
    elif limit == "ignore":
        lim = ax.get_ylim()
    elif limit[-1] == "%":
        fraction = float(limit[:-1])/100.0
        drop = int(numpy.floor(0.5 * (1-fraction) * len(y)))
        ds = numpy.sort(y)
        lim = (ds[drop], ds[-drop])
    else:
        lim = (-float(limit), float(limit))
    ax.set_ylim(lim)


def smoothPlot(ax, x, y, style="", txt="", window_len=0, limit=None):
    if x.ndim != 1 or y.ndim !=1:
        raise ValueError, "Only accepts rank 1 arrays."

    if x.size<2 or y.size<2:
        return

    n = window_len/(86400 * (x[1] - x[0]))
    if n>3 and n<y.size:
        s=numpy.r_[2*y[0]-y[n:1:-1],y,2*y[-1]-y[-1:-n:-1]]
        w=eval('numpy.bartlett(n)')
        y=numpy.convolve(w/w.sum(),s,mode='same')
        y=y[n-1:-n+1]
        y=y[n-1:-n+1]
        x=x[n-1:-n+1]

    if txt == "none":
        ax.plot_date(x, y, style)
    else:
        ax.plot_date(x, y, style, label=txt)


if __name__ == "__main__":
    main()
