#! /usr/bin/env python

"""\
Usage: %prog 1.dat [2.dat ...]

Make a plot from each of the given plot data files.
"""

################
## Command line args:

import os, optparse
op = optparse.OptionParser(usage=__doc__)
op.add_option("-f", "--format", dest="FORMAT", default="PDF",
              help="output format string consisting of desired output formats separated by commas [default=%default]")
op.add_option("-E", "--engine", dest="ENGINE", default="PGF", help="choose rendering engine: 'PGF' = LaTeX PGF plotting, "
              + "'TEX' = TeX text renderer, 'MPL' = matplotlib MathText (fast but very limited)")
op.add_option("-n", "--nproc", dest="NPROC", default=None, type=int, help="number of plotting processes to run in parallel")
op.add_option("--debug", dest="DEBUG", action="store_true", default=False,
              help="run in debug mode with more verbosity and no parallelism")
op.add_option("--quiet", dest="QUIET", action="store_true", default=False,
              help="run in quiet mode with no status output to terminal")
opts, args = op.parse_args()

## Set the verbosity level in response to --debug and --quiet args
opts.VERBOSITY = 1
if opts.DEBUG:
    opts.VERBOSITY = 2
if opts.QUIET:
    opts.VERBOSITY = 0


import numpy as np
import yoda, yoda.plotting

mpl, plt = yoda.plotting.setup_mpl(opts.ENGINE)
COLORS = [plt.cm.jet(i) for i in np.linspace(0.2, 0.8, len(args))]
STYLES = ["-", "--", ":", "-."]


def plot(plotargs):
    idatfile, datfile = plotargs
    numdatfiles = len(args)

    ## Plan for output in (potentially) several different formats
    outfiles = []
    basename = os.path.splitext(datfile)[0]
    formats = opts.FORMAT.upper().split(",")
    if "PDF" in formats:
        outfiles.append(basename+".pdf")
    if "PNG" in formats:
        outfiles.append(basename+".png")
    if "PGF" in formats:
        outfiles.append(basename+".pgf")

    ## Print status update to terminal
    if opts.VERBOSITY > 0:
        outstr = " ".join(outfiles)
        print "Plotting file {f} -> {o} ({i}/{n})".format(f=datfile, o=outstr, i=idatfile+1, n=numdatfiles)

    ## Read PLOT section annotations
    # TODO: Handle regex'd PLOT sections
    plotkeys = yoda.plotting.read_plot_keys(datfile).get('', {})

    ## Load and sort data objects
    aos = yoda.read(datfile)
    hs = sorted([yoda.plotting.NumpyHist(ao) for ao in aos.values()], key=lambda h: h.path)
    # TODO: allow plotting order specification via PlotIndex (-ve = no plot)

    ## Do plotting
    fig, (ax1, ax2) = yoda.plot(hs, plotkeys=plotkeys)
    for of in outfiles:
        fig.savefig(of)


## Do the rendering using a multiprocessing pool (cleaner than threads)
if opts.DEBUG:
    for i_a in enumerate(args):
        plot(i_a)
else:
    import multiprocessing
    nproc = opts.NPROC or multiprocessing.cpu_count()-1 or 1
    pool = multiprocessing.Pool(processes=nproc)
    pool.map(plot, zip(xrange(len(args)), args))
