#! /usr/bin/env python

"""\
%prog - generate histogram comparison plots

USAGE:
 %prog [options] yodafile1[:'PlotOption1=Value':'PlotOption2=Value':...] [path/to/yodafile2 ...] [PLOT:Key1=Val1:...]

where the plot options are described in the make-plots manual in the HISTOGRAM
section.
"""

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


class Plot(dict):
    "A tiny Plot object to help writing out the head in the .dat file"
    def __repr__(self):
        return "# BEGIN PLOT\n" + "\n".join("%s=%s" % (k,v) for k,v in self.iteritems()) + "\n# END PLOT\n\n"


def sanitiseString(s):
    #s = s.replace('_','\\_')
    #s = s.replace('^','\\^{}')
    #s = s.replace('$','\\$')
    s = s.replace('#','\\#')
    s = s.replace('%','\\%')
    return s


def getCommandLineOptions():
    "Parse command line options"
    from optparse import OptionParser, OptionGroup
    parser = OptionParser(usage=__doc__)

    parser.add_option("--no-rivet-refs", dest="RIVETREFS", action="store_false",
                      default=True, help="don't use Rivet reference data files")
    parser.add_option('-o', '--outdir', dest='OUTDIR',
                      default='.', help='write data files into this directory')
    parser.add_option("--hier-out", action="store_true", dest="HIER_OUTPUT", default=False,
                      help="write output dat files into a directory hierarchy which matches the analysis paths")
    parser.add_option('--plotinfodir', dest='PLOTINFODIRS', action='append',
                      default=['.'], help='directory which may contain plot header information (in addition '
                      'to standard Rivet search paths)')

    stygroup = OptionGroup(parser, "Plot style")
    # stygroup.add_option("--refid", dest="REF_ID",
    #                     default="REF", help="ID of reference data set (file path for non-REF data)")
    stygroup.add_option("--linear", action="store_true", dest="LINEAR",
                        default=False, help="plot with linear scale")
    stygroup.add_option("--mc-errs", action="store_true", dest="MC_ERRS",
                        default=False, help="show vertical error bars on the MC lines")
    stygroup.add_option("--no-ratio", action="store_false", dest="RATIO",
                        default=True, help="disable the ratio plot")
    stygroup.add_option("--rel-ratio", action="store_true", dest="RATIO_DEVIATION",
                        default=False, help="show the ratio plots scaled to the ref error")
    stygroup.add_option("--no-plottitle", action="store_true", dest="NOPLOTTITLE",
                        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("--style", dest="STYLE", default="default",
                        help="change plotting style: default|bw|talk")
    stygroup.add_option("-c", "--config", dest="CONFIGFILES", action="append", default=["~/.make-plots"],
                        help="additional plot config file(s). Settings will be included in the output configuration.")
    parser.add_option_group(stygroup)

    selgroup = OptionGroup(parser, "Selective plotting")
    # selgroup.add_option("--show-single", dest="SHOW_SINGLE", choices=("no", "ref", "mc", "all"),
    #                     default="mc", help="control if a plot file is made if there is only one dataset to be plotted "
    #                     "[default=%default]. If the value is 'no', single plots are always skipped, for 'ref' and 'mc', "
    #                     "the plot will be written only if the single plot is a reference plot or an MC "
    #                     "plot respectively, and 'all' will always create single plot files.\n The 'ref' and 'all' values "
    #                     "should be used with great care, as they will also write out plot files for all reference "
    #                     "histograms without MC traces: combined with the -R/--rivet-refs flag, this is a great way to "
    #                     "write out several thousand irrelevant reference data histograms!")
    # selgroup.add_option("--show-mc-only", "--all", action="store_true", dest="SHOW_IF_MC_ONLY",
    #                     default=False, help="make a plot file even if there is only one dataset to be plotted and "
    #                     "it is an MC one. Deprecated and will be removed: use --show-single instead, which overrides this.")
    # # selgroup.add_option("-l", "--histogram-list", dest="HISTOGRAMLIST",
    # #                     default=None, help="specify a file containing a list of histograms to plot, in the format "
    # #                     "/ANALYSIS_ID/histoname, one per line, e.g. '/DELPHI_1996_S3430090/d01-x01-y01'.")
    selgroup.add_option("-m", "--match", action="append",
                        help="Only write out histograms whose $path/$name string matches these regexes. The argument "
                        "may also be a text file.",
                        dest="PATHPATTERNS")
    selgroup.add_option("-M", "--unmatch", action="append",
                        help="Exclude histograms whose $path/$name string matches these regexes",
                        dest="PATHUNPATTERNS")
    parser.add_option_group(selgroup)

    return parser


def getHistos(filelist):
    """Loop over all input files. Only use the first occurrence of any REF-histogram
    and the first occurrence in each MC file for every MC-histogram."""
    refhistos = {}
    mchistos = {}
    for infile in filelist:
        mchistos.setdefault(infile, {})
        analysisobjects = yoda.read(infile, patterns=opts.PATHPATTERNS, unpatterns=opts.PATHUNPATTERNS)
        for path, ao in analysisobjects.iteritems():
            if path.startswith('/REF/'):
                if not refhistos.has_key(path):
                    refhistos[path] = ao
            else:
                if not mchistos[infile].has_key(path):
                    mchistos[infile][path] = ao
    return refhistos, mchistos


def getRivetRefData(refhistos, anas=None):
    "Find all Rivet reference data files"
    rivet_data_dirs = rivet.getAnalysisRefPaths()
    dirlist = []
    for d in rivet_data_dirs:
        if anas is None:
            import glob
            dirlist.append(glob.glob(os.path.join(d, '*.yoda')))
        else:
            dirlist.append([os.path.join(d, a+'.yoda') for a in anas])
    for filelist in dirlist:
        # TODO: delegate to getHistos?
        for infile in filelist:
            analysisobjects = yoda.read(infile, patterns=opts.PATHPATTERNS, unpatterns=opts.PATHUNPATTERNS)
            for path, ao in analysisobjects.iteritems():
                if path.startswith('/REF/'):
                    if not refhistos.has_key(path):
                        refhistos[path] = ao


def parseArgs(args):
    """Look at the argument list and split it at colons, in order to separate
    the file names from the plotting options. Store the file names and
    file specific plotting options."""
    filelist = []
    plotoptions = {}
    for a in args:
        asplit = a.split(':')
        path = asplit[0]
        filelist.append(path)
        plotoptions[path] = []
        has_title = False
        for i in xrange(1, len(asplit)):
            ## Add 'Title' if there is no = sign before math mode
            if '=' not in asplit[i] or ('$' in asplit[i] and asplit[i].index('$') < asplit[i].index('=')):
                asplit[i] = 'Title=%s' % asplit[i]
            if asplit[i].startswith('Title='):
                has_title = True
            plotoptions[path].append(asplit[i])
        if path != "PLOT" and not has_title:
            plotoptions[path].append('Title=%s' % sanitiseString(os.path.basename( os.path.splitext(path)[0] )) )
    return filelist, plotoptions


def setStyle(ao, style):
    """Set default plot styles (color and line width) colors borrowed from Google Ngrams"""
    LINECOLORS = ['{[HTML]{EE3311}}',  # red (Google uses 'DC3912')
                  '{[HTML]{3366FF}}',  # blue
                  '{[HTML]{109618}}',  # green
                  '{[HTML]{FF9900}}',  # orange... weirdly this screws up if the F is lower-case!
                  '{[HTML]{990099}}']  # lilac
    LINESTYLES = ['solid',
                  'dashed',
                  'dashdotted',
                  'dotted']

    if opts.STYLE == 'talk':
        ao.setAnnotation('LineWidth', '1pt')
    if opts.STYLE == 'bw':
        LINECOLORS = ['black!90',
                      'black!50',
                      'black!30']

    c = style % len(LINECOLORS)
    s = (style / len(LINECOLORS)) % len(LINESTYLES)

    ao.setAnnotation('LineStyle', '%s' % LINESTYLES[s])
    ao.setAnnotation('LineColor', '%s' % LINECOLORS[c])


def setOptions(ao, options):
    "Set arbitrary annotations"
    for opt in options:
        key, val = opt.split('=', 1)
        ao.setAnnotation(key, val)


# TODO: move to rivet.utils
def mkoutdir(outdir):
    "Function to make output directories"
    if not os.path.exists(outdir):
        try:
            os.makedirs(outdir)
        except:
            msg = "Can't make output directory '%s'" % outdir
            raise Exception(msg)
    if not os.access(outdir, os.W_OK):
        msg = "Can't write to output directory '%s'" % outdir
        raise Exception(msg)


def writeOutput(output, h):
    "Choose output file name and dir"
    hparts = h.strip("/").split("/")
    if opts.HIER_OUTPUT:
        ana = "_".join(hparts[:-1]) if len(hparts) > 1 else "ANALYSIS"
        outdir = os.path.join(opts.OUTDIR, ana)
        outfile = '%s.dat' % hparts[-1]
    else:
        outdir = opts.OUTDIR
        outfile = '%s.dat' % "_".join(hparts)
    mkoutdir(outdir)
    outfilepath = os.path.join(outdir, outfile)
    f = open(outfilepath, 'w')
    f.write(output)
    f.close()


#--------------------------------------------------------------------------------------------


if __name__ == '__main__':

    ## Command line parsing
    parser = getCommandLineOptions()
    opts, args = parser.parse_args()

    ## Split the input file names and the associated plotting options
    ## given on the command line into two separate lists
    filelist, plotoptions = parseArgs(args)
    ## Remove the PLOT dummy file from the file list
    if "PLOT" in filelist:
        filelist.remove("PLOT")

    ## Read the .plot files
    plotdirs = opts.PLOTINFODIRS + [os.path.abspath(os.path.dirname(f)) for f in filelist]
    plotparser = rivet.mkStdPlotParser(plotdirs, opts.CONFIGFILES)

    ## Create a list of all histograms to be plotted
    refhistos, mchistos = getHistos(filelist)
    hpaths = []
    for aos in mchistos.values():
        for p in aos.keys():
            if p and p not in hpaths:
                hpaths.append(p)

    ## Read the reference data from the Rivet search paths and add them
    ## to the list of reference histograms
    if opts.RIVETREFS:
        getRivetRefData(refhistos)

    ## Now loop over all MC histograms and plot them
    # TODO: factorize much of this into a rivet.utils mkplotfile(mchists, refhist, kwargs) function
    for h in hpaths:
        #print 'Currently looking at', h
        ## A list of all analysis objects to be plotted
        anaobjects = []

        ## Plot object for the PLOT section in the .dat file
        plot = Plot()
        plot['Legend'] = '1'
        plot['LogY'] = '1'
        for key, val in plotparser.getHeaders(h).iteritems():
            plot[key] = val
        if plotoptions.has_key("PLOT"):
            for key_val in plotoptions["PLOT"]:
                key, val = [s.strip() for s in key_val.split("=")]
                plot[key] = val
        if opts.LINEAR:
            plot['LogY'] = '0'
        if opts.NOPLOTTITLE:
            plot['Title'] = ''

        if opts.STYLE == 'talk':
            plot['PlotSize'] = '8,6'
        elif opts.STYLE == 'bw':
            if opts.RATIO:
                plot['RatioPlotErrorBandColor'] = 'black!10'

        ## DrawOnly is needed to keep the order in the Legend equal to the
        ## order of the files on the command line
        drawonly = []

        ## Check if we have reference data for the histogram
        ratioreference = None
        if refhistos.has_key('/REF' + h):
            refdata = refhistos['/REF' + h]
            refdata.setAnnotation('ErrorBars', '1')
            refdata.setAnnotation('PolyMarker', '*')
            refdata.setAnnotation('ConnectBins', '0')
            refdata.setAnnotation('Title', 'Data')
            if opts.RATIO:
                ratioreference = '/REF'+h
            anaobjects.append(refdata)
            drawonly.append('/REF' + h)

        if opts.RATIO and opts.RATIO_DEVIATION:
            plot['RatioPlotMode'] = 'deviation'

        ## Loop over the MC files to plot all instances of the histogram
        styleidx = 0
        for infile in filelist:
            # if infile == "PLOT":
            #     continue  ##< This isn't a real file!
            if mchistos.has_key(infile) and mchistos[infile].has_key(h):
                ## Default linecolor, linestyle
                setStyle(mchistos[infile][h], styleidx)
                styleidx += 1
                if opts.MC_ERRS:
                    mchistos[infile][h].setAnnotation('ErrorBars', '1')
                ## Plot defaults from .plot files
                for key, val in plotparser.getHistogramOptions(h).iteritems():
                    mchistos[infile][h].setAnnotation(key, val)
                ## Command line plot options
                setOptions(mchistos[infile][h], plotoptions[infile])
                mchistos[infile][h].setAnnotation('Path', infile + h)
                anaobjects.append(mchistos[infile][h])
                drawonly.append(infile + h)
                if opts.RATIO and ratioreference is None:
                    ratioreference = infile + h

        plot['DrawOnly'] = ' '.join(drawonly).strip()
        if opts.RATIO and len(drawonly) > 1:
            plot['RatioPlot'] = '1'
            plot['RatioPlotReference'] = ratioreference

        ## Now create the output. We can't use "yoda.writeFLAT(anaobjects, 'foobar.dat')" because
        ## the PLOT and SPECIAL blocks don't have a corresponding analysis object.

        output = ''
        output += str(plot)

        ## Special
        special = plotparser.getSpecial(h)
        if special:
            output += "\n"
            output += "# BEGIN SPECIAL %s\n" % h
            output += special
            output += "# END SPECIAL\n\n"

        from cStringIO import StringIO
        sio = StringIO()
        yoda.writeFLAT(anaobjects, sio)
        output += sio.getvalue()
        ## Write everything into a file
        writeOutput(output, h)
