#! /usr/bin/env python

"""\
%(prog)s <yodafile1> [<yodafile2> ...] -o <yodaoutfile> -r <yodareffile>

Combine the central values of multiple YODA files into an envelope
and write out the resulting envelope into a separate output YODA file.
If a reference file is specified, the resulting central value are taken
from that file, otherwise the midpoint between the edges of the envelope
will be used.

Examples:

%(prog)s -o out.yoda file1.yoda file2.yoda
  Write out a file with the envelope constructed from file1 and file2 and with the central value set to the midpoint
  between the envelope edges.

%(prog)s -o out.yoda -r reffile.yoda file1.yoda file2.yoda
  Write out a file with the envelope constructed from reffile, file1 and file2, with the central value set to the
  central values from reffile.

"""


from __future__ import print_function

## Parse command line args
import argparse, sys
parser = argparse.ArgumentParser(usage=__doc__)
parser.add_argument("INFILES", nargs="+",
                    help="file or folder with reference histos")
parser.add_argument("-r", "--ref", default=None, dest="REF_FILE",
                    help="file to be used for central values, midpoint is used if none if specified")
parser.add_argument("-o", "--output", default="-", dest="OUTPUT_FILE", metavar="PATH",
                    help="write output to specified path")
parser.add_argument("-m", "--match", dest="MATCH", metavar="PATH", default=None,
                    help="only write out histograms whose path matches this regex")
parser.add_argument("-M", "--unmatch", dest="UNMATCH", metavar="PATH", default=None,
                    help="exclude histograms whose path matches this regex")
parser.add_argument("-q", "--quiet", dest="VERBOSITY", action="store_const", const=0, default=1,
                    help="reduce printouts to errors-only")
parser.add_argument("-v", "--debug", dest="VERBOSITY", action="store_const", const=2, default=1,
                   help="increase printouts to include debug info")
args = parser.parse_args()


nfiles = len(args.INFILES) + (not args.REF_FILE is None)
if (nfiles < 2 or args.INFILES[0] == args.REF_FILE):
    sys.stderr.write("ERROR: Need at least two input files to construct an envelope!")
    sys.exit(1)

import yoda


def add2envelope(data, aos, keepScatters = None):
    for aoname, ao in aos.items():
        ao = ao.mkScatter()
        if not keepScatters is None:
            keepScatters[aoname] = ao
        newlo, newhi, cen = None, None, None
        if ao.dim() == 1:
            cen = [ p.x() for p in ao.points() ]
        elif ao.dim() == 2:
            cen = [ p.y() for p in ao.points() ]
        elif ao.dim() == 3:
            cen = [ p.z() for p in ao.points() ]
        if aoname in data:
            lo, hi = data[aoname]
            newlo = list(map(min, zip(lo, cen)))
            newhi = list(map(max, zip(hi, cen)))
        else:
            newlo = newhi = cen
        data[aoname] = [ newlo, newhi ]


def updateErrors(data, aos, useMidPoint):
    for aoname, ao in aos.items():
        all_lo, all_hi = data[aoname]
        for i,p in enumerate(ao.points()):
            lo = all_lo[i]
            hi = all_hi[i]
            if ao.dim() == 1:
                cen = p.x()
                if useMidPoint:
                    cen = 0.5 * ( lo + hi )
                    p.setX(cen)
                p.setXErrs( (cen - lo, hi - cen) )
            elif ao.dim() == 2:
                cen = p.y()
                if useMidPoint:
                    cen = 0.5 * ( lo + hi )
                    p.setY(cen)
                p.setYErrs( (cen - lo, hi - cen) )
            elif ao.dim() == 3:
                cen = p.z()
                if useMidPoint:
                    cen = 0.5 * ( lo + hi )
                    p.setZ(cen)
                p.setZErrs( (cen - lo, hi - cen) )

aos_out = { }
envdata = { }
aos_in = None
for i, filename in enumerate(args.INFILES):
    if args.VERBOSITY > 0:
        msg = "Adding data file {:s} [{:d}/{:d}]".format(filename, i+1, nfiles)
        sys.stdout.write(msg + "\n")

    del aos_in
    aos_in = yoda.read(filename, True, args.MATCH, args.UNMATCH)
    add2envelope(envdata, aos_in, None if i else aos_out)

if not args.REF_FILE is None:
    # Take the central values from the specified reference file
    if args.VERBOSITY > 0:
        msg = "Adding data file {:s} [{:d}/{:d}]".format(args.REF_FILE, nfiles, nfiles)
        sys.stdout.write(msg + "\n")

    del aos_in
    aos_in = yoda.read(args.REF_FILE, True, args.MATCH, args.UNMATCH)
    add2envelope(envdata, aos_in, aos_out)

useMidPoint = args.REF_FILE is None
updateErrors(envdata, aos_out, useMidPoint)

yoda.writeYODA(aos_out, args.OUTPUT_FILE)
