#!/usr/bin/python

"""\
%prog [options] <yodafile1> [<yodafile2> <yodafile3>...] [PLOT:Key1=Val1:...]

Make web pages from histogram files written out by Rivet.  You can specify
multiple Monte Carlo YODA files to be compared in the same syntax as for
rivet-cmphistos, i.e. including plotting options.

Reference data, analysis metadata, and plot style information should be found
automatically (if not, set the RIVET_ANALYSIS_PATH or similar variables
appropriately).

Any existing output directory will be overwritten.

ENVIRONMENT:
 * RIVET_ANALYSIS_PATH: list of paths to be searched for plugin
     analysis libraries at runtime
 * RIVET_DATA_PATH: list of paths to be searched for data files
"""

from __future__ import print_function

import rivet, sys, os
rivet.util.check_python_version()
rivet.util.set_process_name(os.path.basename(__file__))

import glob, shutil
from subprocess import Popen,PIPE


from optparse import OptionParser, OptionGroup
parser = OptionParser(usage=__doc__)
parser.add_option("-o", "--outputdir", dest="OUTPUTDIR",
                  default="./rivet-plots", help="directory for Web page output")
parser.add_option("-c", "--config", dest="CONFIGFILES", action="append", default=["~/.make-plots"],
                  help="plot config file(s) to be used with rivet-cmphistos")
parser.add_option("-n", "--num-threads", metavar="NUMTHREADS", dest="NUMTHREADS", type=int,
                  default=None, help="request make-plots to use a specific number of threads")
parser.add_option("--ignore-missing", dest="IGNORE_MISSING", action="store_true",
                  default=False, help="ignore missing YODA files")
parser.add_option("-i", "--ignore-unvalidated", dest="IGNORE_UNVALIDATED", action="store_true",
                  default=False, help="ignore unvalidated analyses")
# parser.add_option("--ref", "--refid", dest="REF_ID",
#                   default=None, help="ID of reference data set (file path for non-REF data)")
parser.add_option("--dry-run", help="don't actually do any plotting or HTML building", dest="DRY_RUN",
                  action="store_true", default=False)
parser.add_option("--no-cleanup", dest="NO_CLEANUP", action="store_true", default=False,
                  help="keep plotting temporary directory")
parser.add_option("--no-subproc", dest="NO_SUBPROC", action="store_true", default=False,
                  help="don't use subprocesses to render the plots in parallel -- useful for debugging")
parser.add_option("--pwd", dest="PATH_PWD", action="store_true", default=False,
                  help="append the current directory (pwd) to the analysis/data search paths (cf. $RIVET_ANALYSIS_PATH)")

stygroup = OptionGroup(parser, "Style options")
stygroup.add_option("-t", "--title", dest="TITLE",
                    default="Plots from Rivet analyses", help="title to be displayed on the main web page")
stygroup.add_option("--reftitle", dest="REFTITLE",
                    default="Data", help="legend entry for reference data")
stygroup.add_option("--no-plottitle", dest="NOPLOTTITLE", action="store_true",
                    default=False, help="don't show the plot title on the plot "
                        "(useful when the plot description should only be given in a caption)")
stygroup.add_option("-s", "--single", dest="SINGLE", action="store_true",
                    default=False, help="display plots on single webpage.")
stygroup.add_option("--no-ratio", dest="SHOW_RATIO", action="store_false",
                    default=True, help="don't draw a ratio plot under each main plot.")
stygroup.add_option("--errs", "--mc-errs", dest="MC_ERRS", action="store_true",
                    default=False, help="plot error bars.")
stygroup.add_option("--offline", dest="OFFLINE", action="store_true",
                    default=False, help="generate HTML that does not use external URLs.")
stygroup.add_option("--pdf", dest="VECTORFORMAT", action="store_const", const="PDF",
                    default="PDF", help="use PDF as the vector plot format.")
stygroup.add_option("--ps", dest="VECTORFORMAT", action="store_const", const="PS",
                    default="PDF", help="use PostScript as the vector plot format. DEPRECATED")
stygroup.add_option("--booklet", dest="BOOKLET", action="store_true",
                    default=False, help="create booklet (currently only available for PDF with pdftk or pdfmerge).")
stygroup.add_option("--font", dest="OUTPUT_FONT", choices="palatino,cm,times,helvetica,minion".split(","),
                    default="palatino", help="choose the font to be used in the plots")
stygroup.add_option("--palatino", dest="OUTPUT_FONT", action="store_const", const="palatino", default="palatino",
                    help="use Palatino as font (default). DEPRECATED: Use --font")
stygroup.add_option("--cm", dest="OUTPUT_FONT", action="store_const", const="cm", default="palatino",
                    help="use Computer Modern as font. DEPRECATED: Use --font")
stygroup.add_option("--times", dest="OUTPUT_FONT", action="store_const", const="times", default="palatino",
                    help="use Times as font. DEPRECATED: Use --font")
stygroup.add_option("--helvetica", dest="OUTPUT_FONT", action="store_const", const="helvetica", default="palatino",
                    help="use Helvetica as font. DEPRECATED: Use --font")
stygroup.add_option("--minion", dest="OUTPUT_FONT", action="store_const", const="minion", default="palatino",
                    help="use Adobe Minion Pro as font. DEPRECATED: Use --font")
parser.add_option_group(stygroup)

selgroup = OptionGroup(parser, "Selective plotting")
selgroup.add_option("-m", "--match", action="append", dest="PATHPATTERNS", default=[],
                    help="only write out histograms whose $path/$name string matches any of these regexes")
selgroup.add_option("-M", "--unmatch", action="append", dest="PATHUNPATTERNS", default=[],
                    help="exclude histograms whose $path/$name string matches any of these regexes")
selgroup.add_option(#"-a", #< these were confusing, and -m should be enough
                    "--ana-match", action="append", dest="ANAPATTERNS", default=[],
                    help="only write out histograms from analyses whose name matches any of these regexes")
selgroup.add_option(#"-A", #< these were confusing, and -M should be enough
                    "--ana-unmatch", action="append", dest="ANAUNPATTERNS", default=[],
                    help="exclude histograms from analyses whose name matches any of these regexes")
parser.add_option_group(selgroup)

vrbgroup = OptionGroup(parser, "Verbosity control")
vrbgroup.add_option("-v", "--verbose", help="add extra debug messages", dest="VERBOSE",
                  action="store_true", default=False)
parser.add_option_group(vrbgroup)

opts, yodafiles = parser.parse_args()


## Add pwd to search paths
if opts.PATH_PWD:
    rivet.addAnalysisLibPath(os.path.abspath("."))
    rivet.addAnalysisDataPath(os.path.abspath("."))


## Check that there are some arguments!
if not yodafiles:
    print("Error: You need to specify some YODA files to be plotted!")
    sys.exit(1)


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

## Get set of analyses involved in the runs
plotarg = None
analyses = set()
blocked_analyses = set()
import yoda
for yodafile in yodafiles:
    if yodafile.startswith("PLOT:"):
        plotarg = yodafile
        continue
    yodafilepath = os.path.abspath(yodafile.split(":")[0])
    if not os.access(yodafilepath, os.R_OK):
        print("Error: cannot read from %s" % yodafilepath)
        if opts.IGNORE_MISSING:
            continue
        else:
            sys.exit(2)

    try:
        ## Note: we use -m/-M flags here as well as when calling rivet-cmphistos, to potentially speed this initial loading
        analysisobjects = yoda.read(yodafilepath, patterns=opts.PATHPATTERNS, unpatterns=opts.PATHUNPATTERNS)
    except IOError as e:
        print("File reading error: ", e.strerror)
        sys.exit(1)

    for path, ao in analysisobjects.items():
        ## Make a path object and ensure the path is in standard form
        try:
            aop = rivet.AOPath(path)
        except Exception as e:
            #print(e)
            print("Found analysis object with non-standard path structure:", path, "... skipping")
            continue

        ## We don't plot data objects with path components hidden by an underscore prefix
        if aop.istmp():
            continue

        ## Identify analysis/histo name parts
        analysis = "ANALYSIS"
        if aop.basepathparts(keepref=False):
            analysis = aop.basepathparts(keepref=False)[0] #< TODO: for compatibility with rivet-cmphistos... generalise?
            #analysis = "_".join(aop.dirnameparts(keepref=False)[:-1]) #< TODO: would this be nicer? Currently incompatible with rivet-cmphistos

        ## Optionally veto on analysis name pattern matching
        if analysis in analyses.union(blocked_analyses):
            continue
        import re
        matched = True
        if opts.ANAPATTERNS:
            matched = False
            for patt in opts.ANAPATTERNS:
                if re.match(patt, analysis) is not None:
                    matched = True
                    break
        if matched and opts.ANAUNPATTERNS:
            for patt in opts.ANAUNPATTERNS:
                if re.match(patt, analysis) is not None:
                    matched = False
                    break
        if matched:
            analyses.add(analysis)
        else:
            blocked_analyses.add(analysis)


## Sort analyses: group ascending by analysis name (could specialise grouping by collider), then
## descending by year, and finally descending by bibliographic archive ID code (INSPIRE first).
def anasort(name):
    rtn = (1, name)
    if name.startswith("MC"):
        rtn = (99999999, name)
    else:
        stdparts = name.split("_")
        try:
            year = int(stdparts[1])
            rtn = (0, stdparts[0], -year, 0)
            idcode = (0 if stdparts[2][0] == "I" else 1e10) - int(stdparts[2][1:])
            rtn = (0, stdparts[0], -year, idcode)
            if len(stdparts) > 3:
                rtn += stdparts[3:]
        except:
            pass
    return rtn

analyses = sorted(analyses, key=anasort)

## Uncomment to test analysis ordering on index page
# print(analyses)
# sys.exit(0)


## Run rivet-cmphistos to get plain .dat files from .yoda
## We do this here since it also makes the necessary directories
ch_cmd = ["rivet-cmphistos"]
if opts.MC_ERRS:
    ch_cmd.append("--mc-errs")
if not opts.SHOW_RATIO:
    ch_cmd.append("--no-ratio")
if opts.NOPLOTTITLE:
    ch_cmd.append("--no-plottitle")
# if opts.REF_ID is not None:
#     ch_cmd.append("--refid=%s" % os.path.abspath(opts.REF_ID))
if opts.REFTITLE:
    ch_cmd.append("--reftitle=%s" % opts.REFTITLE )
if opts.PATHPATTERNS:
    for patt in opts.PATHPATTERNS:
        ch_cmd += ["-m", patt] #"'"+patt+"'"]
if opts.PATHUNPATTERNS:
    for patt in opts.PATHUNPATTERNS:
        ch_cmd += ["-M", patt] #"'"+patt+"'"]
ch_cmd.append("--hier-out")
# TODO: Need to be able to override this: provide a --plotinfodir cmd line option?
ch_cmd.append("--plotinfodir=%s" % os.path.abspath("../"))
for af in yodafiles:
    yodafilepath = os.path.abspath(af.split(":")[0])
    if af.startswith("PLOT:"):
        yodafilepath = "PLOT"
    elif not os.access(yodafilepath, os.R_OK):
        continue
    newarg = yodafilepath
    if ":" in af:
        newarg += ":" + af.split(":", 1)[1]
    # print(newarg)
    ch_cmd.append(newarg)

## Pass rivet-mkhtml -c args to rivet-cmphistos
for configfile in opts.CONFIGFILES:
    configfile = os.path.abspath(os.path.expanduser(configfile))
    if os.access(configfile, os.R_OK):
        ch_cmd += ["-c", configfile]

if opts.VERBOSE:
    ch_cmd.append("--verbose")
    print("Calling rivet-cmphistos with the following command:")
    print(" ".join(ch_cmd))

## Run rivet-cmphistos in a subdir, after fixing any relative paths in Rivet env vars
if not opts.DRY_RUN:
    for var in ("RIVET_ANALYSIS_PATH", "RIVET_DATA_PATH", "RIVET_REF_PATH", "RIVET_INFO_PATH", "RIVET_PLOT_PATH"):
        if var in os.environ:
            abspaths = [os.path.abspath(p) for p in os.environ[var].split(":")]
            os.environ[var] = ":".join(abspaths)
    subproc = Popen(ch_cmd, cwd=opts.OUTPUTDIR, stdout=PIPE, stderr=PIPE)
    out, err = subproc.communicate()
    retcode = subproc.returncode
    if opts.VERBOSE or retcode != 0:
        print('Output from rivet-cmphistos:\n', out)
    if err :
        print('Errors from rivet-cmphistos:\n', err)
    if retcode != 0:
        print('Crash in rivet-cmphistos code = ', retcode, ' exiting')
        exit(retcode)



## Write web page containing all (matched) plots
## Make web pages first so that we can load it locally in
## a browser to view the output before all plots are made
if not opts.DRY_RUN:

    style = '''<style>
      html { font-family: sans-serif; }
      img { border: 0; }
      a { text-decoration: none; font-weight: bold; }
    </style>
    '''

    ## Include MathJax configuration
    script = ''
    if not opts.OFFLINE:
        script = '''\
        <script type="text/x-mathjax-config">
        MathJax.Hub.Config({
          tex2jax: {inlineMath: [["$","$"]]}
        });
        </script>
        <script type="text/javascript"
          src="https://cdn.mathjax.org/mathjax/latest/MathJax.js?config=TeX-AMS-MML_HTMLorMML">
        </script>
        '''

    ## A helper function for metadata LaTeX -> HTML conversion
    from rivet.util import htmlify

    ## A timestamp HTML fragment to be used on each page:
    import datetime
    timestamp = '<p>Generated at %s</p>\n' % datetime.datetime.now().strftime("%A, %d. %B %Y %I:%M%p")

    index = open(os.path.join(opts.OUTPUTDIR, "index.html"), "w")
    index.write('<html>\n<head>\n<title>%s</title>\n%s</head>\n<body>' % (opts.TITLE, style + script))
    if opts.BOOKLET and opts.VECTORFORMAT == "PDF":
        index.write('<h2><a href="booklet.pdf">%s</a></h2>\n\n' % opts.TITLE)
    else:
        index.write('<h2>%s</h2>\n\n' % opts.TITLE)

    if opts.SINGLE:
        ## Write table of contents
        index.write('<ul>\n')
        for analysis in analyses:
            summary = analysis
            ana = rivet.AnalysisLoader.getAnalysis(analysis)
            if ana:
                summary = "%s (%s)" % (ana.summary(), analysis)
                if opts.IGNORE_UNVALIDATED and ana.status() != "VALIDATED":
                    continue
            index.write('<li><a href="#%s">%s</a>\n' % (analysis, htmlify(summary)) )
        index.write('</ul>\n')

    for analysis in analyses:
        references = []
        summary = htmlify(analysis)
        description, inspireid, spiresid = None, None, None

        if analysis.find("_I") > 0:
            inspireid = analysis[analysis.rfind('_I')+2:len(analysis)]
        elif analysis.find("_S") > 0:
            spiresid = analysis[analysis.rfind('_S')+2:len(analysis)]

        ana = rivet.AnalysisLoader.getAnalysis(analysis)
        if ana:
            if ana.summary():
                summary = htmlify("%s (%s)" % (ana.summary(), analysis))
            references = ana.references()
            description = htmlify(ana.description())
            spiresid = ana.spiresId()
            if opts.IGNORE_UNVALIDATED and ana.status().upper() != "VALIDATED":
                continue

        if opts.SINGLE:
            index.write('\n<h3 style="clear:left; padding-top:2em;"><a name="%s">%s</a></h3>\n' % (analysis, summary))
        else:
            index.write('\n<h3><a href="%s/index.html" style="text-decoration:none;">%s</a></h3>\n' % (analysis, summary))

        reflist = []
        if inspireid:
            reflist.append('<a href="http://inspire-hep.net/record/%s">Inspire</a>' % inspireid)
            reflist.append('<a href="http://hepdata.cedar.ac.uk/view/ins%s">HepData</a>' % inspireid)
        elif spiresid:
        # elif ana.spiresId():
            reflist.append('<a href="http://inspire-hep.net/search?p=find+key+%s">Inspire</a>' % spiresid)
            reflist.append('<a href="http://hepdata.cedar.ac.uk/view/irn%s">HepData</a>' % spiresid)
        reflist += references
        index.write('<p>%s</p>\n' % " &#124; ".join(reflist))

        if description:
            try:
                index.write('<p style="font-size:smaller;">%s</p>\n' % description)
            except UnicodeEncodeError as ue:
                    print("Unicode error in analysis description for " + analysis + ": " + str(ue))

        anapath = os.path.join(opts.OUTPUTDIR, analysis)
        if not opts.SINGLE:
            if not os.path.exists(anapath):
                os.makedirs(anapath)
            anaindex = open(os.path.join(anapath, "index.html"), 'w')
            anaindex.write('<html>\n<head>\n<title>%s &ndash; %s</title>\n%s</head>\n<body>\n' %
                           (htmlify(opts.TITLE), analysis, style + script))
            anaindex.write('<h3>%s</h3>\n' % htmlify(analysis))
            anaindex.write('<p><a href="../index.html">Back to index</a></p>\n')
            if description:
                try:
                    anaindex.write('<p>\n  %s\n</p>\n' % description)
                except UnicodeEncodeError as ue:
                    print("Unicode error in analysis description for " + analysis + ": " + str(ue))
        else:
            anaindex = index

        datfiles = glob.glob("%s/*.dat" % anapath)
        #print(datfiles)

        anaindex.write('<div style="float:none; overflow:auto; width:100%">\n')
        for datfile in sorted(datfiles):
            obsname = os.path.basename(datfile).replace(".dat", "")
            pngfile = obsname+".png"
            vecfile = obsname+"."+opts.VECTORFORMAT.lower()
            srcfile = obsname+".dat"
            if opts.SINGLE:
                pngfile = os.path.join(analysis, pngfile)
                vecfile = os.path.join(analysis, vecfile)
                srcfile = os.path.join(analysis, srcfile)

            anaindex.write('  <div style="float:left; font-size:smaller; font-weight:bold;">\n')
            anaindex.write('    <a href="#%s-%s">&#9875;</a><a href="%s">&#8984</a> %s:<br/>\n' %
                           (analysis, obsname, srcfile, os.path.splitext(vecfile)[0]) )
            anaindex.write('    <a name="%s-%s"><a href="%s">\n' % (analysis, obsname, vecfile) )
            anaindex.write('      <img src="%s">\n' % pngfile )
            anaindex.write('    </a></a>\n')
            anaindex.write('  </div>\n')
        anaindex.write('</div>\n')

        if not opts.SINGLE:
            anaindex.write('<div style="float:none">%s</body>\n</html></div>\n' % timestamp)
            anaindex.close()
    index.write('<br>%s</body>\n</html>' % timestamp)
    index.close()

# http://stackoverflow.com/questions/377017/test-if-executable-exists-in-python
def which(program):
    import os
    def is_exe(fpath):
        return os.path.isfile(fpath) and os.access(fpath, os.X_OK)

    fpath, fname = os.path.split(program)
    if fpath:
        if is_exe(program):
            return program
    else:
        for path in os.environ["PATH"].split(os.pathsep):
            path = path.strip('"')
            exe_file = os.path.join(path, program)
            if is_exe(exe_file):
                return exe_file

    return None

## Run make-plots on all generated .dat files
# sys.exit(0)
mp_cmd = ["make-plots"]
if opts.NUMTHREADS:
    mp_cmd.append("--num-threads=%d" % opts.NUMTHREADS)
if opts.NO_CLEANUP:
    mp_cmd.append("--no-cleanup")
if opts.NO_SUBPROC:
    mp_cmd.append("--no-subproc")
if opts.VECTORFORMAT == "PDF":
    mp_cmd.append("--pdfpng")
elif opts.VECTORFORMAT == "PS":
    mp_cmd.append("--pspng")
if opts.OUTPUT_FONT:
    mp_cmd.append("--font=%s" % opts.OUTPUT_FONT)
# if opts.OUTPUT_FONT.upper() == "PALATINO":
#     mp_cmd.append("--palatino")
# if opts.OUTPUT_FONT.upper() == "CM":
#     mp_cmd.append("--cm")
# elif opts.OUTPUT_FONT.upper() == "TIMES":
#     mp_cmd.append("--times")
# elif opts.OUTPUT_FONT.upper() == "HELVETICA":
#     mp_cmd.append("--helvetica")
# elif opts.OUTPUT_FONT.upper() == "MINION":
#     mp_cmd.append("--minion")
datfiles = []
for analysis in analyses:
    anapath = os.path.join(opts.OUTPUTDIR, analysis)
    #print(anapath)
    anadatfiles = glob.glob("%s/*.dat" % anapath)
    datfiles += sorted(anadatfiles)
if datfiles:
    mp_cmd += datfiles
    if opts.VERBOSE:
        mp_cmd.append("--verbose")
        print("Calling make-plots with the following options:")
        print(" ".join(mp_cmd))
    if not opts.DRY_RUN:
        Popen(mp_cmd).wait()
        if opts.BOOKLET and opts.VECTORFORMAT=="PDF":
            if which("pdftk") is not None:
                bookletcmd = ["pdftk"]
                for analysis in analyses:
                    anapath = os.path.join(opts.OUTPUTDIR, analysis)
                    bookletcmd += sorted(glob.glob("%s/*.pdf" % anapath))
                bookletcmd += ["cat", "output", "%s/booklet.pdf" % opts.OUTPUTDIR]
                print(bookletcmd)
                Popen(bookletcmd).wait()
            elif which("pdfmerge") is not None:
                bookletcmd = ["pdfmerge"]
                for analysis in analyses:
                    anapath = os.path.join(opts.OUTPUTDIR, analysis)
                    bookletcmd += sorted(glob.glob("%s/*.pdf" % anapath))
                bookletcmd += ["%s/booklet.pdf" % opts.OUTPUTDIR]
                print(bookletcmd)
                Popen(bookletcmd).wait()
            else:
                print("Neither pdftk nor pdfmerge available --- not booklet output possible")
