import ROOT,math,argparse

def getGConst(charge,sigma):
    return charge/(g_sigma*ROOT.TMath.Sqrt(ROOT.TMath.Pi()))

def fillTuple(ptuple,sig,con,N=5000,binw=0,swidth=10):
    gauss=ROOT.TF1('f','gaus')
    gauss.SetParameter(0,con)
    gauss.SetParameter(1,0)
    gauss.SetParameter(2,sig)
    i=0
    posarr=[]
    if binw==-2:
        N=int(2*swidth*sig)-1
        
    
    while i<N:
        if binw==-2:
            pos=int(i-swidth*sig)
        else:
            pos=ROOT.gRandom.Uniform(-swidth*sig,swidth*sig)
            
        if binw>0:
            pos=int(pos/binw)*binw

        if binw==-1 and ( pos in posarr ):
            continue

        val=gauss.Eval(pos)
        if val<1e-2 and binw>0:
            continue

        posarr.append(pos)
        i+=1

        ptuple.Fill(ROOT.Double(pos),ROOT.Double(val))

def rebinTuple(intuple,outtuple,binw):    
    for e in intuple:
        pos=e.x
        val=e.amp
        newpos=pos
        if binw>0:
            newpos=int(pos/binw)*binw
        outtuple.Fill(newpos,val)

def smearTuple(intuple,outtuple,binw):
    bincounts={}
    #find bins to smear:
    for e in intuple:
        pos=e.x
        counts_in_bin=bincounts.get(str(pos),0)+1
        bincounts[str(pos)]=counts_in_bin
        
    for e in intuple:
        pos=e.x
        val=e.amp
        if bincounts.get(str(pos),1)!=1:
            newpos=pos+ROOT.gRandom.Uniform(-binw/2.,binw/2.)
        else:
            newpos=pos
        outtuple.Fill(newpos,val)

def fitRooFit(ptuple,opt=''):
    opts=opt.split(' ')
    xmin=ptuple.GetMinimum('x')
    xmax=ptuple.GetMaximum('x')

    rooX=ROOT.RooRealVar('x','',0.,xmin,xmax)
    rooA=ROOT.RooRealVar('amp','',0.)
    varset=ROOT.RooArgSet('varset')
    varset.add(rooX)
    varset.add(rooA)

    fitMean=ROOT.RooRealVar('fitMean','',0.,xmin,xmax)
    fitSigma=ROOT.RooRealVar('fitSigma','',g_sigma,0,xmax-xmin)
    fitgauss=ROOT.RooGaussian("fitgauss",'roogauss',rooX,fitMean,fitSigma)


    dataset=ROOT.RooDataSet('dataset','pointstuple',ptuple,varset,'','amp')

    if opt.find('plot')!=-1:
        frame=rooX.frame(ROOT.RooFit.Title('title'))
        dataset.plotOn(frame,ROOT.RooFit.Binning(10000),ROOT.RooFit.DataError(ROOT.RooAbsData.SumW2))
        frame.Draw()

    sig=0
    mean=0
    suma=0
    suma2=0
    sigoff=0
    if opt.find('sigoff'):
        for o in opts:
            if o.find('sigoff')!=-1:
                sigoff=float(o[6:len(o)])


    bincounts={}
    for e in ptuple:
        pos=e.x
        mean+=e.x*e.amp
        suma+=e.amp
        suma2+=e.amp**2
        counts_in_bin=bincounts.get(str(pos),0)+1
        bincounts[str(pos)]=counts_in_bin

    mean/=suma
    for e in ptuple:
        res=e.x-mean
        pos=e.x
        #decide=ROOT.gRandom.Uniform(float(-1./bincounts[str(pos)]),float(1./bincounts[str(pos)]))
        if sigoff!=0:
            decide=ROOT.gRandom.Gaus(0,sigoff/2)
            res+=decide
        '''
        if res<0:
            res-=decide*sigoff/2.#(1.+ bincounts[str(pos)])
        else:
            res+=decide*sigoff/2.#(1.+  bincounts[str(pos)])
           '''
        
        sig+=e.amp*(res)**2
    
    sig/=suma
    #sig*=(suma)/(suma**2-suma2)
    #sig+=sigoff
    sig=math.sqrt(sig)

    fitSigma.setVal(sig)
    #fitSigma.setVal(100)
    fitMean.setVal(mean)
    if args.plot:
        print 'the found sigma is: {0}'.format(fitSigma.getVal())

    if opt.find('fixsig')!=-1:
        fitSigma.setConstant(True)

    if opt.find('sumw2')!=-1:

        result=fitgauss.fitTo(dataset,
                              ROOT.RooFit.SumW2Error(True),
                              ROOT.RooFit.Save(),
                              ROOT.RooFit.Verbose(False),
                              ROOT.RooFit.PrintLevel(-1),
                              ROOT.RooFit.PrintEvalErrors(0),
                              ROOT.RooFit.Warnings(False)
                              )
    else:
        result=fitgauss.fitTo(dataset,
                              ROOT.RooFit.SumW2Error(False),
                              ROOT.RooFit.Save())
        print 'no sumw2!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!'

    if opt.find('plot')!=-1:
        fitgauss.plotOn(frame,ROOT.RooFit.Range("Full"),ROOT.RooFit.LineStyle(ROOT.kDashed),ROOT.RooFit.LineColor(ROOT.kRed)) 
        frame.Draw()

    error=fitMean.getError()
    mean=fitMean.getVal()
    sigma=fitSigma.getVal()
    return_value=()
    return_value=(error,)
    if opt.find('getMean')!=-1:
        return_value = return_value + (mean,)
    if opt.find('getSigma')!=-1:
        return_value = return_value + (sigma,)
    
    return return_value

def fitTGraph(ptuple,opt='',):
    xmin=ptuple.GetMinimum('x')
    xmax=ptuple.GetMaximum('x')
    maxamp=ptuple.GetMaximum('amp')

    amperr=0
    if opt.find('ampError')!=-1:
        amperr=0.05

    graph=ROOT.TGraphErrors()
    counter=-1
    mean=0
    amp=0
    for e in ptuple:
        counter+=1
        graph.SetPoint(counter,e.x,e.amp)        
        graph.SetPointError(counter,0,e.amp*amperr)
        mean+=e.x*e.amp
        amp+=e.amp
    mean/=amp
    #second looop for variance
    var=0
   
    for e in ptuple:
        var+=e.amp*(e.x-mean)**2
    var/=amp
    sig=math.sqrt(var)
    tfitgauss=ROOT.TF1('fitgauss','gaus',xmin,xmax)
    tfitgauss.SetParameter(0,maxamp)
    tfitgauss.SetParLimits(0,0,1.1*maxamp)
    tfitgauss.SetParameter(1,mean)
    tfitgauss.SetParLimits(1,mean-sig,mean+sig)
    tfitgauss.SetParameter(2,sig)
    tfitgauss.SetParLimits(2,0.5*sig,2*sig)
    graph.Fit(tfitgauss)
    graph.SetMarkerStyle(20)
    if opt.find('plot')!=-1:
        graph.Draw('AP')
        tfitgauss.Draw('same')
    return tfitgauss.GetParError(1),graph,tfitgauss

def fittuple(ptuple,opt=''):
    print 'starting fittuple:'
    xmin=ptuple.GetMinimum('x')
    xmax=ptuple.GetMaximum('x')
    maxamp=ptuple.GetMaximum('amp')

    amperr=0
    if opt.find('ampError')!=-1:
        amperr=0.05

    counter=-1
    mean=0
    amp=0
    for e in ptuple:
        counter+=1
        mean+=e.x*e.amp
        amp+=e.amp
    mean/=amp
    #second looop for variance
    var=0
    for e in ptuple:
        var+=e.amp*(e.x-mean)**2
    var/=amp
    sig=math.sqrt(var)

    tupfitgauss=ROOT.TF1('tuplefitgauss','gaus',xmin,xmax)
    tupfitgauss.SetParameter(0,maxamp)
    #tupfitgauss.SetParameter(0,1)
    #tupfitgauss.FixParameter(0,1)
    tupfitgauss.SetParLimits(0,0,1.1*maxamp)
    tupfitgauss.SetParameter(1,mean)
    #tupfitgauss.SetParLimits(1,mean-sig,mean+sig)
    tupfitgauss.SetParameter(2,sig)
    #tupfitgauss.SetParLimits(2,0.5*sig,2*sig)
    ptuple.UnbinnedFit(tupfitgauss.GetName(),'x')
    return tupfitgauss.GetParError(1),tupfitgauss

parser=argparse.ArgumentParser(description='Mickey Mouse MC for error on Mean')
parser.add_argument('--diffusion',help='analyze diffusion effects',action='store_true')
parser.add_argument('--binwidth',help='analyze binwidth effects',action='store_true')
parser.add_argument('--plot',help='plot in between',action='store_true')
parser.add_argument('--events',help='number of events',type=int,default=50)
args=parser.parse_args()

ROOT.gROOT.ProcessLine(".x rootlogon.C") 
ROOT.RooMsgService.instance().setGlobalKillBelow(ROOT.RooFit.ERROR)
ROOT.RooMsgService.instance().setSilentMode(True)


charge=1000
g_sigma=10
g_const=getGConst(charge,g_sigma)


pointstuple=ROOT.TNtuple('points','','x:amp')




canv=ROOT.TCanvas('Canvas','Fitting',2000,500)
canv.Divide(2,1)
errorVsSigCanv=ROOT.TCanvas()

fillTuple(pointstuple,g_sigma,g_const,100,-2)
canv.cd(1)
fitRooFit(pointstuple,'plot')

canv.cd(2)
e,g,f=fitTGraph(pointstuple)
g.Draw('AP')
f.Draw('same')
canv.Update()

del pointstuple

if args.diffusion:
    errorVsSigRoo=ROOT.TGraph()
    errorVsSigTGr=ROOT.TGraph()
    errorVsSigTGrErr=ROOT.TGraph()
    errorVsSigTGrErrErr=ROOT.TGraph()
    errorVsSigRooSumw2=ROOT.TGraph()
    thismin=999
    thismax=0
    for isig in range(0,70,10):
        canv.Clear()
        canv.Divide(5,1)

        g_sigma=10*math.sqrt(100*(isig+1))
        g_const=getGConst(charge,g_sigma)
        pointstuple=ROOT.TNtuple('points','','x:amp')
        fillTuple(pointstuple,g_sigma,g_const,1000,1)

        canv.cd(1)
        rooerr,=fitRooFit(pointstuple,'plot')

        canv.cd(2)
        rooerrSumw2,=fitRooFit(pointstuple,'plot sumw2')

        canv.cd(3)
        tgrerr,g2,f2=fitTGraph(pointstuple)
        g2.Draw('AP')
        f2.Draw('same')

        canv.cd(4)
        tgrerrErr,g3,f3=fitTGraph(pointstuple,'ampError')
        g3.Draw('AP')
        f3.Draw('same')

        canv.cd(5)
        tupErr,f4=fittuple(pointstuple)
        pointstuple.Draw('amp:x')
        f4.Draw('same')

        print rooerr
        errorVsSigRoo.SetPoint(isig,g_sigma,rooerr)
        errorVsSigRooSumw2.SetPoint(isig,g_sigma,rooerrSumw2)
        errorVsSigTGr.SetPoint(isig,g_sigma,tgrerr)
        errorVsSigTGrErr.SetPoint(isig,g_sigma,tgrerrErr)


        thismin=min([thismin,rooerr,tgrerr,tgrerrErr,rooerrSumw2])
        thismax=max([thismax,rooerr,tgrerr,tgrerrErr,rooerrSumw2])
        del pointstuple
        print 'at sigma',isig,g_sigma,g_const
        canv.Update()
        #u=raw_input('next sigma')

    errorVsSigCanv.cd()
    errorVsSigRoo.SetMarkerStyle(20)
    errorVsSigTGr.SetMarkerStyle(20)
    errorVsSigTGrErr.SetMarkerStyle(20)
    errorVsSigRooSumw2.SetMarkerStyle(20)
    errorVsSigRoo.SetMarkerColor(ROOT.kMagenta)
    errorVsSigTGr.SetMarkerColor(ROOT.kCyan)
    errorVsSigTGrErr.SetMarkerColor(ROOT.kSpring)
    errorVsSigRooSumw2.SetMarkerColor(ROOT.kOrange+8)
    errorVsSigRoo.SetMaximum(thismax*1.1)
    errorVsSigRoo.SetMinimum(thismin*0.9)

    errorVsSigRoo.GetXaxis().SetTitle('sigma')
    errorVsSigRoo.GetYaxis().SetTitle('Error on mean')
    errorVsSigRoo.SetTitle("constant charge, constant N")
    errorVsSigRoo.Draw('AP')
    errorVsSigTGr.Draw('P')
    errorVsSigTGrErr.Draw('P')
    errorVsSigRooSumw2.Draw('P')
    leg=ROOT.TLegend(0.1,0.7,0.4,0.9)
    leg.AddEntry(errorVsSigRoo,'RooFit no sumW2','lp')
    leg.AddEntry(errorVsSigRooSumw2,'RooFit with sumW2','lp')
    leg.AddEntry(errorVsSigTGr,'TGraph without errors','lp')
    leg.AddEntry(errorVsSigTGrErr,'TGraph with 5% error on amp','lp')
    leg.SetFillColor(0)
    leg.Draw()
    errorVsSigCanv.Update()

if args.binwidth:
    binstep=5
    maxbinw=200
    
    errhistmax=200

    errorgraph=ROOT.TGraph()
    errorhist=ROOT.TH2D('errorhist','Error of Mean Vs binwidth',int(maxbinw/binstep),0,maxbinw,100,0,errhistmax)
    meanhist=ROOT.TH2D('meanhist','Residual of Mean Vs binwidth',int(maxbinw/binstep),0,maxbinw,1000,-100,100)
    sigmahist=ROOT.TH2D('sigmahist','Sigma-Gsigma Vs binwidth',int(maxbinw/binstep),0,maxbinw,100,-100,100)
    errorhist_s=ROOT.TH2D('errorhist_s','Error of Mean Vs binwidth smeared',int(maxbinw/binstep),0,maxbinw,100,0,errhistmax)
    meanhist_s=ROOT.TH2D('meanhist_s','Residual of Mean Vs binwidth smeared',int(maxbinw/binstep),0,maxbinw,1000,-100,100)
    sigmahist_s=ROOT.TH2D('sigmahist_s','Sigma-Gsigma Vs binwidth smeared',int(maxbinw/binstep),0,maxbinw,100,-100,100)

    bcanv=ROOT.TCanvas('bcanv','binwidth',1500,1500)
    bcanv.DivideSquare(int(maxbinw/binstep))
    bcanv_s=ROOT.TCanvas('bcanv_s','binwidth_s',1500,1500)
    bcanv_s.DivideSquare(int(maxbinw/binstep))
    bcanv2=ROOT.TCanvas('binwidth','binwidth',1000,1000)
    bcanv2_s=ROOT.TCanvas('binwidth_s','binwidth_s',1000,1000)

    charge=1000
    g_sigma=100
    g_const=getGConst(charge,g_sigma)

    for iset in range(args.events):
        print 'start a set',iset
        pointstuple=ROOT.TNtuple('points','','x:amp')
        #npoints=ROOT.gRandom.Uniform(5,15)
        npoints=100
        fillTuple(pointstuple,g_sigma,g_const,npoints,-1,4)
        
        rebin_tuples=[]
        smear_tuples=[]
        ibin=-1
        for binwidth in range(0,maxbinw,binstep):
            ibin+=1
            rebin_tuples.append(ROOT.TNtuple('points_rebin{0}'.format(binwidth),'','x:amp'))
            smear_tuples.append(ROOT.TNtuple('points_smeared{0}'.format(binwidth),'','x:amp'))
            rebinTuple(pointstuple,rebin_tuples[-1],binwidth)
            smearTuple(rebin_tuples[-1],smear_tuples[-1],binwidth)
            opt='sumw2 getMean getSigma fixsig sigoff0'.format(binwidth)
            if args.plot:
                opt='plot '+opt
                
            bcanv.cd(ibin+1)
            rooerrSumw2,mean,sigma=fitRooFit(rebin_tuples[-1],opt)
            bcanv_s.cd(ibin+1)
            rooerrSumw2_s,mean_s,sigma_s=fitRooFit(smear_tuples[-1],opt)
            
            
        
            errorgraph.SetPoint(ibin+iset*maxbinw,binwidth,rooerrSumw2)

            errorhist.Fill(binwidth,rooerrSumw2)
            meanhist.Fill(binwidth,mean)
            sigmahist.Fill(binwidth,sigma-g_sigma)
            
            errorhist_s.Fill(binwidth,rooerrSumw2_s)
            meanhist_s.Fill(binwidth,mean_s)
            sigmahist_s.Fill(binwidth,sigma_s-g_sigma)
            
            if args.plot:
                print 'at binwidth: {0}'.format(binwidth)
                bcanv.Update()
                bcanv_s.Update()
            
        if args.plot:
            u=raw_input('next')
        del pointstuple
        del rebin_tuples
   
    bcanv2.Divide(2,2)
    bcanv2.cd(1)
    sigmahist.Draw('colz')
    #errorgraph.SetMarkerStyle(20)
    #errorgraph.Draw('AP')
    bcanv2.cd(2)
    errorhist.Draw('colz')
    bcanv2.cd(3)
    #meanhist.Draw('colz')
    prof_sig=sigmahist.ProfileX(sigmahist.GetName()+'_profile')
    prof_sig.Draw()
    bcanv2.cd(4)
    prof=errorhist.ProfileX(errorhist.GetName()+'_profile')
    prof.Draw()
    bcanv2.Update()

    bcanv2_s.Divide(2,2)
    bcanv2_s.cd(1)
    #errorgraph.SetMarkerStyle(20)
    #errorgraph.Draw('AP')
    sigmahist_s.Draw('colz')
    bcanv2_s.cd(2)
    errorhist_s.Draw('colz')
    bcanv2_s.cd(3)
    #meanhist_s.Draw('colz')
    prof_sig_s=sigmahist_s.ProfileX(sigmahist_s.GetName()+'_profile')
    prof_sig_s.Draw()
    bcanv2_s.cd(4)
    prof_s=errorhist_s.ProfileX(errorhist_s.GetName()+'_profile_s')
    prof_s.Draw()
    bcanv2_s.Update()
u='w'
while u!='q':
    u=raw_input('done?')