#! /usr/bin/env python3

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

Convert a YODA data file to another YODA file (for convenient tidying and filtering).
Like yodacnv, but allows output to stdout since the output format isn't auto-detected.
"""

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("--make-inert", dest="MAKE_INERT", action="store_true", default=False,
                    help="convert all input analysis objects to their inert equivalents")
parser.add_argument("--as-scatters", dest="AS_SCATTERS", action="store_true", default=False,
                    help="convert all input analysis objects to Scatter types")
parser.add_argument("--rename-source", dest="SOURCE", default="",
                    help="rename error source of the inert objects (requires --make-inert flag)")
parser.add_argument("--split-by-variation", "--split-variations", dest="SPLIT_VARS", action="store_true", default=False,
                    help="separate multiweight streams into separate outputfiles")

def extractWeightName(path):
    import re
    re_weight = re.compile(r".*\[(.*?)\].*")
    m = re_weight.match(path)
    if not m:
        return ""
    else:
        return m.group(1)

args = parser.parse_args()
in_out = parse_x2y_args(args.ARGS, [".yoda", ".yoda.gz", ".yoda.h5"], [".yoda", ".yoda.gz", ".yoda.h5"])
if not in_out:
    sys.stderr.write("You must specify the YODA input and output file names\n")
    sys.exit(1)
# print(in_out)

for i, o in in_out:
    analysisobjects = yoda.read(i, False, args.MATCH, args.UNMATCH)
    if args.AS_SCATTERS:
        analysisobjects = [ao.mkScatter(ao.path()) for ao in analysisobjects]
    elif args.MAKE_INERT:
        analysisobjects = [ao.mkInert(ao.path(), args.SOURCE) for ao in analysisobjects]
    if args.SPLIT_VARS:
        for wname in set([ extractWeightName(ao.path()) for ao in analysisobjects ]):
            sub_annos = [ ao for ao in analysisobjects if (wname in ao.path() if wname else \
                                                           not ao.path().endswith(']')) ]
            ow = o.replace(".yoda", "_"+wname+".yoda") if wname else o
            yoda.write(sub_annos, ow)
    else:
        yoda.write(analysisobjects, o)
