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

    did_plot = False
    did_load = False

    def check_no_file():
        if not did_load:
            from warnings import warn
            warn("No file loaded -- is this what you want?")

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

            col0_len = max(len(k) for k, v in items) + 1

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

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

            col0_len = max(len(k) for k, v in items) + 1

            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"] = expr_y
            if options.legend_descr:
                kwargs["label"] = 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 == "table":
            check_no_file()

            if not did_load:
                from warnings import warn
                warn("No file loaded -- is this what you want?")

            print logmgr.quantity_data[args.pop(0)].table
        else:
            # not a known command, interpret as file name
            logmgr.load(cmd)
            did_load = True

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

        xmin, xmax, ymin, ymax = axis()
        if options.scale_x:
            xmin, xmax = [float(x) for x in options.scale_x.split(",")]
        if options.scale_y:
            ymin, ymax = [float(x) for x in options.scale_y.split(",")]
        axis((xmin, xmax, ymin, ymax))

        title(options.title)
        show()

if __name__ == "__main__":
    main()
