#!/usr/bin/env python

import ROOT

def add_friend(infile, intree, outfile, outtree):

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

    friend_name = outtree+"_friend_"+intree

    for k in root_outfile.GetListOfKeys():
        if k.GetName() == friend_name:
            print("Tree with name {} already exists in outputfile - writing new cycle!".format(intree))

    root_outfile.cd()
    root_outtree = root_outfile.Get(outtree)

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


    if root_outtree.GetListOfFriends():
        for k in root_outtree.GetListOfFriends():
            if k.GetName() == friend_name:
                print("Tree with name {} is already friend of {} - writing new cycle!".format(intree, outtree))

    root_infile.cd()
    root_intree = root_infile.Get(intree)

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

    # Add friend and write friend tree and original tree to outfile
    root_outfile.cd()
    clonetree = root_intree.CloneTree(-1, "fast")
    clonetree.SetName(friend_name)
    clonetree.SetTitle(friend_name)
    clonetree.Write(friend_name)
    root_outtree.AddFriend(clonetree)
    root_outtree.Write(root_outtree.GetName())
    root_infile.Close()
    root_outfile.Close()


if __name__ == "__main__":

    import argparse

    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()

    add_friend(args.infile, args.intree, args.outfile, args.outtree)