#!/usr/bin/env python

import argparse

import ROOT

parser = argparse.ArgumentParser(description='add a friend tree to a tree in another file')
parser.add_argument("infile", help="input file that contains the friend tree")
parser.add_argument("intree", help="name of the friend tree")
parser.add_argument("outfile", help="output file where the friend tree should be added")
parser.add_argument("outtree", help="name of the tree (in output file) to which the friend should be added")

args = parser.parse_args()

outfile = ROOT.TFile.Open(args.outfile, "UPDATE")
infile = ROOT.TFile.Open(args.infile)

for k in outfile.GetListOfKeys():
    if k.GetName() == args.intree:
        raise ValueError("Tree with name {} already exists in outputfile".format(args.intree))

outfile.cd()
outtree = outfile.Get(args.outtree)

if not outtree:
    raise KeyError("Tree {} not found in file {}".format(args.outtree, args.outfile))


if outtree.GetListOfFriends():
    for k in outtree.GetListOfFriends():
        if k.GetName() == args.intree:
            raise ValueError("Tree with name {} is already friend of {}".format(args.intree, args.outtree))

infile.cd()
intree = infile.Get(args.intree)

if not intree:
    raise KeyError("Tree {} not found in file {}".format(args.intree, args.infile))

# Add friend and write friend tree and original tree to outfile
outfile.cd()
outtree.AddFriend(intree)
outtree.Write(outtree.GetName())
outfile.cd()
clonetree = intree.CloneTree(-1, "fast")
clonetree.Write(intree.GetName())
infile.Close()
outfile.Close()