#! /usr/bin/env python

"""\
%(prog)s data1 [data2 ...]

Make a plot from each of the given plot data files (can be .yoda or make-plots .dat).

TODO:
 - Should be able to do default plotting of all plots loaded from a .yoda file
 - Overlay option (print all plots from a file, overlaid)
 - Command-line legend specification cf. Rivet cmphistos and mkhtml
 - Add a "--outdir" option and
 - Add a "collate" option to gather output into a single PDF?
 - Handle regex'd PLOT sections
 - Allow CLI specification of default plotkeys
"""

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

from __future__ import print_function

import sys, os, argparse
parser = argparse.ArgumentParser(usage=__doc__)
parser.add_argument("DATFILES", nargs="+",
                    help=".dat plot-specification files")
parser.add_argument("-f", "--format", dest="FORMAT", default="PDF",
                    help="output format string consisting of desired output formats separated by commas [default=%(default)s]")
parser.add_argument("--mode", dest="MODE", default="CMP",
                    help="mode of plot combination: CMP=compare same histograms across files -> one file per histo path; "
                    + "FILE=overlay all histograms per file arg -> one file per arg  [default=%(default)s]")
parser.add_argument("-m", "--match", dest="MATCH", metavar="PATT", default=None,
                    help="only use histograms whose path matches this regex")
parser.add_argument("-M", "--unmatch", dest="UNMATCH", metavar="PATT", default=None,
                    help="exclude histograms whose path matches this regex")
parser.add_argument("-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)")
# parser.add_argument("-n", "--nproc", dest="NPROC", default=None, type=int,
#               help="number of plotting processes to run in parallel")
parser.add_argument("--debug", dest="DEBUG", action="store_true", default=False,
                    help="run in debug mode with more verbosity and no parallelism")
parser.add_argument("--quiet", dest="QUIET", action="store_true", default=False,
                    help="run in quiet mode with no status output to terminal")
args = parser.parse_args()

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


import yoda
try:
    mpl = yoda.mplinit(args.ENGINE)
except AttributeError:
    sys.stderr.write('YODA plotting not available. Is matplotlib installed?\n')
    sys.exit(1)

import matplotlib
matplotlib.use('Agg')

from matplotlib import cm
COLORS = [cm.jet(i) for i in yoda.linspace(len(args.DATFILES), 0.2, 0.8)]
STYLES = ["-", "--", ":", "-."]


def plot(plotargs):
    i_n, name, hists, plotkeys = plotargs

    ## Plan for output in (potentially) several different formats
    formats = args.FORMAT.lower().split(",")
    outfiles = [name+"."+f for f in formats]

    ## Print status update to terminal
    if args.VERBOSITY > 0:
        outstr = " ".join(outfiles)
        print("Plotting to {o} ({i:d}/{n:d})".format(o=outstr, i=i_n[0]+1, n=i_n[1]))

    ## Do plotting
    # TODO: allow plotting order specification via PlotIndex (-ve = no plot)
    fig, (ax1, ax2) = yoda.plot(hists, **plotkeys)
    for of in outfiles:
        fig.savefig(of)

    import matplotlib.pyplot as plt
    plt.close()


## Assemble plotting arguments depending on mode
plotargs = []
if args.MODE.upper() == "CMP":
    hists, plotkeys = {}, {}
    for datfile in args.DATFILES:
        aos = yoda.read(datfile, patterns=args.MATCH, unpatterns=args.UNMATCH)
        hists.update(aos)
        plotkeys.update(yoda.plotting.read_plot_keys(datfile))
    for i, (aopath, aos) in enumerate(sorted(hists.items())):
        name = aopath.replace("/", "_")
        if name.startswith("_"):
            name = name[1:]
            plotargs.append([(i, len(hists)), name, aos, plotkeys])
elif args.MODE.upper() == "FILE":
    for i, datfile in enumerate(args.DATFILES):
        import os
        aos = yoda.read(datfile, asdict=False, patterns=args.MATCH, unpatterns=args.UNMATCH)
        name = os.path.splitext(os.path.basename(datfile))[0]
        plotkeys = yoda.plotting.read_plot_keys(datfile)
        plotargs.append([(i, len(args.DATFILES)), name, aos, plotkeys])


## Distribute the plotting jobs
# TODO: fix the multiprocessing
# if args.DEBUG:
for pa in plotargs:
    plot(pa)
# else:
#     import multiprocessing
#     nproc = args.NPROC or multiprocessing.cpu_count()-1 or 1
#     print(nproc)
#     pool = multiprocessing.Pool(processes=nproc)
#     pool.map(plot, plotargs)
