#! /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.
"""
    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("--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

    legend_prefix = ""

    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

            plot(data_x, data_y, hold=True, **kwargs)

            did_plot = True
        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))
        else:
            # not a known command, interpret as file name
            from os import access, R_OK
            if access(cmd, R_OK):
                logmgr = LogManager(cmd)
            else:
                raise IOError, "file '%s' not found" % cmd

    if did_plot:
        from pylab import show, title, legend, axis, grid
        if options.legend_expr or options.legend_descr:
            legend()

        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)
        show()

if __name__ == "__main__":
    main()