import ROOT as r
# Read data from NumPy arrays
# The column names in the RDataFrame are taken from the dictionary keys
import numpy as np
#Define some constants
mainFileName="file.root"
mainTreeName="MainTree"
friendFileName="friend_file.root"
friendTreeName="FriendTree"
#Create and save the main tree RDF
idx, val = np.arange(5), np.array([1]*5)
df = r.RDF.FromNumpy({"idx": idx, "val1": val})
df.Snapshot(mainTreeName, mainFileName)
#Create and save the friend tree RDF
idx, val = np.arange(5), np.array([2]*5)
friend_df = r.RDF.FromNumpy({"idx": idx, "val2": val})
friend_df.Snapshot(friendTreeName, friendFileName)
#make the trees in the two files friends
mainFile = r.TFile.Open(mainFileName, "UPDATE")
mainTree = mainFile.Get(mainTreeName)
friendFile = r.TFile.Open(friendFileName)
friendTree = friendFile.Get(friendTreeName)
friendTree.BuildIndex("idx")
mainTree.AddFriend(friendTree)
mainFile.Write()
mainFile.Close()
#check that friend_file.root is a friend of file.root
mainFile = r.TFile.Open(mainFileName)
mainTree = mainFile.Get(mainTreeName)
lof = mainTree.GetListOfFriends()
for i, elt in enumerate(lof):
    print(lof.At(i).GetTitle())
mainFile.Close()
#create a folder, move both files to it
import os
os.mkdir("folder")
os.system(f"mv {mainFileName} folder")
os.system(f"mv {friendFileName} folder")
#create a folder, move both files in it, check that the friend still exists and its path
mainFile = r.TFile.Open(f"folder/{mainFileName}")
mainTree = mainFile.Get(mainTreeName)
lof = mainTree.GetListOfFriends()
for i, elt in enumerate(lof):
    print(lof.At(i).GetTitle())
#can access branches in the main tree
mainTree.val1
#can't access branches in the friend tree as the path where it's been looked for is still the old one
mainTree.val2
#can we make this path relative rather than absolute so we can pass our trees with their friends to
#colleagues working on different storages?