#! /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:
 - Overlay option (print all plots from a file, overlaid)
 - Command-line legend specification cf. Rivet cmphistos and mkhtml
 - 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
import subprocess, multiprocessing

parser = argparse.ArgumentParser(usage=__doc__)
parser.add_argument("INPUTFILES", nargs="+",
                    help="yoda input files")
parser.add_argument("-o", "--outputdir", dest="OUTPUTDIR",
                    default="./yoda-plots", help="output directory for plotting output")
parser.add_argument("-f", "--format", action="append", dest="FORMATS", default=[],
                    help="output format string consisting of desired output formats separated by commas [default=PDF]")
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", action="append", dest="MATCH", default=[],
                    help="only use histograms whose path matches this regex")
parser.add_argument("-M", "--unmatch", action="append", dest="UNMATCH", default=[],
                    help="exclude histograms whose path matches this regex")
parser.add_argument("-j", "--jobs", dest="JOBS", 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")


import yoda
from yoda.plotting import script_generator

def executeScript(script):
    if not os.path.isfile(script):
      raise FileNotFoundError("Python script not found!")
    try:
        subprocess.check_call([script], shell=True)
    except:
        print('Unexpected error when executing', script)

## Assemble plotting arguments depending on mode
def generateScripts(args):
    pyScripts = []
    formats = args.FORMATS if args.FORMATS else [ "PDF" ]
    if args.MODE.upper() == "CMP":
        hists = {}; plotContent = { 'histograms' : {}, }
        for yodafile in args.INPUTFILES:
            aos = yoda.read(yodafile, patterns=args.MATCH, unpatterns=args.UNMATCH)
            for aopath, ao in aos.items():

                # TODO find more permanent solution for Counter/Scatter1D
                if ao.type() == 'Counter' or ao.type() == 'Scatter1D': continue
                
                if aopath in hists:
                    hists[aopath]['histograms'][yodafile] = { 'nominal' : ao, }
                else:
                    hists[aopath] = { 'histograms' : { yodafile : { 'nominal' : ao, }, }, }
        for i, (aopath, aos) in enumerate(sorted(hists.items())):
            name = aopath.replace("/", "_").replace(":", "-")
            if name.startswith("_"):
                name = name[1:]
                pyScript = script_generator.process(aos, name, args.OUTPUTDIR, formats)
                pyScripts += [os.path.abspath(pyScript)]
    elif args.MODE.upper() == "FILE":
        for i, yodafile in enumerate(args.INPUTFILES):
            plotContent = { 'histograms' : {}, }
            aos = yoda.read(yodafile, patterns=args.MATCH, unpatterns=args.UNMATCH)
            name = yodafile.split('.')[0]
            for aopath, ao in aos.items():
                if ao.type() == 'Counter' or ao.type() == 'Scatter1D': continue
                plotContent['histograms'][aopath] = { 'nominal' : ao, }
            pyScript = script_generator.process(plotContent, name, args.OUTPUTDIR, formats)
            pyScripts += [os.path.abspath(pyScript)]
    return pyScripts

## Make output directory
def makeOutputDir(outdir):
    if os.path.exists(outdir) and not os.path.realpath(outdir)==os.getcwd():
        import shutil
        shutil.rmtree(outdir)
    try:
        os.makedirs(outdir)
    except:
        print("Error: failed to make new directory '%s'" % outdir)
        sys.exit(1)

# safeguard to prevent child processes importing this main function and 
# therefore recursively starting the Multiprocessing pool
# see https://stackoverflow.com/a/18205006
if __name__ == '__main__':
    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

    # create directory for scripts/plots
    makeOutputDir(args.OUTPUTDIR)
    pyScripts = generateScripts(args)

    ## Distribute the plotting jobs
    sys.stdout.write('Plotting...')
    nproc = args.JOBS or multiprocessing.cpu_count()-1 or 1
    p = multiprocessing.Pool(processes=nproc)
    try:
      import tqdm
      for _ in tqdm.tqdm(p.imap_unordered(executeScript, pyScripts), total=len(pyScripts)):
          pass
    except ImportError:
      for i, _ in enumerate(p.imap_unordered(executeScript, pyScripts), 1):
          sys.stderr.write(f'\rdone {100*i/len(pyScripts):.1f}%')

    sys.stdout.write('\r')
    sys.stdout.write('Plotting... done!\n')
