#! /usr/bin/env python

"""\
%prog [options] flatfile [flatfile2 ...]

Convert make-plots data files to AIDA XML format. The output is by default
written out to a file with the same name as the input (out.aida in the case of
stdin) unless the --output option is specified. When specifying either input or
output filenames, a '-' is used to refer to stdin or stdout as appropriate.

Histograms can also be filtered by histo path, using the -m or -M options for a
positive or negative regex pattern patch respectively.
"""

import sys
if sys.version_info[:3] < (2,4,0):
    print "rivet scripts require Python version >= 2.4.0... exiting"
    sys.exit(1)


import os, logging
import lighthisto


##########################################################


if __name__ == "__main__":

    ## Default plot file search paths
    default_plotdirs = ["."]
    try:
        import rivet
        default_plotdirs += rivet.getAnalysisPlotPaths()
    except:
        pass


    ## Parse command line options
    from optparse import OptionParser, OptionGroup
    parser = OptionParser(usage=__doc__)
    parser.add_option("-o", "--output", default=None,
                      help="Write all histos to a single output file. "
                      "stdout can be explicitly specified by setting '-' as the output filename. This option will "
                      "be disregarded if --split is specified.", dest="OUTPUT")
    parser.add_option("-s", "--split", action="store_true", default=False,
                      help="Split histograms into individual files", dest="SPLITOUTPUT")
    parser.add_option("--plotinfodir", dest="PLOTINFODIR", action="append",
                      default=default_plotdirs, help="directory which may contain plot header information")
    parser.add_option("-m", "--match", action="append",
                      help="Only write out histograms whose $path/$name string matches these regexes",
                      dest="PATHPATTERNS")
    parser.add_option("-M", "--unmatch", action="append",
                      help="Exclude histograms whose $path/$name string matches these regexes",
                      dest="PATHUNPATTERNS")
    verbgroup = OptionGroup(parser, "Verbosity control")
    verbgroup.add_option("-v", "--verbose", action="store_const", const=logging.DEBUG, dest="LOGLEVEL",
                         default=logging.INFO, help="print debug (very verbose) messages")
    verbgroup.add_option("-q", "--quiet", action="store_const", const=logging.WARNING, dest="LOGLEVEL",
                         default=logging.INFO, help="be very quiet")
    opts, args = parser.parse_args()


    ## Configure logging
    logging.basicConfig(level=opts.LOGLEVEL, format="%(message)s")


    ## Initialise steering variables which need a bit more care
    import re
    if opts.PATHPATTERNS is None:
        opts.PATHPATTERNS = []
    opts.PATHPATTERNS = [re.compile(r) for r in opts.PATHPATTERNS]
    if opts.PATHUNPATTERNS is None:
        opts.PATHUNPATTERNS = []
    opts.PATHUNPATTERNS = [re.compile(r) for r in opts.PATHUNPATTERNS]


    ## Check that at least one file has been supplied
    if len(args) < 1:
        sys.stderr.write("Must specity at least one histogram file (or stdin)\n")
        sys.exit(1)


    ## Add directories to the plotinfo path
    for flatfile in args:
        if flatfile != "-":
            flatdir = os.path.dirname(flatfile)
            if flatdir not in opts.PLOTINFODIR:
                opts.PLOTINFODIR.append(flatdir)
    ## Remove empty path entries
    opts.PLOTINFODIR = filter(lambda s: len(s) > 0, opts.PLOTINFODIR)
    ## Create plot file parser
    plotparser = lighthisto.PlotParser(opts.PLOTINFODIR)


    ## Run over the files and build histo objects selected by the pattern filtering
    histos = []
    for flatfile in args:
        if flatfile != "-" and not os.access(flatfile, os.R_OK):
            logging.error("%s can not be read" % flatfile)
            sys.exit(1)
        try:
            allhistos = lighthisto.Histo.fromFlat(flatfile)
        except Exception, e:
            logging.error("%s can not be parsed" % flatfile)
            print e
            sys.exit(1)

        for hist in allhistos:
            histpath = hist.fullPath()
            useThis = True
            if opts.PATHPATTERNS:
                useThis = False
                for regex in opts.PATHPATTERNS:
                    if regex.search(histpath):
                        useThis = True
                        break
            if useThis and opts.PATHUNPATTERNS:
                for regex in opts.PATHUNPATTERNS:
                    if regex.search(histpath):
                        useThis = False
                        break
            if useThis:
                try:
                    plotparser.updateHistoHeaders(hist)
                except ValueError, err:
                    logging.debug(err)
                histos.append((flatfile, hist))


    ## Write output
    if histos:
        ## Split output per-histogram
        if opts.SPLITOUTPUT:
            for f, h in histos:
                histo = h.fullPath()[1:].replace("/", "_")
                outfile = "%s.aida" % histo
                out = open(outfile, "w")
                out.write('<?xml version="1.0" ?>\n')
                out.write('<!DOCTYPE aida SYSTEM "http://aida.freehep.org/schemas/3.3/aida.dtd">\n')
                out.write('<aida version="3.3">\n')
                out.write('  <implementation version="1.1" package="Rivet"/>\n')
                out.write(h.asAIDA())
                out.write('</aida>\n')
                out.close()
        ## Write all output to a single file (stdout by default)
        elif opts.OUTPUT:
            outfile = opts.OUTPUT
            if outfile == "-":
                out = sys.stdout
            else:
                out = open(outfile, "w")
            out.write('<?xml version="1.0" ?>\n')
            out.write('<!DOCTYPE aida SYSTEM "http://aida.freehep.org/schemas/3.3/aida.dtd">\n')
            out.write('<aida version="3.3">\n')
            out.write('  <implementation version="1.1" package="Rivet"/>\n')
            for f, h in histos:
                out.write(h.asAIDA())
            out.write('</aida>\n')
            if outfile != "-":
                out.close()
        ## Split output per-infile
        else:
            histodict = dict()
            for f, h in histos:
                histodict.setdefault(f, [])
                histodict[f].append(h)
            for f, hs in histodict.iteritems():
                outfile = os.path.basename(f).replace(".dat", ".aida")
                if f == "-":
                    outfile = "out.dat"
                out = open(outfile, "w")
                out.write('<?xml version="1.0" ?>\n')
                out.write('<!DOCTYPE aida SYSTEM "http://aida.freehep.org/schemas/3.3/aida.dtd">\n')
                out.write('<aida version="3.3">\n')
                out.write('  <implementation version="1.1" package="Rivet"/>\n')
                for h in hs:
                    out.write(h.asAIDA())
                out.write('</aida>\n')
                out.close()
