import ROOT as r 
from array import array


#Set up two trees
print("Setting up trees.")
tree_1 = r.TTree("tree_1", "tree_1")
tree_2 = r.TTree("tree_2", "tree_2")

x_1 = array('i', [0 for x in range(0,100)])
x_2 = array('i', [0 for x in range(0,100)])

tree_1.Branch("x", x_1, "x/I")
tree_2.Branch("x", x_2, "x/I")

#Fill trees
for x in range(0,200):
    x_1[0] = x
    
    tree_1.Fill()
for x in range(0,100):
    x_2[0] = x+50
    tree_2.Fill()

canvas = r.TCanvas("canvas", "canvas",800, 600)

print("Add as a friend.")
#Add as friend 
tree_2.BuildIndex("x")
tree_1.AddFriend(tree_2, "tree_2")

print("DF method.")
#DF method
df = r.RDataFrame(tree_1)
model = r.RDF.TH2DModel("df_hist", "", 200, 0, 200, 200, 0, 200)
hist = df.Histo2D(model, "x", "tree_2.x")
hist.Draw("colz")
canvas.Print("test_df_out.png")

canvas.Clear()

print("Normal mehtod.")
#Normal tree method
tree_hist = r.TH2D("tree_hist", "", 200, 0, 200, 200, 0, 200)
tree_1.Draw("x:tree_2.x>>tree_hist", "", "colz")
tree_1.Print()
canvas.Print("test_tree_out.png")