#! /usr/bin/env python
'''
Find FSRs on a file, check they are OK

First check the FileRecords exist in the Root file,
then check they can be read by Gaudi
Then check they give positive Lumi
'''
import sys
filename=''
isVerbose=False
year="2012"

def usage():
    print '''
Find FSRs on a file, check they are OK

First check the FileRecords exist in the Root file,
then check they can be read by Gaudi
Then check they give positive Lumi

ROOT files are assumed to have the key 'Refs' and have less than 10 keys
POOL files are assumed to have the keys '\#\#Shapes', '\#\#Links', '\#\#Params'

useage: '''+sys.argv[0]+''' [--verbose, -V, --debug, -D] [--help, -H] [--year=<year>] <filename>

option: --verbose, -V, --debug, -D, print more verbose information
        --help, -H, print this help and exit
        --year=<year>, "datatype" (year)
    '''

fsrDetected = { "TimeSpanFSR" : 0,
                "LumiFSRBeamCrossing" : 0,
                "LumiFSRBeam1" : 0,
                "LumiFSRBeam2" : 0,
                "LumiFSRNoBeam" : 0,
                "EventCountFSR" : 0
                }

def parseArgs():
    global isVerbose
    global filename
    global year
    for arg in sys.argv[1:]:
        if arg.lower() in ['--verbose', '-v', '--debug', '-d']:
            isVerbose=True
        elif arg.lower() in ['--help', '-h']:
            usage()
            sys.exit(0)
        elif len(arg)>7 and "--year=" in arg.lower()[:7]:
            year=arg.lower().split("=")[-1]
        elif filename!='':
            usage()
            raise ValueError, 'one file at a time please'
        else:
            filename=arg
    if filename=='':
        usage()
        raise ValueError, 'please supply a filename'

# a simple class with a write method
class WritableObject:
    def __init__(self):
        self.content = []
    def write(self, string):
        self.content.append(string)


def depth(top, mydepth):
    deepest=mydepth
    deepestFSR=0
    branches=top.GetListOfBranches()
    for branch in branches:
        #detect FSR depth
        flag=False
        for name in fsrDetected.keys():
            if name in branch.GetName():
                if isVerbose:
                    print "found", mydepth, branch.GetName()
                if mydepth>deepestFSR:
                    deepestFSR=mydepth
                flag=True
        #detect branch depth
        if not flag and "_FileRecords_" in branch.GetName() and not ".refs" in branch.GetName() and not ".links" in branch.GetName() and "_R." not in branch.GetName()[-3:]:
            if isVerbose:
                print "recurse to", mydepth+1, branch.GetName()
            adepth,adeepestFSR=depth(branch,mydepth+1)
            if adepth>deepest:
                deepest=adepth
            if adeepestFSR>deepestFSR:
                deepestFSR=adeepestFSR
    return deepest, deepestFSR


try:
    import ROOT
except:
    usage()
    raise ImportError, 'ROOT modules do not exist, call this script from within an Environment with ROOT'

try:
    from Configurables import DaVinci
except:
    usage()
    raise ImportError, 'DaVinci configurable does not exist, call this script from within a DaVinci environment'

parseArgs()

if not isVerbose:
    sys.stdout=WritableObject()
    sys.stderr=WritableObject()

thefile=ROOT.TFile.Open(filename)

if not isVerbose:
    sys.stdout=sys.__stdout__
    sys.stderr=sys.__stderr__

if not thefile or (hasattr(thefile, "isZombie") and thefile.isZombie) or (hasattr(thefile,"IsZombie") and thefile.IsZombie()):
    raise IOError, "the file could not be opened by Root"

keynames=[key.GetName() for key in thefile.GetListOfKeys()]

if isVerbose:
    print keynames

#three keys must be there in a pool file
poolkeys=['##Shapes', '##Links', '##Params']

isPool=(len([akey for akey in keynames if akey in poolkeys])==len(poolkeys))

#one key must be there in a Root file
rootkeys=['Refs']

isRoot=(len(keynames)<10 and len([akey for akey in keynames if akey in rootkeys])==len(rootkeys))

persistency=None

def atEnd():
    global thefile
    if thefile:
        thefile.Close()

persistency=None

if isRoot==isPool:
    print "Error: UNKNOWN file type. Fsrs can only exist on a root-like file"
    atEnd()
    sys.exit(1)
    
elif isRoot:
    persistency="ROOT"
elif isPool:
    persistency="POOL"

print "Success: "+persistency+" type detected"

hasFSLocation=(isRoot and "FileRecords" in keynames) or (isPool and len([akey for akey in keynames if "FileRecords" in akey])>0)

if not hasFSLocation:
    print "Error: No file record location on this file!"
    atEnd()
    sys.exit(2)

print "Success: File has File Records location"

if isRoot:
    top=ROOT.gDirectory.Get("FileRecords")
    adepth,anFSRdepth=depth(top,0)
    if adepth==0 and anFSRdepth==0:
        print "Success: FSR tree seems to be merged and cleaned up correctly."
    elif anFSRdepth<adepth:
        print "Error: empty FSR levels exist, merging without cleanup!"
    else:
        print "Success: FSR tree depth "+str(adepth)+" does not seem to be merged."
    if isVerbose or anFSRdepth<adepth:
        print "Info: FSR tree depth is "+str(adepth)
        print "Info: Deepest known FSR is at "+str(anFSRdepth)
    thefile.cd()

def gaudirunOptions(options):
    import commands
    result=commands.getoutput('gaudirun.py --option="'+options+'"')
    if isVerbose:
        print "---------------------- GAUDIRUN --------------------"
        print result
        print "-------------------- End GAUDIRUN ------------------"
    
    if "FATAL" in result and "LFC host has not been defined" in result:
        print "Error: Looks like you must run within a DaVinci environment with --use-grid and a valid grid proxy"
        #usage()
        atEnd()
        sys.exit(3)
    elif "FATAL" in result and "Bad credentials" in result:
        print "Error: Looks like you must run within a DaVinci environment with a valid grid proxy"
        #usage()
        atEnd()
        sys.exit(4)
    elif "FATAL" in result:
        print "Error: generic Gaudi error, check result (add --debug)"
        atEnd()
        sys.exit(8)
    
    if "Traceback (most recent call last):" in result:
        print "Error: generic python traceback, check traceback (add --debug)"
    
    return result


dumpoptions="from LumiAlgs.LumiFsrReaderConf import LumiFsrReaderConf as LumiFsrReader; LumiFsrReader().OutputLevel =  INFO; LumiFsrReader().inputFiles = ['%s'] ; LumiFsrReader().Persistency = '%s'; LumiFsrReader().EvtMax = 1; from Configurables import LHCbApp; LHCbApp().Persistency='%s';  from Configurables import CondDB, DDDBConf; CondDB().UseLatestTags=['%s']; DDDBConf(DataType='%s');"% (filename,persistency,persistency,year,year)

dumpResult=gaudirunOptions(dumpoptions)

totFSRs=0

for aline in dumpResult.strip().split('\n'):
    for fsr in fsrDetected:
        if fsr in aline:
            fsrDetected[fsr]+=1
            totFSRs+=1

missingFSRs=[fsr for fsr in fsrDetected if fsrDetected[fsr]==0]

if len(missingFSRs)==len(fsrDetected.keys()):
    print "Error: All FSRs are missing from this file"
elif len(missingFSRs):
    print "Error: Some FSRs are missing from this file"
else:
    print "Success: All required FSRs exist"

print fsrDetected
print " -- total = ", totFSRs

lumioptions="from Configurables import DaVinci, LHCbApp; DaVinci().DataType='%s'; from GaudiConf import IOHelper; IOHelper('%s','%s').inputFiles(['%s']); LHCbApp().Persistency='%s'; DaVinci().Lumi=True; from Configurables import CondDB, DDDBConf; CondDB().UseLatestTags=['%s']; DDDBConf(DataType='%s'); DaVinci().EvtMax=1;" % (year,persistency, persistency, filename,persistency,year,year)

lumiResult=gaudirunOptions(lumioptions)

alumi=None
for aline in lumiResult.strip().split('\n'):
    if "Integrated luminosity:" in aline:
        aline=aline.replace('  ',' ')
        words=aline.strip().split(' ')
        for word in range(len(words)):
            if words[word]=="luminosity:":
                alumi=words[word+1].strip()
                if '.' in alumi:
                    alumi=float(alumi)
                else:
                    alumi=int(alumi)
        break

if alumi is None:
    print "Error: Lumi printout not found!"
    atEnd()
    sys.exit(5)
elif alumi<0:
    print "Error: Negative lumi!", alumi
    atEnd()
    sys.exit(6)
elif alumi==0:
    print "Error: Lumi is Zero!"
    atEnd()
    sys.exit(7)
else:
    print "Success: Lumi is positive, ", alumi

print "Success: Everything seems OK with this file and its FSRs"

atEnd()
sys.exit(0)
