import ROOT,time,math
ROOT.gROOT.ProcessLine(".x rootlogon.C") 

c1=ROOT.TCanvas("c","roofit gauss test",1200,600)
c1.Divide(3,1)
#c2=ROOT.TCanvas()
#c3=ROOT.TCanvas()
rand=ROOT.TRandom3()
rand.SetSeed(int(time.time()))

prod_mean=rand.Uniform(-2,2)
prod_sigma=rand.Uniform(1,5)
prod_amp=rand.Uniform(0,200)

prod_gaus=ROOT.TF1('prod_gaus','gaus',-12,12)
prod_gaus.SetParameter(0,prod_amp)
prod_gaus.SetParameter(1,prod_mean)
prod_gaus.SetParameter(2,prod_sigma)
tuple1=ROOT.TNtuple('2d gauss','','x:amp')
tuple2=ROOT.TNtuple('2d gauss x smeared','','x:amp')
tuple3=ROOT.TNtuple('gauss x smeared','','x')

results=[]
for i in range(100):
    x=rand.Uniform(-12,12)
    amp=prod_gaus.Eval(x)
    tuple2.Fill(x+prod_sigma*rand.Gaus(),amp)
    amp+=prod_amp*0.1*rand.Gaus()
    while amp<0:
        amp=prod_gaus.Eval(x)+prod_amp*0.1*rand.Gaus()
    tuple1.Fill(x,amp)
    

for i in range(1000):
    x=prod_gaus.GetRandom()    
    tuple3.Fill(x+prod_sigma*rand.Gaus())

x=ROOT.RooRealVar('x','x',-20,20)
amp=ROOT.RooRealVar('amp','amp',0,120)
gamp=ROOT.RooRealVar('gamp','gamp',80,0,120)
mu=ROOT.RooRealVar('mu','mu',4,-20,20)
sigma=ROOT.RooRealVar('sigma','sigma',5,-100,100)
sigma2=ROOT.RooRealVar('sigma2','sigma2',4,-10,10)

varset1=ROOT.RooArgSet()
varset1.add(x)
varset1.add(amp)


dataset1=ROOT.RooDataSet('dataset1','dataset1',tuple1,varset1)

gausvar=ROOT.RooGaussVar('gausvar','gausvar',x,gamp,mu,sigma)
pdf1=ROOT.RooGaussian('gaus1','gaus1',amp,gausvar,sigma2)


start=time.time()
pdf1.fitTo(dataset1)
results.append(['pdf1 mean',mu.getVal()])
results.append(['pdf1 sigma',sigma.getVal()])
results.append(['pdf1 mean error',mu.getError()])
results.append(['pdf1 sigma error',sigma.getError()])

print '1d fittime',time.time()-start
print 'mc mean={0} amp={1} sigma={2}'.format(prod_mean,prod_amp,prod_sigma)
h1=dataset1.createHistogram(x,amp,1000,1000)
h1_g=gausvar.createHistogram('x',1000,1000)
c1.cd(1)
h1.Draw()
h1_g.Draw('same')

dataset2=ROOT.RooDataSet('dataset2','dataset2',tuple2,varset1,'','amp')
sigma3=ROOT.RooRealVar('sigma3','sigma3',4,-100,100)
mean3=ROOT.RooRealVar('mean3','mean3',0,-100,100)

sig_res=ROOT.RooRealVar('sig_res','sigma of resolution',2,0,10)
m_res=ROOT.RooRealVar('m_res','mean of resolution',0)
s=ROOT.RooRealVar('s','resolution',2,0,20)
resolution=ROOT.RooGaussModel('gauss_res','gauss resolution',x,x,s,sig_res)
pdf2=ROOT.RooGaussian('gaus2','gaus2',x,mean3,sigma3)


pdf2_res=ROOT.RooProdPdf('pdf2_res','pdf2 with resolution',ROOT.RooArgSet(pdf2),ROOT.RooFit.Conditional(ROOT.RooArgSet(resolution),ROOT.RooArgSet(x)))

pdf2.fitTo(dataset2,ROOT.RooFit.SumW2Error(True),ROOT.RooFit.Save())
results.append(['pdf2 mean',mean3.getVal()])
results.append(['pdf2 sigma',sigma3.getVal()])
results.append(['pdf2 mean error',mean3.getError()])
results.append(['pdf2 sigma error',sigma3.getError()])
pdf2.fitTo(dataset2,ROOT.RooFit.SumW2Error(True),ROOT.RooFit.Save(),ROOT.RooFit.ConditionalObservables(ROOT.RooArgSet(x)))
results.append(['pdf2_res mean',mean3.getVal()])
results.append(['pdf2_res sigma',sigma3.getVal()])
results.append(['pdf2_res mean error',mean3.getError()])
results.append(['pdf2_res sigma error',sigma3.getError()])

print 'mc mean={0} amp={1} sigma={2}'.format(prod_mean,prod_amp,prod_sigma)
c1.cd(2)
h2=x.frame()
dataset2.plotOn(h2,ROOT.RooFit.DataError(ROOT.RooAbsData.SumW2))
pdf2_res.plotOn(h2)
pdf2.plotOn(h2,ROOT.RooFit.LineColor(ROOT.kRed))
h2.Draw()


sigma4=ROOT.RooRealVar('sigma4','sigma4',4,-100,100)
mean4=ROOT.RooRealVar('mean4','mean4',0,-100,100)
pdf3=ROOT.RooGaussian('gaus2','gaus2',x,mean4,sigma4)
pdf3_res=ROOT.RooProdPdf('pdf3_res','pdf3 with resolution',ROOT.RooArgSet(pdf3),ROOT.RooFit.Conditional(ROOT.RooArgSet(resolution),ROOT.RooArgSet(x)))

c1.cd(3)
varset3=ROOT.RooArgSet()
varset3.add(x)
dataset3=ROOT.RooDataSet('dataset3','dataset3',tuple3,varset3)


pdf3.fitTo(dataset3)
results.append(['pdf3 mean',mean4.getVal()])
results.append(['pdf3 sigma',sigma4.getVal()])
results.append(['pdf3 mean error',mean4.getError()])
results.append(['pdf3 sigma error',sigma4.getError()])
pdf3_res.fitTo(dataset3)
results.append(['pdf3_res mean',mean4.getVal()])
results.append(['pdf3_res sigma',sigma4.getVal()])
results.append(['pdf3_res mean error',mean4.getError()])
results.append(['pdf3_res sigma error',sigma4.getError()])

h3=x.frame()
dataset3.plotOn(h3)
pdf3_res.plotOn(h3)
pdf3.plotOn(h3,ROOT.RooFit.LineColor(ROOT.kRed))
h3.Draw()
c1.Update()

for r in results:
    if r[0].find('mean')!=-1 and r[0].find('error')==-1:
        print '{0} {1} {2}'.format(r[0],r[1],r[1]-prod_mean)
    elif r[0].find('sigma')!=-1 and r[0].find('error')==-1:
        print '{0} {1} {2}'.format(r[0],r[1],r[1]-prod_sigma)
    else:
        print r[0],r[1]
print 'mc mean={0} amp={1} sigma={2}'.format(prod_mean,prod_amp,prod_sigma)
u=raw_input('done?')