#!/usr/bin/env python
"""
Produce a batch file for plotting a lifeline from a netlogger log.
Supported output formats are R (eventually!) and gnuplot.
"""
__author__ = "Dan Gunter <dkgunter@lbl.gov>"
__rcsid__ = "$Id: nl_lifeline 802 2008-06-06 18:15:21Z dang $"

import logging
import optparse
import sys
import time
from warnings import warn
from netlogger.util import getLogger
from netlogger.parsers.base import NLFastParser
from netlogger.analysis.lifeline import Lifeline, LifelineSet

# Logging
log = getLogger("netlogger.nl_lifeline")

class Plot:
    """Superclass for plots.
    """
    def __init__(self, ofile, events=None, title='Lifeline plot'):
        self.ofile = ofile
        self.events = events
        self.title = title
        self.ymin, self.ymax = 0, len(self.events) - 1
        self.xmin, self.xmax = None, None
        self.lifelines = [ ]

    def addLifeline(self, lifeline):
        print "@@add",lifeline.times
        self.lifelines.append(lifeline)
        _ = min(filter(None,lifeline.times))
        if self.xmin is None or _ < self.xmin:
            self.xmin = _
        _ = max(filter(None,lifeline.times))
        if self.xmax is None or _ > self.xmax:
            self.xmax = _
        
    def create(self):
        """Subclasses should override this to output the plot.
        """
        pass

    def howto(self):
        """Subclasses should override this to return an example of
        on how to finish creating the plot (as a string).
        """
        return "<plotting-program> %s" % self.ofile.name

class Gnuplot(Plot):
    def __init__(self, ofile, events=None, **kw):
        Plot.__init__(self, ofile, events=events, **kw)
        self.ytics = ','.join(['"%s" %d' % (e, i) 
                               for i, e in enumerate(self.events)])

    def _hdr(self):
        self.xrange = self.xmax - self.xmin
        return """reset
set terminal png
set output "plot.png"
set multiplot
set pointsize 1.5
set nokey
set xrange [0:%(xrange)lf]
set xlabel "time (seconds)"
set yrange [%(ymin)lf:%(ymax)lf]
set ylabel "events"
set noautoscale x
set noautoscale y
set ytics (%(ytics)s)
set grid xtics ytics
set title "%(title)s"
""" % self.__dict__

    def create(self):
        self.ofile.write(self._hdr())
        for lifeline in self.lifelines:
            self.ofile.write("plot '-' with linespoints\n")
            for i, t in enumerate(lifeline.times):
                if t is not None:
                    self.ofile.write("%lf %d\n" % (t - self.xmin, i))
            self.ofile.write("e\n")
        self.ofile.flush()

    def howto(self):
        return "gnuplot %s" % self.ofile.name

class Rplot(Plot):
    def __init__(self, ofile, events=None, **kw):
        Plot.__init__(self, ofile, events=events, **kw)

    def create(self):
        xdata, ydata = [ ], [ ]
        for lifeline in self.lifelines:
            for i, t in enumerate(lifeline.times):
                if t is not None:
                    xdata.append('%lf' % (t - self.xmin,))
                    ydata.append("'%s'" % self.events[i])
            xdata.append('NA')
            ydata.append('NA')
        self.xstr = "c(" + ','.join(xdata) + ")"
        self.ystr = "factor(x = c(%s), levels=c(%s))" % (
            ','.join(ydata), ','.join(["'%s'" % e for e in self.events]))
        self.ofile.write("""
d <- data.frame(x = %(xstr)s, y = %(ystr)s)
g <- xyplot(y ~ x, data=d, xlab='time (seconds)',
            ylab='events', main='%(title)s', type=c('b','g'))
pdf('plot.pdf')
print(g)
dev.off()
""" % self.__dict__)

    def howto(self):
        return "R CMD BATCH %s" % self.ofile.name

def run(infile_names, options):
    # initialize lifeline set
    llset = LifelineSet(events=options.events, lineids=options.lineids,
                        groupids=options.groupids, prefix=options.prefix)    
    # initialize plot
    if options.output_file == sys.stdout.name:
        ofile = sys.stdout
    else:
        ofile = file(options.output_file,'w')
    if options.plot_type == 'gnuplot':
        plot_class = Gnuplot
    elif options.plot_type == 'R':
        plot_class = Rplot
    else:
        raise ValueError("unknown plot type '%s'" % options.plot_type)
    plot = plot_class(ofile, options.events)
    # process input data
    for filename in infile_names:
        if filename == sys.stdin.name:
            infile = sys.stdin
        else:
            infile = file(filename)
        # parse this file
        parser = NLFastParser(infile)
        for record in parser:
            lifeline = llset.process(record)
            if lifeline is not None:
                plot.addLifeline(lifeline)
        for lifeline in llset.lifelines:
            plot.addLifeline(lifeline)
    # write all lifelines
    plot.create()
    # show how-to
    print "To create the final plot, run:"
    print plot.howto()

def main():
    usage = "%prog [options] [files..]"
    desc = ' '.join(__doc__.split('\n')[1:3])
    parser = optparse.OptionParser(usage=usage, version="%prog 0.1", 
                                   description=desc)
    parser.add_option('-e', '--events',
                      action='store', default="", dest="event_str",
                      help="Events in lifeline, as a comma-separated"
                      " list (required)")
    parser.add_option('-g', '--group-id',
                      action='append', default=[], dest='groupids',
                      help="Field used to group lifelines, e.g., to "
                      "color-code them, repeatable (default=guid)")
    parser.add_option('-l', '--line-id',
                      action='append', default=[], dest='lineids',
                      help="Field used to place events "
                      "in the same lifeline, repeatable (default=guid)")
    parser.add_option('-o', '--output', metavar="FILE", 
                      action="store", dest="output_file", default=None,
                      help="output file (default=stdout)")
    parser.add_option('-p', '--prefix', metavar="PREFIX", default='',
                      action="store", dest="prefix",
                      help="strip event prefix PREFIX where found "
                      "(default=%default)")
    parser.add_option('-t', '--type',
                      action="store", type="choice",
                      choices=('gnuplot', 'R'),
                      dest='plot_type', default="gnuplot",
                      help="Output type (default=%default)")
    options, args= parser.parse_args()
    # configuration
    if not options.event_str:
        parser.error("event list (-e/--events) is required")
    options.events = options.event_str.split(',')
    if not options.lineids:
        options.lineids = ('guid',)
    if not options.groupids:
        options.groupids = ('guid',)
    # input files
    if args:
        infile_names = args
    else:
        infile_names = [sys.stdin.name]
    # output file
    if options.output_file is None:
        options.output_file = sys.stdout.name
    # main method
    run(infile_names, options)
    
if __name__ == "__main__": 
    main()
