#! /usr/bin/env python

import time
start = time.time()

from ROOT import gRandom, TFile, TTree, TH1D
from array import array
import numpy # alternative to array

# Add branch to existing tree:
#    https://root.cern.ch/phpBB3/viewtopic.php?t=9661
#
# Write directories into root files:
#    https://root.cern.ch/root/html/tutorials/io/dirs.C.html


def writeTree():
    print ">>> Writing a tree."
    
    # make new root file with new tree
    f = TFile("tree.root", 'recreate')
    t = TTree("tree_name", "tree title")
    
    # create 1 dimensional float arrays as fill variables, in this way the float array serves
    # as a pointer which can be passed to the branch
    n = array('d',[0])
    u = array('d',[0])
    
    # using numpy instead of array class (note python's float datatype corresponds to C++ doubles):
    #n = numpy.zeros(1, dtype=float)
    #u = numpy.zeros(1, dtype=float)
    
    # create the branches and assign the fill-variables to them as doubles (D)
    t.Branch("normal", n, 'normal/D')
    t.Branch("uniform", u, 'uniform/D')
    
    # create some random numbers, assign them into the fill variables and call Fill()
    for i in xrange(10000):
        n[0] = gRandom.Gaus()
        u[0] = gRandom.Uniform()
        t.Fill()

    # write the tree into the output file and close the file
    f.Write()
    f.Close()



def addBranch():
    print ">>> Add branch."

    # reopen tree file
    f = TFile("tree.root",'update')
    t = f.Get("tree_name")
    
    # make new variable
    n = array('d',[0])
    #n = numpy.zeros(1, dtype=float)
    b = t.Branch("normal2", n, 'normal2/D')
    
    # fill the branch
    N = t.GetEntries()
    for i in xrange(N):
        t.GetEntry(i)
        n[0] = gRandom.Gaus(3,2)
        b.Fill()

    # overwrite the tree in the output file and close the file
    f.Write("",TFile.kOverwrite)
    f.Close()



def addHist():
    print ">>> Add histograms to a tree and one to a new branch."

    # reopen tree file
    f = TFile("tree.root",'update')
    t = f.Get("tree_name")
    
    # make new histogram
    h1 = TH1D("normal3","Haters be histin'",50,-2,6)

    # fill histogram
    for i in xrange(10000):
        h1.Fill(gRandom.Gaus(2,2))

    f.Write("",TFile.kOverwrite)
    f.Close()



def addHistToDirectory():
    print ">>> Add histogram to directory."

    # reopen tree file
    f = TFile("tree.root",'update')
    
    # make directory if it does not exist
    dir = f.GetDirectory("dir1")
    if not dir:
        print ">>>   created dir1"
        dir = f.mkdir("dir1")
    
    # make this directory the current one: everything you write will be saved here
    dir.cd()
    
    # declare histogram
    h1 = TH1D("normal4","Hist in folder",50,-1,9)
    
    # write histogram
    for i in xrange(10000):
        h1.Fill(gRandom.Gaus(4,3))

    f.Write("",TFile.kOverwrite)
    f.Close()



def addBranchToDirectory():
    print ">>> Add branch to directory."
    
    f = TFile("tree.root",'update')
    
    # make directory if it does not exist
    dir = f.GetDirectory("dir1")
    if not dir:
        print ">>>   created dir1"
        dir = f.mkdir("dir1")
    
    # make this directory the current one: everything you write will be saved here
    dir.cd()
    
    # make new tree with a branch
    t = TTree("tree2", "tree2")
    n = array('d',[0])
    #n = numpy.zeros(1, dtype=float)
    t.Branch("normal5", n, 'normal5/D')
    
    # fill tree
    for i in xrange(10000):
        n[0] = gRandom.Gaus(3,2)
        t.Fill()

    f.Write("",TFile.kOverwrite)
    f.Close()



def loopThroughTree():
    print ">>> Loop through tree."

    f = TFile("tree.root",'read')
    t = f.Get("tree_name")
    
    # loop over tree in the easiest way
    i = 0
    for event in t:
        if (i%2000) == 0:
            print ">>>   event %5i: n = %5.2f" % (i,t.normal)
            print ">>>   event %5i: u = %5.2f" % (i,t.uniform)
        i += 1
    print ">>>   event %5i: n = %5.2f" % (i,t.normal)
    print ">>>   event %5i: u = %5.2f" % (i,t.uniform)
    
    # alternatively
    # N = t.GetEntries()
    # for i in xrange(N):
    #     if (i%2000) == 0:
    #         t.GetEntry(i)
    #         print ">>>   event %5i: n = %5.2f" % (i,t.normal)
    #         print ">>>   event %5i: u = %5.2f" % (i,t.uniform)
    #     i += 1
    # print ">>>   event %5i: n = %5.2f" % (i,t.normal)
    # print ">>>   event %5i: u = %5.2f" % (i,t.uniform)
    
    f.Close()



def main():
    print "\n>>> Hello, world!"
    writeTree()
    addBranch()
    addBranchToDirectory()
    addHist()
    addHistToDirectory()
    loopThroughTree()



if __name__ == '__main__':
    main()
    end = time.time()
    print ">>>\n>>> The program lasted %.3f seconds.\n" % (end-start)