# 
# this is the main root interface
# modified to write different trees
# and branches with arrays
import sys, os, string, math, re
# import debugger
import   time
import pdb as deb
#
from array import array
import numpy as num 
#import ctypes as cc
#
import ROOT as R
from ROOT import gROOT, AddressOf





class myROOTmod:

    def __init__(self,keys,treeName,arr1,arr2):
	"""
	this is the init routine
	keys is the list of variables
	n is the number of different branches we will need to create
	"""
	self.keys=keys
	self.treeName=treeName # every packet will have its own tree name
	# some of the packes have arrays.
        self.nv = len(self.keys)  # number of variables which will also be number of 
	self.array1=arr1
	self.array2=arr2	
		     		     # branches
	self.Make_Tree()	# we just create a tree whenn we call root_manager	
	self.branchDict={}

    def Make_Tree(self):

        # create the tree
	self.title = " tree %s with %d variables" %(self.treeName,self.nv)
	self.tree=R.TTree(self.treeName,self.title)
	# Create a list of arrays o length 0
        self.var_address=[]
        for i in range(self.nv):
            self.var_address.append(array('L',[0]))

	# create a list
	self.variable_name=[]
	self.variable_type=[]
	for k in range(self.nv):
            	self.variable_name.append(self.keys[k])
            	self.variable_type.append(self.variable_name[k]+"/i")


        for k in range(self.nv):
            self.tree.Branch(self.variable_name[k],self.var_address[k], self.variable_type[k])

    def Append_Branch(self,keys):
    	"""
		add branches to existing tree
	"""
	self.keys1=keys
	self.variable_name1=[]
	self.variable_type1=[]
	self.var_address1=[]
        self.nv1 = len(self.keys1)  # number of variables which will also be number of 
        for i in range(self.nv1):
            self.var_address1.append(array('L',[0]))



	if (self.array1>0) and (self.array2 == 0):
		for k in range(self.nv1):
            		self.variable_name1.append(self.keys1[k])
			temp="[%d]"%(self.array1)
			print self.keys1, self.variable_name1[k]
            		self.variable_type1.append(self.variable_name1[k]+temp+"/i")


	create array of array
 	self.var_address1[k]=num.zeros((1,32),num.int64)
        for k in range(self.nv1):

            self.tree.Branch(self.variable_name1[k],self.var_address1[k], self.variable_type1[k])
    
	

    def Tree_Fill1(self,array1):
    	"""
	fill tree with array
	"""
	
	#for i in range(0,32):
	self.var_address1[0]=array1
	print self.tree.Fill()
 	
    
    
    

 
    def Write_Tree(self,filename):
    	"""
	opens a file and write the tree to a root file
	"""
	self.filename=filename
	self.root_file=R.TFile(self.filename,"RECREATE")
	self.root_file.Write()
	self.tree.Print()
	self.tree.Write()
	self.root_file.Write()
	print "writing the tree"
	self.root_file.Close()