#! /usr/bin/env python3

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

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

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

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("-g", "--as-graphs", dest="AS_GRAPHS", action="store_true", default=False,
                    help="convert to ROOT TGraphs rather than native types")
parser.add_argument("-f", "--use-focus", dest="USE_FOCUS", action="store_true", default=False,
                    help="use the bin focus rather than centre when creating TGraph points")
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, ".yoda", ".root")
if not in_out:
    sys.stderr.write("You must specify the YODA and ROOT file names\n")
    sys.exit(1)

try:
    import ROOT
    ROOT.gROOT.SetBatch(True)
    ROOT.TH1.AddDirectory(False)
except ImportError:
    sys.stderr.write("Could not load ROOT Python module, exiting...\n")
    sys.exit(2)

for i, o in in_out:
    of = ROOT.TFile(o, "recreate")
    analysisobjects = yoda.read(i, False, args.MATCH, args.UNMATCH)
    rootobjects = []
    for ao in analysisobjects:
        try:
            rootobjects.append( yoda.root.to_root(ao,
                                                  asgraph=args.AS_GRAPHS,
                                                  usefocus=args.USE_FOCUS,
                                                  widthscale=args.DIVBINSIZE) )
        except Exception as ex:
            print(f"Warning: could not convert object '{ao.path()}': {ex} ... skipping")

    ## Protect against "/" in the histogram name, which ROOT does not like
    for obj in rootobjects:
        ## It's possible for the ROOT objects to be null, if conversion failed
        if obj is None:
            continue
        ## Split the name on "/" directory separators
        parts = obj.GetName().split("/")
        ## Walk down dir tree, making it as we go
        d = of
        for part in parts[:-1]:
            subdir = d.GetDirectory(part)
            if not subdir:
                subdir = d.mkdir(part)
            d = subdir
        ## Write the histo into the leaf dir
        d.WriteTObject(obj, parts[-1])
    of.Close()
