import sys
import ROOT as r

def getall(d, basepath="/"):
    "Generator function to recurse into a ROOT file/dir and yield (path, obj) pairs"
    keylist = d.GetListOfKeys()
    for key in keylist:
        kname = key.GetName()
        if key.IsFolder() and d.Get(kname).ClassName() == 'TDirectoryFile' :
            for i in getall(d.Get(kname), basepath+kname+"/"):
                yield i
        else:
            yield basepath+kname, d.Get(kname)

def getLeaves(t):
    keylist = t.GetListOfLeaves()
    for key in keylist:
        #print key
        yield t.GetName() + '/' + key.GetName(), key

def getAllLeaves(f):
    for k, o in getall(f):
        #print o.ClassName(), k
        yield k, o
        if 'TTree' in o.ClassName():
            for l, lo in getLeaves(o):
                #print lo.ClassName(), l
                yield k+'/'+l, lo


fname = str(sys.argv[1])
basepath = '/disk/users/elena/B2KShh/ntuples/jobs/'

f = r.TFile(basepath + fname, 'read')

# dirs = []
# print '%s :' % fname
# for k in f.GetListOfKeys():
#   print '\t %s :' % k.GetName()
#   dirs.append(k.GetName())
#   f.cd(k.GetName())
# f.cd()
#   for item in r.gDirectory.GetList():
#       print '\t\t %s' % item.GetName()

#for k, o in getall(f):
#    print o.ClassName(), k

leaves = []
for k, o in getAllLeaves(f):
    print o.ClassName(), k
    leaves.append((k, o))

polarity = [l for l in leaves if 'polarity'.lower() in l[0].lower()]
polarity += [l for l in leaves if 'pol'.lower() in l[0].lower()]
polarity += [l for l in leaves if 'magnet'.lower() in l[0].lower()]
polarity += [l for l in leaves if 'mag'.lower() in l[0].lower()]
#polarity += [l for l in leaves if 'h1'.lower() in l[0].lower()]