import ROOT,math
from functions import *

class anaFile:
    def __init__(self,name):
        self.fname=name
        self.Rfile=ROOT.TFile(self.fname,'UPDATE')
        if not self.Rfile.IsOpen():
            self.Rfile.Close()
            print 'could not open file in update mode!'
            print 'opening read only'
            self.Rfile=ROOT.TFile(self.fname,'READ')
        self.zSlices=[]
        self.makeZSlices()
        self.info={}
        self.makeInfo()
        print 'created anafile from',name

    def __del__(self):
        self.Rfile.Close()
        
    def getTH1D(self):
        key=self.Rfile.FindKeyAny()
        #self.RFile.Get(key.)

    def makeZSlices(self):
        print "anaFile: creating z-slices"
        hslices=self.Rfile.FindObjectAny("hslices")
        if(type(hslices)!=ROOT.TH1D):
            print '***************************************************'
            print '***************************************************'
            print '***************************************************'
            print 'WARNING! could not find slices histo. using default'
            print '***************************************************'
            print '***************************************************'
            print '***************************************************'
            self.zSlices=( 12.0, 24.0, 36.0,  48.0, 60.0, 73.0 )
            return
        
        nslices=hslices.GetNbinsX()
        for i in range(1,nslices+1):
            self.zSlices.append(float(hslices.GetBinContent(i)))
            if hslices.GetEntries()>nslices:
                self.zSlices[-1]/=(hslices.GetEntries()/nslices)
            print "found slice",i,"at",self.zSlices[-1]
        if self.zSlices.count(0)>1:
            print "no slices found, setting to standard"
            self.zSlices=( 12.0, 24.0, 36.0,  48.0, 60.0, 73.0 ) # cosmic ne
#accessors
    def makeInfo(self):
        infod=self.Rfile.GetDirectory('info')
        if infod:
            for k in infod.GetListOfKeys():
                info=k.ReadObj()
                self.info[info.GetName()]=info.GetTitle()

    def getInfos(self):
        return self.info
    
    def getInfo(self,theinfo):
        if self.info.get(theinfo,None)!=None:
            return self.info[theinfo]
        else:
            print "this information is not available (",theinfo,")"
            return None
    def printInfos(self):
        for inf in self.info:
            print inf,self.info[inf]
        
    def getZSlices(self):
        return self.zSlices

    def getDriftSlices(self):
        dslices=[]
        for j in range(len(self.zSlices)):
            if j==0:
                offset=(abs(-62.5-self.zSlices[j])/2)
            elif j+1==len(self.zSlices) and self.zSlices[j]>10.78:
                offset =(72.8-abs(10.78-self.zSlices[j-1])/2)
            else:
                offset = offset+abs(self.zSlices[j-1]-self.zSlices[j])
        dslices.append(offset)
        return self.zSlices

    def getHres(self,coord):
        histos=[]
        for i in self.zSlices:
            print "StatsRes{0} Z: {1}".format(coord,i)
            histos.append(self.Rfile.FindObjectAny("StatsRes{0} Z:{1}".format(coord,i)))
            if type(histos[-1])!=ROOT.TH1D :
                histos[-1]=(self.Rfile.FindObjectAny("StatsRes{0} Z:{1}".format(coord,i)))
            if type(histos[-1])!=ROOT.TH1D :
                print "error while retrieveing Residual histo StatsRes{0} Z:{1}".format(coord,i)
                return None
        return histos

    def getMcHres(self,coord):
        histos=[]
        for i in self.zSlices:
            print "StatsMCRes{0} Z:{1}".format(coord,i)
            histos.append(self.Rfile.FindObjectAny("MCStatsRes{0} Z:{1}".format(coord,i)))
            if type(histos[-1])!=ROOT.TH1D :
                histos[-1]=(self.Rfile.FindObjectAny("MCStatsRes{0} Z:{1}".format(coord,i)))
            if type(histos[-1])!=ROOT.TH1D :
                print "error while retrieveing MC Residual histo {0}".format(i)
                return None
        return histos

    def getHXres(self):
        histos=[]
        for i in self.zSlices:
            print "StatsResX Z:"+str(i)
            histos.append(self.Rfile.FindObjectAny("StatsResX Z:"+str(i)))
            if type(histos[-1])!=ROOT.TH1D :
                histos[-1]=(self.Rfile.FindObjectAny("StatsResX Z:"+str(int(i))))
            if type(histos[-1])!=ROOT.TH1D :
                print "error while retrieveing Residual histo StatsResX Z:"+str(int(i))
                return None
        return histos

    def getMcHXres(self):
        histos=[]
        for i in self.zSlices:
            print "StatsMCResX Z:"+str(i)
            histos.append(self.Rfile.FindObjectAny("MCStatsResX Z:"+str(i)))
            if type(histos[-1])!=ROOT.TH1D :
                histos[-1]=(self.Rfile.FindObjectAny("MCStatsResX Z:"+str(int(i))))
            if type(histos[-1])!=ROOT.TH1D :
                print "error while retrieveing MC Residual histo"+str(int(i))
                return None
        return histos


    def getHClpercm(self):
        histos=[]
        for i in self.zSlices:
            histos.append(self.Rfile.FindObjectAny("hclusterpercm"+str(i)))
            if type(histos[-1])!=ROOT.TH1D :
                histos[-1]=(self.Rfile.FindObjectAny("hclusterpercm"+str(i)))
            if type(histos[-1])!=ROOT.TH1D :
                print "error while retrieveing hclusterpercm histo"+str(i)
                return None
        return histos

    def getHDigipercm(self):
        histos=[]
        for i in self.zSlices:
            histos.append(self.Rfile.FindObjectAny("hdigipercm"+str(i)))
            if type(histos[-1])!=ROOT.TH1D :
                histos[-1]=(self.Rfile.FindObjectAny("hdigipercm"+str(i)))
            if type(histos[-1])!=ROOT.TH1D :
                print "error while retrieveing hdigipercm histo"+str(i)
                return None
        return histos

    def getHTracklength(self):
        histos=[]
        for i in self.zSlices:
            histos.append(self.Rfile.FindObjectAny("htracklength"+str(i)))
            if type(histos[-1])!=ROOT.TH1D :
                histos[-1]=(self.Rfile.FindObjectAny("htracklength"+str(i)))
            if type(histos[-1])!=ROOT.TH1D :
                print "error while retrieveing htracklength histo"+str(i)
                return None
        return histos

    def getHClpertrk(self):
        histos=[]
        for i in self.zSlices:
            histos.append(self.Rfile.FindObjectAny("hclusterpertrack"+str(i)))
            if type(histos[-1])!=ROOT.TH1D :
                histos[-1]=(self.Rfile.FindObjectAny("hclusterpertrack"+str(i)))
            if type(histos[-1])!=ROOT.TH1D :
                print "error while retrieveing hclusterpertrack histo"+str(i)
                return None
        return histos

    def getHDigipertrk(self):
        histos=[]
        for i in self.zSlices:
            histos.append(self.Rfile.FindObjectAny("hdigipertrack"+str(i)))
            if type(histos[-1])!=ROOT.TH1D :
                histos[-1]=(self.Rfile.FindObjectAny("hdigipertrack"+str(i)))
            if type(histos[-1])!=ROOT.TH1D :
                print "error while retrieveing hdigipertrack histo"+str(i)
                return None
        return histos

    def getHClsize2D(self):
        histos=[]
        for i in self.zSlices:
            histos.append(self.Rfile.FindObjectAny("clustersize2D"+str(i)))
            if type(histos[-1])!=ROOT.TH1D :
                histos[-1]=(self.Rfile.FindObjectAny("clustersize2D"+str(i)))
            if type(histos[-1])!=ROOT.TH1D :
                print "error while retrieveing clustersize2D histo"+str(i)
                return None
        return histos

    def getHClsize3D(self):
        histos=[]
        for i in self.zSlices:
            histos.append(self.Rfile.FindObjectAny("clustersize3D"+str(i)))
            if type(histos[-1])!=ROOT.TH1D :
                histos[-1]=(self.Rfile.FindObjectAny("clustersize3D"+str(i)))
            if type(histos[-1])!=ROOT.TH1D :
                print "error while retrieveing clustersize3D histo"+str(i)
                return None
        return histos

    def getHChi2(self):
        histos=[]
        for i in self.zSlices:
            histos.append(self.Rfile.FindObjectAny("hchi2ndf"+str(i)))
            if type(histos[-1])!=ROOT.TH1D :
                histos[-1]=(self.Rfile.FindObjectAny("hchi2ndf"+str(i)))
            if type(histos[-1])!=ROOT.TH1D :
                print "error while retrieveing hchi2ndf histo"+str(i)
                return None
        return histos

    def getXYsliceHists(self):
        histos=[]
        for i in range(len(self.zSlices)):
            hist=self.Rfile.FindObjectAny("hxyresnorm"+str(self.zSlices[i]))
            if type(hist)==ROOT.TObject:
                print "error while getting histo: hxyresnorm"+str(self.zSlices[i])
                return None
            histos.append(hist)
        return histos

    def getSlicedHist(self,prefix):
        histos=[]
        for i in range(len(self.zSlices)):
            hist=self.Rfile.FindObjectAny(prefix+str(self.zSlices[i]))
            if type(hist)==ROOT.TObject:
                print"error while getting sliced histo:",prefix,str(self.zSlices[i])
                return None
            histos.append(hist)
        return histos

    def getPhiSlicedHist(self,prefix):
        histos=[]
        phiSlices=(30,60,90,120,150,180)
        for i in range(len(phiSlices)):
            hist=self.Rfile.FindObjectAny(prefix+str(phiSlices[i]))
            if type(hist)==ROOT.TObject:
                print "error while getting phi sliced histo:",prefix,str(phiSices[i])
                return None
            histos.append(hist)
        return histos

    def getHist(self,name,cd=False):
        if cd:
            key=self.Rfile.FindKeyAny(name)
            if key!=None:
                hist=key.ReadObj()
            else:
                print "Could not find histo:",name
                return None
        else:
            hist=self.Rfile.FindObjectAny(name)

        if type(hist)==ROOT.TObject:
            print "Could not find histo:",name
            return None
        return hist

    def get(self,name):
        return self.Rfile.Get(name)

    def getHistGroup(self,prefix,loop,suffix=''):
        loop=loop.split(';')
        histset=[]
        for l in loop:
            histset.append(self.Rfile.FindObjectAny('{0}{1}{2}'.format(prefix,l,suffix)))
        return histset

    def getHistsByPattern(self,pattern="",exclude="",folder=''):
        histset=[]
        if(exclude!=""):
            if exclude.find(" "):
                exclude=exclude.split(" ")
            else:
                exclude=[exclude]
        toscan=self.Rfile
        if folder!='':
            toscan=self.Rfile.GetDirectory(folder)
            for h in toscan.GetListOfKeys():
                do_append=True
                if( h.GetName().find(pattern)!=-1):
                    for exc in exclude:
                        if h.GetName().find(exc)!=-1:
                            do_append=False
                            continue
                    if do_append:
                        histset.append(self.getHist(h.GetName()))
        return histset
 
    def makeSliceGraph(self,hname,fitopt=0,fitfunc=0,force=False,ropt='c',rebin=-1):
        pullprojs={}
        pullgraphs={}

        #check if graphs exist
        hist=self.getHist(hname,True)
        graph=self.getHist(hname+"_Mean")
        if graph!=None and not force:
            print 'graph for {0} already existing and will not be recreated'.format(hname)
            return

        print 'generating graph for {0} with fitopt {1}, fitfunc {2}, ropt={3}'.format(hname,fitopt,fitfunc,ropt)
        projs=[]
        graph=getSlices(hist,projs,fitopt,fitfunc,ropt,rebin)
        currentDir=self.Rfile.CurrentDirectory().GetName()
        if self.Rfile.GetDirectory('{0}/graphs'.format(currentDir))==None:
            self.Rfile.mkdir('{0}/graphs'.format(currentDir))
        self.Rfile.cd('{0}/graphs'.format(currentDir))
        graph['Mean'].Write()
        graph['RMS'].Write()
        if( graph.get('Const',None)!=None):
            graph['Const'].Write()
        if self.Rfile.GetDirectory('{0}/projs'.format(currentDir))==None:
            self.Rfile.mkdir('{0}/projs'.format(currentDir))
        if self.Rfile.GetDirectory('{0}/projs/{1}_projs'.format(currentDir,hname))==None:
            self.Rfile.mkdir('{0}/projs/{1}_projs'.format(currentDir,hname))
        self.Rfile.cd('{0}/projs/{1}_projs'.format(currentDir,hname))
        for proj in projs:
            proj.Write()
            
        self.Rfile.cd()
              
    def makeSliceGraphGroup(self,prefix,loop,suffix,fitopt=0,fitfunc=0,force=False,ropt='c',rebin=-1):
        loop=loop.split(';')
        for l in loop:
            self.makeSliceGraph('{0}{1}{2}'.format(prefix,l,suffix),fitopt,fitfunc,force,ropt,rebin)

    def getSliceGraph(self,hname,fitopt=0,fitfunc=0,force=False,ropt='c',rebin=-1):
        graphs={}
        self.makeSliceGraph(hname,fitopt,fitfunc,force,ropt,rebin)
        graphnameadd=[]
        graphnameadd.append('Mean')
        if fitopt==0:
            graphnameadd.append('Sigma')
            graphnameadd.append('Const')
        else:
            graphnameadd.append('RMS')

        for gt in graphnameadd:
            graphs[gt]=self.getHist('{0}_{1}'.format(hname,gt))
        return graphs

    def getSliceGraphGroup(self,prefix,loop,suffix,fitopt=0,fitfunc=0,force=False,ropt='c',rebin=-1):
        graphs=[]
        self.makeSliceGraphGroup(prefix,loop,suffix,fitopt,fitfunc,force,ropt,rebin)
        loop=loop.split(';')
        graphnameadd=[]
        graphnameadd.append('Mean')
        if fitopt==0:
            graphnameadd.append('Sigma')
            graphnameadd.append('Const')
        else:
            graphnameadd.append('RMS')


        for l in loop:
            graphs.append({})
            for gt in graphnameadd:
                graphs[-1][gt]=self.getHist('{0}{1}{2}_{3}'.format(prefix,l,suffix,gt))
                
        return graphs

    def getDigiStdDef(self,force=False):
        graphs={}
        hvar=self.getHist('hdigiVariance')
        hnd=self.getHist('hndigis',True)

        graph=self.getHist('graph_Digivariance')
        if graph!=None and not force:
            print 'variance and diffusion graphs already existing and will not be recreated'
            graphs['Digivariance']=self.getHist('graph_Digivariance')
            graphs['Digidiffusion']=self.getHist('graph_Digidiffusion')
            graphs['DigivarDifFrac']=self.getHist('graph_DigivarDifFrac')
            return graphs

        currentDir=self.Rfile.CurrentDirectory().GetName()
        if self.Rfile.GetDirectory('{0}/graphs'.format(currentDir))==None:
            self.Rfile.mkdir('{0}/graphs'.format(currentDir))
        self.Rfile.cd('{0}/graphs'.format(currentDir))
        
        graphs['Digivariance']=ROOT.TGraph()
        graphs['Digivariance'].SetName('graph_Digivariance')
        graphs['Digidiffusion']=ROOT.TGraph()
        graphs['Digidiffusion'].SetName('graph_Digidiffusion')
        graphs['DigivarDifFrac']=ROOT.TGraph()
        graphs['DigivarDifFrac'].SetName('graph_DigivarDifFrac')

        for ivar in range(hvar.GetNbinsX()):
            num=hnd.GetBinContent(ivar)
            var=hvar.GetBinContent(ivar)
            if num!=0:
                var=var/num
            else:
                var=0
            sig=math.sqrt(var)
            graphs['Digivariance'].SetPoint(ivar,ivar+0.5,sig)
            elecDiff=float(self.info.get('Dt',-1))*math.sqrt(ivar+0.5)
            graphs['Digidiffusion'].SetPoint(ivar,ivar+0.5,elecDiff)
            if sig!=0:
                graphs['DigivarDifFrac'].SetPoint(ivar,ivar+0.5,elecDiff/sig)
            else:
                graphs['DigivarDifFrac'].SetPoint(ivar,ivar+0.5,0)

        for g in graphs:
            graphs[g].Write()

        self.Rfile.cd()
        return graphs