import ROOT,math, array

class pyFCNClusterCalib(ROOT.TPyMultiGenFunction):
    def __init__( self ):
        ROOT.TPyMultiGenFunction.__init__( self, self )
        self.first=False
        self.data=[]
        self.calibClusterArray=[]
        self.params=[]
        self.histbins=21
        self.fitOffset=True
        self.fitSlope=True
        self.fitGauss=True
        self.tuplefile=ROOT.TFile("fitTuple.root","recreate")
        self.chi2Tuple=ROOT.TNtuple("chi2Tuple",'','p1:p2:p3:chi2G:chi2BinOff:chi2Slope')
    
    def setFitOffset(self,opt):
        self.fitOffset=opt

    def setFitSlope(self,opt):
        self.fitSlope=opt

    def setFitGauss(self,opt):
        self.fitGauss=opt
        
    def NDim( self ):
        return 3
    
    def DoEval( self, x ):
        ret=0
        self.params=x   
        self.getFakeData()
        ret=self.getChi2()
        print '********************************************'
        print 'params=',x[0],x[1],x[2],ret
        print '********************************************'
        
        print
        return ret 
    
    def binData(self,data):
        binned=[]
        data_range=[min(data),max(data)]
        binwidth=(data_range[1]-data_range[0])/self.histbins
        #print 'binning data',data_range
        for i in range(self.histbins):
            binned.append(0)
        for i in range(len(data)):
            bin=int((data[i]-data_range[0])/binwidth)
            bin=min(self.histbins-1,bin)
            binned[bin]+=1
        return binned,binwidth,data_range
    
    def getGausExpected(self,x,height,mean=0):
        gval=height*math.exp(-((x-mean)**2)/2.)
        return gval
    
    def getChi2BinOff(self,mean):
        #print 'chi2 binoff'
        return (mean-0)**2
        
    def getChi2G(self,data,bw,mean,height):
        #print 'chi2 gauss'
        ret=0
        
        '''
        canv=ROOT.TCanvas('tmp','tmp',1000,500)
        h1=ROOT.TH1D('h1','data',40,0,0)
        h2=ROOT.TH1D('h2','expect',40,0,0)
        '''
        
        for i in range(len(data)):
            #print 'at histbin',i,'bin width=',bw,'binpos:',i*bw-mean,'mean=',mean
            expected=self.getGausExpected(i*bw,height,mean*bw)
            #print 'exp:',expected,'meas:',data[i]
            #h1.Fill(i,data[i])
            #h2.Fill(i,expected)
            diff=expected-data[i]
            chi2=diff**2 
            ret+=chi2
            #print 'diff=',diff,'chi2=',chi2
            #print
        '''
        canv.cd(1)
        h1.Draw()
        h2.SetLineColor(2)
        h2.Draw('same')
        canv.Update()
        u=raw_input('next')
        '''
        return ret
    
    def getChi2Slope(self,data):
        #print 'chi2slope'
        ret=0
        avgOffset=0
        for i in range(len(data)-1):
            slope=data[i]-data[i+1]
            #print 'slope:',slope
            ret+=abs(slope)
            #avgOffset+=abs(data[i])
        #avgOffset+=data[-1]
        #avgOffset/=len(data)
        #ret+=avgOffset**2
        return ret
    
    def getChi2(self):
        means=[]
        chi2s=[]
        for i in range(len(self.data)):            
            #print 'at zbin ',i,self.params
            binned,bw,dr=self.binData(self.data[i])
        
            means.append(0)
            for i in range(len(binned)):
                means[-1]+=i*binned[i]
            means[-1]/=sum(binned)
            
            if self.fitOffset:
                chi2s.append(self.getChi2BinOff(means[-1]*bw+dr[0]))
            else:
                chi2s.append(0)
        
            if self.fitGauss:
                height=0
                #for i in range(-1,1,1):
                #    height+=binned[ int(self.histbins/2)+i ]
                #height/=3.
                height=max(binned)
                chi2s.append(self.getChi2G(binned, bw, means[-1], height))
            else:
                chi2s.append(0)
        
        if self.fitSlope:
            chi2s.append(self.getChi2Slope(means))
        else:
            chi2s.append(0)
        
        self.fillTuple(chi2s)
        return sum(chi2s)
        
    def fillTuple(self,fillArray):
        if self.chi2Tuple!=None:
            self.chi2Tuple.Fill(array.array("f",fillArray))
     
    def writeTuple(self):
        self.tuplefile.cd()
        self.chi2Tuple.Write()
        self.tuplefile.Close()
    
    def getFakeData(self):
        self.data=[]
        for i in range(100):
            self.data.append([])
        for icc in range(len(self.calibClusterArray)):
            self.calibClusterArray[icc].setParams(self.params)
            tmparr=self.calibClusterArray[icc].produceFakePulls()
            for i in range(len(tmparr)):                
                self.data[i]+=tmparr[i]
        return self.data
        
    def plotData(self):
        hist=ROOT.TH2D('datahist','data',100,0,100,200,-500,500)
        for i in range(len(self.data)):
            for point in self.data[i]:
                hist.Fill(i,point)
                
    def setCalibClusterArray(self,cca):
        self.calibClusterArray=cca
        
    def setStartParams(self,params):
        self.params=params