#######################################################################
import user  # look for .pythonrc.py for user init
import ROOT
import PyCintex
import AthenaROOTAccess.transientTree

v = ROOT.TLorentzVector()
spartyjetdir = '/lapp_data/atlas/delsart/testARA/spartyjet/libs/'
ROOT.gSystem.Load(spartyjetdir+"libJetCore.so")
# ROOT.gSystem.Load(spartyjetdir+"libCDFJet.so")
ROOT.gSystem.Load(spartyjetdir+"libFastJet.so")
# ROOT.gSystem.Load(spartyjetdir+"libExtras.so")
# ROOT.gSystem.Load(spartyjetdir+"libATLASJet.so")


SJ = ROOT.SpartyJet
#######################################################################


## Ususal ARA set-up ##################################################
# Put your AOD file here.
aodFile = 'test.AOD.root'
f = ROOT.TFile.Open (aodFile)
assert f.IsOpen()


# Fill this in if you want to change the names of the transient branches.
branchNames = {}
#branchNames['ElectronCollection'] = 'ele'
#branchNames['PhotonCollection'] = 'gam'

tt = AthenaROOTAccess.transientTree.makeTree(f, branchNames = branchNames)

# tt is the transient tree "CollectionTree_trans" containing the (proxies) to
# all available transient data object in the file f.
# The original, persistent tree is declared as a friend
# of CollectionTree_trans, so that the transient tree will provide
# access to both transient data objects and to their persistent counterparts.

#######################################################################


## Set-up SpartyJet TTree reader ######################################
read = ROOT.ARASpartyJetReader_array_float("FastJet_")
read.set_var_names('eta','phi','p_T','e')
read.set_constituent_num_name('CaloCalTopoCluster_N') # the variableholding constits num in SpartyJet TTree

SJfile = ROOT.TFile('outN.root')
SJtree = SJfile.Get('myTree')

print 'Preparing SJ TTree'
read.SetBranchAddress(SJtree)
#######################################################################

def event_analysis(ievt):
   print 'analyzing ', ievt
   tt.GetEntry(ievt)
   clusters = tt.CaloCalTopoCluster
   fj_coll  = ROOT.JetCollection()
   # fill the collection :
   print '    Filling '
   read.fill_collection_withconstit(ievt, fj_coll, clusters)
   print 'jet 0 : ',fj_coll[0].e() ,fj_coll[0].eta() 
   print 'jet 1 : ',fj_coll[1].e() ,fj_coll[1].eta() 

##############################################################################
#
# Now you should be able to do things like:
#
event_analysis(0)