#! /usr/bin/python

def main():
    from pytools.log import LogManager
    import sys
    from optparse import OptionParser

    description = """Operate on data gathered during code runs.
FILE is a log saved from a code run. COMMANDS may be one of the
following:
"list" to list the available time-series and constants,
"plot expr_x,expr_y" to plot a graph,
"datafile outfile expr_x,expr_y" to write out a data file.
"table variable" to print the full data table for a time series variable.
"prefix string" to set the legend prefix for all following plot commands.
"warnings" to list the warnings that were issued during the logged run.
"saveplot filename" to save the current plot to a file.
"print" to print the current plot using "lp".
"""
    parser = OptionParser(usage="%prog FILE COMMANDS FILE COMMANDS...",
            description=description)

    parser.add_option("--scale-x", metavar="XMIN,XMAX",
            help="Set the scale of the X axis")
    parser.add_option("--scale-y", metavar="YMIN,YMAX",
            help="Set the scale of the Y axis")
    parser.add_option("--skip", metavar="N", type="int",
            help="Only use every N'th data point", default=1)
    parser.add_option("--units-x", 
            help="Show only units, not descriptions on the X axis",
            action="store_true")
    parser.add_option("--units-y", 
            help="Show only units, not descriptions on the Y axis",
            action="store_true")
    parser.add_option("--grid", 
            help="Show a grid",
            action="store_true")
    parser.add_option("--legend-expr", 
            help="Generate a legend from the expression",
            action="store_true")
    parser.add_option("--legend-descr", 
            help="Generate a legend from the description",
            action="store_true")
    parser.add_option("--title", 
            help="Set the title of a plot",
            default="Log evaluation")
    parser.add_option("--start-step", metavar="STEP", type="int",
            help="Start the plot at this timestep number")
    parser.add_option("--end-step",  metavar="STEP", type="int",
            help="End the plot at this timestep number")
    options, args = parser.parse_args()

    if len(args) < 1:
        parser.print_help()
        sys.exit(1)

    logmgr = None

    did_plot = False
    saveplot_filename = None
    print_plot = False

    legend_prefix = ""

    from pytools import cartesian_product, Record
    class PlotStyle(Record):
        pass

    styles = [
            PlotStyle(dashes=dashes, color=color)
            for dashes, color in cartesian_product(
                [(), (12, 2), (4, 2),  (2,2), (2,8) ],
                ["blue", "green", "red", "magenta", "cyan"],
                )]

    def check_no_file():
        if logmgr is None:
            raise RuntimeError, "no file loaded"

    while args:
        cmd = args.pop(0)
        if cmd == "list":
            check_no_file()

            print "Time series"
            print "-----------"

            items = list(logmgr.quantity_data.iteritems())
            items.sort(lambda a,b: cmp(a[0], b[0]))

            if items:
                col0_len = max(len(k) for k, v in items) + 1
            else:
                col0_len = 0

            for key, qdat in items:
                print "%s\t%s" % (key.ljust(col0_len), qdat.description)

            print
            print "Constants"
            print "---------"
            items = list(logmgr.constants.iteritems())
            items.sort(lambda a,b: cmp(a[0], b[0]))

            if items:
                col0_len = max(len(k) for k, v in items) + 1
            else:
                col0_len = 0

            for key, value in items:
                print "%s\t%s" % (key.ljust(col0_len), str(value))
        elif cmd == "plot":
            check_no_file()

            expr_x, expr_y = args.pop(0).split(",")

            from pylab import xlabel, ylabel, plot
            (data_x, descr_x, unit_x), (data_y, descr_y, unit_y) = \
                    logmgr.get_plot_data(expr_x, expr_y,
                            options.start_step, options.end_step)

            if options.units_x:
                xlabel(unit_x)
            else:
                xlabel("%s [%s]" % (descr_x, unit_x))
            if options.units_y:
                ylabel(unit_y)
            else:
                ylabel("%s [%s]" % (descr_y, unit_y))

            kwargs = {}

            if options.legend_expr:
                kwargs["label"] = legend_prefix+expr_y
            if options.legend_descr:
                kwargs["label"] = legend_prefix+descr_y

            style = styles.pop(0)
            plot(data_x[::options.skip], 
                    data_y[::options.skip], 
                    dashes=style.dashes, color=style.color, 
                    hold=True,  **kwargs)

            did_plot = True
        elif cmd == "warnings":
            check_no_file()
            print logmgr.get_warnings()

        elif cmd == "datafile":
            check_no_file()

            expr_x, expr_y = args.pop(0).split(",")

            logmgr.write_datafile(args.pop(0), expr_x, expr_y)
        elif cmd == "prefix":
            legend_prefix = args.pop(0)
        elif cmd == "table":
            check_no_file()

            print logmgr.get_table(args.pop(0))
        elif cmd == "saveplot":
            saveplot_filename = args.pop(0)
        elif cmd == "print":
            print_plot = True
        else:
            # not a known command, interpret as file name
            from os import access, R_OK
            if access(cmd, R_OK):
                logmgr = LogManager(cmd, "r")
            else:
                raise IOError, "file '%s' not found" % cmd

    if did_plot:
        from pylab import show, title, legend, axis, grid, savefig
        if options.legend_expr or options.legend_descr:
            from matplotlib.font_manager import FontProperties
            legend(pad=0.04, prop=FontProperties(size=8), loc="best",
                    labelsep=0)

        def float_or_none(str):
            if str == "*":
                return None
            else:
                return float(str)

        xmin, xmax, ymin, ymax = axis()
        if options.scale_x:
            xmin_new, xmax_new = [float_or_none(x) for x in options.scale_x.split(",")]
            if xmin_new is not None:
                xmin = xmin_new
            if xmax_new is not None:
                xmax = xmax_new
        if options.scale_y:
            ymin_new, ymax_new = [float_or_none(x) for x in options.scale_y.split(",")]
            if ymin_new is not None:
                ymin = ymin_new
            if ymax_new is not None:
                ymax = ymax_new

        axis((xmin, xmax, ymin, ymax))

        if options.grid:
            grid()

        title(options.title)

        if print_plot:
            from tempfile import gettempprefix, gettempdir
            import os.path
            tmpname = os.path.join(gettempdir(), gettempprefix()+"logtoolprint.ps")
            savefig(tmpname, orientation="landscape", papertype="letter")

            from os import system, unlink
            system("lp %s" % tmpname)

            unlink(tmpname)

        if saveplot_filename:
            savefig(saveplot_filename)

        if not print_plot and not saveplot_filename:
            show()


if __name__ == "__main__":
    main()
