#!/usr/bin/env python
# $Id$
#
# 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", dest="prnHL", 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="int",
                      default=0,
                      help="Smooth highlighted PRNs and clock data with specified \
                      window length. Window length is specified in epochs.")

    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 importance 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.")

    (options, args) = parser.parse_args()

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

    debugLevel = options.debugLevel

    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.prnHL:
        if options.prnHL[0] == 'all':
            highlightPrns = range(1,33)
        else:
            highlightPrns = [int(p) for p in options.prnHL]

    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 epochs" % 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=[]

    if options.debugLevel:
        print "Reading ",inputFile.name
    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='sans-serif',
                  fontsize=12,
                  x=-0.01)

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

    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])

    if plotOrds:
        if options.debugLevel:
            print "Plotting ORDs"
        label="All prns"
        for prn in prns:
            if prn in highlightPrns: continue
            if options.wartiness < 3:
                onePrn = ords[:,ords[1]==prn]
                smoothPlot(ax1, onePrn[0], onePrn[2], 'g.', label, options.smoothing)
            if options.wartiness and warts.size > 0:
                oneWart = warts[:,warts[1]==prn]
                smoothPlot(ax1, oneWart[0], oneWart[2], 'r.', label, 0)
            label="none"

        for prn in highlightPrns:
            if options.wartiness < 3:
                onePrn = ords[:,ords[1]==prn]
                smoothPlot(ax1, onePrn[0], onePrn[2], '.', "prn %2d"%prn, options.smoothing)
            if options.wartiness and warts.size > 0:
                oneWart = warts[:,warts[1]==prn]
                smoothPlot(ax1, oneWart[0], oneWart[2], '.', "prn %2dW"%prn, 0)

        if options.wartiness >= 2:
           ylim=strip_lim(ords[2], options.yRange, warts[2])
        else:
           ylim=strip_lim(ords[2], options.yRange)

        if options.clockDelta:
            # Plot the first deritive of the clock...
            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", options.smoothing)
            yl2 = strip_lim(dc, options.yRange)
            ylim = (min(ylim[0],yl2[0]), max(ylim[1],yl2[1]))
        elif len(ocds[0]):
            # otherwise plot the clock deviations
            smoothPlot(ax1, ocds[0], ocds[1], 'b-', "clk res", options.smoothing)
            yl2 = strip_lim(ocds[1], options.yRange)
            ylim = (min(ylim[0],yl2[0]), max(ylim[1],yl2[1]))

        ax1.set_ylim(ylim)

        # 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=1, markerscale=100, loc=(1,0),
                       handlelength=0.5, handletextpad=0.3, labelspacing=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=4)
            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)
    
        pylab.setp(ax1.get_yticklabels(), fontsize=10, family='sans-serif')
        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 "Plotting clocks"
        clockOffset = int(numpy.average(clocks[1]))
        clocks[1] = clocks[1] - clockOffset

        # if only plotting clocks, put string at top. otherwise, put it above
        # clock plot
        clkStringLoc = 0.49
        if not plotOrds : clkStringLoc = 0.95        

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

        smoothPlot(ax2, clocks[0], clocks[1], 'g.', "clock offset", options.smoothing)
        ax2.set_ylim(strip_lim(clocks[1],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:
            y0 -= clockOffset
            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=1, markerscale=100,  loc=(1,0),
                       handlelength=0.5, handletextpad=0.3, labelspacing=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(pylab.DateFormatter("%H:%M\n%j"))
    ax2.xaxis.set_minor_formatter(pylab.NullFormatter())
    ax2.xaxis.set_major_locator(pylab.MaxNLocator(10))
    ax2.xaxis.set_minor_locator(pylab.MaxNLocator(50))

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

    
    # 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, dpi=122)
# 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 strip_lim(y, limit, warts=numpy.array([])):
    d=numpy.append(y, warts)
    if not len(d):
        d=numpy.array((0))
    if not limit or limit == "ignore" or limit == "none":
        lim = (d.min(), d.max())
    elif limit[-1] == "%":
        fraction = float(limit[:-1])/100.0
        drop = int(numpy.floor(0.5 * (1-fraction) * len(d)))
        ds = numpy.sort(d)
        lim = (ds[drop], ds[-(drop+1)])
    elif limit[-1] == "s":
        mult = float(limit[:-1])
        sigma = numpy.std(d) * mult        
        lim = (-sigma , sigma)
    else:
        lim = (-float(limit), float(limit))

    return lim;
#


def smoothPlot(ax, x, y, style="", txt="", window_len=0):
    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 = int(window_len)
    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, markersize=2)
    else:
        ax.plot_date(x, y, style, label=txt, markersize=2)
#


if __name__ == "__main__":
    main()
