import ROOT
import ctypes
from array import array
import numpy as np

TMP_DIRECTORY = "./"

def create_fake_tuple():
    filename = TMP_DIRECTORY + "my_file.root"

    output_file = ROOT.TFile.Open( filename, "RECREATE")
    Pi_TRACK_CHI2NDOF = ctypes.c_double()
    my_tuple = ROOT.TTree("myTuple", "myTuple")
    chi2_branch = my_tuple.Branch("Pi_TRACK_CHI2NDOF", Pi_TRACK_CHI2NDOF, "Pi_TRACK_CHI2NDOF/D")

    for j in range(5000):
        Pi_TRACK_CHI2NDOF.value = np.random.uniform(0.0, 4.0)
        my_tuple.Fill()
    
    my_tuple.Write()
    output_file.Close()

    return filename


whatever = ROOT.TH1F("whatever", "", 5, 0.0, 4.0)
whatever.SetBinContent(1, 0.5)
whatever.SetBinContent(2, 0.01)
whatever.SetBinContent(3, 2.5)
whatever.SetBinContent(4, 3.5)
whatever.SetBinContent(5, 0.0)


unique = 0;
def loop_tree(tree, outputfile):
    global unique;

    unique += 1;

    sum_w = 0.0
    second_sum = 0.0;
    
    A = ctypes.c_double()
    #A = array('d', [0.]);

    tree.SetBranchAddress("Pi_TRACK_CHI2NDOF", A)

    weight = ctypes.c_double();
    friend_tree = ROOT.TTree("myFriendlyTree_{}".format(unique), "myFriendlyTree")
    friend_tree.Branch("weight", weight, "weight/D");

    outputfile.cd();
    for i in range(0, tree.GetEntries()):
            tree.GetEntry(i)
            sum_w += whatever.GetBinContent(whatever.FindBin(tree.Pi_TRACK_CHI2NDOF))
            weight.value = whatever.GetBinContent(whatever.FindBin(tree.Pi_TRACK_CHI2NDOF))
            friend_tree.Fill();
    
    tree.Show(2)
    tree.AddFriend( friend_tree );

    for i in range(0, tree.GetEntries()):
            tree.GetEntry(i)
            second_sum += whatever.GetBinContent(whatever.FindBin(tree.Pi_TRACK_CHI2NDOF))
    
    tree.ResetBranchAddresses();

    return (sum_w, second_sum, friend_tree)

if __name__ == "__main__":
    np.random.seed(5);
    filename = create_fake_tuple();
    rootfile = ROOT.TFile(filename,"READ")
    tree = rootfile.Get("myTuple")

    outfile = ROOT.TFile.Open(TMP_DIRECTORY + "my_working_tuples.root", "RECREATE")

    sumA = loop_tree(tree, outfile);
    sumB = loop_tree(tree, outfile);

    outfile.Close();

    print(sumA, sumB)