#! /usr/bin/env python

"""\
%prog <yodafile1> [<yodafile2> ...]

Make plot files from YODA histograms

PROTOTYPE VERSION! INTERFACE AND BEHAVIOUR WILL CHANGE!


Perform basic plotting of data objects stored in YODA-readable data files, using
the LaTeX pgfplots package as a backend. May develop into something more fully
fledged, with support for PLOT, FUNCTION, etc. objects in YODA files, but we're
being cautious about that since we don't want to get stuck with a turkey.

TODO:
 * style control from .plot files and PLOT sections
 * decide on automatic lin vs. log axis (include zero/-ve?)
 * work out automatic legend placement by displayed density of points
 * line colours and styling
 * ratio plots and plot grids
 * axis labels at _ends_ of axes
 * error bands
 * tick label format
 * histogram comparison
 * function (sampling with uniform display density), stack, labels, ...
 * speed up (batch processing, threading, etc.)
"""

import optparse
op = optparse.OptionParser(usage=__doc__)
op.add_option("-s", "--split", dest="SPLIT", action="store_true", default=False)
op.add_option("-f", "--format", dest="FORMAT", default="pdf")
op.add_option("-m", "--match", dest="MATCH", default=None)
op.add_option("-M", "--unmatch", dest="UNMATCH", default=None)
opts, args = op.parse_args()

head = \
r"""\documentclass[multi=tikzpicture,border=2mm]{standalone}
\usepackage[T1]{fontenc}

\usepackage{underscore}
\usepackage{url}
\usepackage{microtype}

\usepackage{mathpazo}

\usepackage{tikz}
\usepackage{pgfplots}
% TODO: use siunitx for number typesetting?
\pgfplotsset{every axis/.append style={thick,tick style={thick,black}, error bars/error bar style={thick}}}

\begin{document}%
"""

foot = \
r"""\end{document}
"""

import re
re_match = re.compile(opts.MATCH) if opts.MATCH else None
re_unmatch = re.compile(opts.UNMATCH) if opts.UNMATCH else None

for arg in args:

    ## Separate filename and style overrides from each arg
    infile = arg
    argstyles = {"LineColor":"red", "LineStyle":"solid"}
    if ":" in arg:
        tmp = arg.split(":")
        infile = tmp[0]
        argstyles = dict(s.split("=") if "=" in s else ("Title", s) for s in tmp[1:])

    import yoda, os
    basename = os.path.splitext(os.path.basename(infile))[0] if infile != "-" else "plots"
    aos = yoda.read(infile) if infile != "-" else yoda.readYODA(infile)
    names = []
    bodies = []
    for path, ao in sorted(aos.iteritems()):
        if (re_match and not re_match.match(path)) or (re_unmatch and re_unmatch.match(path)):
            continue

        s = ao.mkScatter()
        aostyles = {ann : ao.annotation(ann) for ann in ao.annotations}
        aostyles.update(argstyles)

        body = "\\begin{tikzpicture}\n"
        ## TODO: dynamically determine default axis types based on distribution of points
        # xmin = min(p.x - p.xErrs[0] for p in s.points)
        # xmax = max(p.x + p.xErrs[1] for p in s.points)
        ## TODO: want a symlog axis option like matplotlib offers...
        axistypeidx = 2*int(ao.annotation("LogX", False)) + int(ao.annotation("LogY", False))
        axistype = ("axis", "semilogyaxis", "semilogxaxis", "loglogaxis")[axistypeidx]
        body += "  \\begin{%s}[" % axistype
        if ao.hasAnnotation("Title"):
            body += "title={%s}, " % ao.annotation("Title")
        if ao.hasAnnotation("XLabel"):
            body += "xlabel={%s}, " % ao.annotation("XLabel")
        if ao.hasAnnotation("YLabel"):
            body += "ylabel={%s}, " % ao.annotation("YLabel")
        body += "legend style={draw=none}, "
        body += "enlarge x limits=false]\n"

        isref = ao.path.startswith("/REF")
        aostyles.setdefault("LineColor", "black" if isref else "red")
        aostyles.setdefault("LineWidth", "very thick" if isref else "thick")
        aostyles.setdefault("LineStyle", "none" if isref else "solid")
        aostyles.setdefault("Marker", "*" if isref else "none")
        aostyles.setdefault("MarkerScale", 0.8)
        aostyles.setdefault("ErrorBars", isref)
        aostyles.setdefault("Legend", True)

        if type(s) is yoda.Scatter2D:
            body += "    \\addplot["
            body += "color=%s, " % aostyles["LineColor"]
            body += "%s, " % aostyles["LineWidth"]
            if aostyles["LineStyle"] == "none":
                body += "only marks, "
            else:
                body += "const plot mark mid, %s" % aostyles["LineStyle"]
            if aostyles["Marker"] != "none":
                body += "mark=%s, mark options={scale=%g}, " % (aostyles["Marker"], aostyles["MarkerScale"])
            if aostyles["ErrorBars"]:
                body += "error bars/.cd, error mark=, x dir=both, x explicit, y dir=both, y explicit"
            body += "] coordinates {\n"
            # body += "    \\addplot[color=black, very thick, only marks, mark=*, mark options={scale=0.8}, error bars/.cd, error mark=, x dir=both, x explicit, y dir=both, y explicit] coordinates {\n"
            # body += "    \\addplot[color=red, const plot mark mid] coordinates {\n"
            for p in s.points:
                body += "      ({x:g}, {y:g}) -= ({ex[0]:g}, {ey[0]:g}) += ({ex[1]:g}, {ey[1]:g})\n".format(x=p.x, y=p.y, ex=p.xErrs, ey=p.yErrs)
            body += "    };\n"

        elif type(s) is yoda.Scatter3D:
            body += "    \\addplot3["
            body += "mesh/row=%d, " % ao.numBinsX
            body += "surf, "
            body += "] coordinates {\n"
            for p in s.points:
                body += "      ({x:g}, {y:g}, {z:g})\n".format(x=p.x, y=p.y, z=p.z)
            body += "    };\n"

        if aostyles["Legend"]:
            body += "  \\addlegendentry{%s}%%\n" % basename

        body += "  \\end{%s}%%\n" % axistype
        body += "\\end{tikzpicture}%\n"

        name = ao.path.replace("/", "_").lstrip("_")
        names.append(name)
        bodies.append(body)


    ## Make TeX and PDF files in a temporary directory
    import tempfile, subprocess, shutil, glob
    rundir = os.getcwd()
    tmpdir = tempfile.mkdtemp(".yodaplot")
    os.chdir(tmpdir)

    ## Make LaTeX source
    texname = basename + ".tex"
    with open(texname, "w") as f:
        f.write(head + "".join(bodies) + foot)
    if "TEX" in opts.FORMAT.upper():
        shutil.copy(texname, rundir)

    ## Make PDF via pdflatex
    pdfname = basename + ".pdf"
    if "PDF" in opts.FORMAT.upper():
        subprocess.check_call(["pdflatex", texname])
        if not opts.SPLIT:
            shutil.copy(pdfname, rundir)
        else:
            subprocess.check_call(["pdftk", pdfname, "burst"])
            for i, n in enumerate(names, start=1):
                subpdfname = basename + "_" + n + ".pdf"
                shutil.copy("pg_{:04d}.pdf".format(i), os.path.join(rundir, subpdfname))

    os.chdir(tmpdir)
    shutil.rmtree(tmpdir)
