#!/usr/bin/env python
# Copyright 2012 the rootpy developers
# distributed under the terms of the GNU General Public License

from rootpy.extern.argparse import ArgumentParser

parser = ArgumentParser()
parser.add_argument('-n', type=int, default=500,
        help="number of files to merge at once")
parser.add_argument('dest')
parser.add_argument('files', nargs='+')
args = parser.parse_args()

import sys
import os
from subprocess import call
from tempfile import mkstemp

tmppath = None
tmpdest = None

if os.path.exists(args.dest):
    sys.exit("destination file %s exists" % destination)

if len(args.files) == 1:
    # simple copy
    import shutil
    shutil.copy(args.files[0], args.dest)
    sys.exit(0)

if args.n < 2:
    sys.exit("you must merge at least two files at once")

while args.files:
    if tmppath:
        files = args.files[:args.n - 1]
        args.files = args.files[args.n - 1:]
        files.insert(0, tmppath)
    else:
        files = args.files[:args.n]
        args.files = args.files[args.n:]
    print "merging %i files... (%i remaining)" % (len(files), len(args.files))
    dest = args.dest
    hadd_args = []
    if args.files:
        tmpfd, tmpdest = mkstemp(suffix=".root")
        dest = tmpdest
        os.close(tmpfd)
        hadd_args.append("-f")
    files.insert(0, dest)
    call(["hadd"] + hadd_args + files)
    if tmppath:
        os.unlink(tmppath)
    tmppath = tmpdest
