#! /usr/bin/env python

"""\
%(prog)s rootfile [yodafile]

Convert a ROOT data file to the YODA data format.
"""

import yoda, os, sys, argparse
from yoda.script_helpers import parse_x2y_args, filter_aos

parser = argparse.ArgumentParser(usage=__doc__)
parser.add_argument("ARGS", nargs="+", help="infile [outfile]")
parser.add_argument("-m", "--match", dest="MATCH", metavar="PATT", default=None,
                    help="only write out histograms whose path matches this regex")
parser.add_argument("-M", "--unmatch", dest="UNMATCH", metavar="PATT", default=None,
                    help="exclude histograms whose path matches this regex")
parser.add_argument("-d", "--div-by-binsize", dest="DIVBINSIZE", action="store_true", default=False,
                    help="divide bin values by bin size (width, area, etc.)")

args = parser.parse_args()
in_out = parse_x2y_args(args.ARGS, ".root", ".yoda")
if not in_out:
    sys.stderr.write("You must specify the ROOT and YODA file names\n")
    sys.exit(1)

import ROOT
ROOT.gROOT.SetBatch(True)
for i, o in in_out:
    rf = ROOT.TFile(i)
    rootobjects_raw = list(yoda.root.getall(rf))
    rootobjects     = [(path, ro) for (path, ro) in rootobjects_raw if not isinstance(ro, ROOT.TH1F)]
    th1f            = [(path, ro) for (path, ro) in rootobjects_raw if     isinstance(ro, ROOT.TH1F)]
    # Conversion of TH1F into TH1D
    for path, ro in th1f:
        temp = ROOT.TH1D()
        ro.Copy(temp)
        rootobjects.append((path, temp))

    def to_yoda(path, ro):
        if path is None or ro is None:
            return None
        ao = yoda.root.to_yoda(ro, widthscale=args.DIVBINSIZE)
        if ao is not None:
            ao.setPath(path)
        return ao

    analysisobjects = [to_yoda(path, ro) for (path, ro) in rootobjects]
    rf.Close()

    analysisobjects = [ao for ao in analysisobjects if ao is not None]
    filter_aos(analysisobjects, args.MATCH, args.UNMATCH)
    yoda.writeYODA(analysisobjects, o)
