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

c1=ROOT.TCanvas("c","roofit gauss test",1000,1000)
c1.Divide(2,1)
#c2=ROOT.TCanvas()
#c3=ROOT.TCanvas()


prod_gaus=ROOT.TF1('prod_gaus','gaus',-12,12)
prod_gaus.SetParameter(0,100)
prod_gaus.SetParameter(1,2)
prod_gaus.SetParameter(2,3)
tuple1=ROOT.TNtuple('2d gauss','','x:amp')
tuple1_1=ROOT.TNtuple('2d gauss v2','','x:amp_sig')

cov=ROOT.TMatrixD(2,2)
cov[0][0]=0
cov[1][1]=0
cov[1][0]=cov[0][1]=0
x_mean=0
sumAmp=0
for i in range(10):

    #x=i/10.
    x=ROOT.gRandom.Uniform(-12,12)

    
    amp=prod_gaus.Eval(x)+8*ROOT.gRandom.Gaus()
    while amp<0:
        amp=prod_gaus.Eval(x)+8*ROOT.gRandom.Gaus()
    x_mean+=amp*x
    sumAmp+=amp
    sign=1
    if ROOT.gRandom.Uniform(-1,1)<0:
        sign=-1
    
    for i in range(2):
        tuple1.Fill(x,amp)

x_mean/=sumAmp
maxamp=tuple1.GetMaximum('amp')
print 'xmean=',x_mean
varX=0
for e in tuple1:
    if ROOT.gRandom.Uniform(-1,1)<0:
        sign=-1
    thisamp=sign*(e.amp-maxamp)
    print thisamp,maxamp,e.amp
    tuple1_1.Fill(e.x,thisamp)

    cm=ROOT.TMatrixD(2,1)
    cm[0][0]=e.x-x_mean
    cm[1][0]=sign*e.amp
    cm_t=ROOT.TMatrixD(ROOT.TMatrixD.kTransposed,cm)
    acov=ROOT.TMatrixD(cm,ROOT.TMatrixD.kMult,cm_t)
    cov+=acov


cov*=1./tuple1_1.GetEntries()
cov[1][0]=0
cov[0][1]=0
cov=ROOT.TMatrixDSym(2,cov.GetMatrixArray())

x=ROOT.RooRealVar('x','x',-50,50)
amp=ROOT.RooRealVar('amp','amp',0,120)
amp_sig=ROOT.RooRealVar('amp_sig','amp_sig',-120,120)
gamp=ROOT.RooRealVar('gamp','gamp',80,0,120)
mu=ROOT.RooRealVar('mu','mu',4,-20,20)
mux=ROOT.RooRealVar('mux','mux',4,-20,20)
mua=ROOT.RooRealVar('mua','mua',0)
mua.setConstant(True)
sigma=ROOT.RooRealVar('sigma','sigma',5,-100,100)
sigmax=ROOT.RooRealVar('sigmax','sigmax',5,-100,100)
sigmaa=ROOT.RooRealVar('sigmaa','sigmaa',50,-200,200)
xVec=ROOT.RooArgList()
xVec.add(x)
xVec.add(amp_sig)
muVec=ROOT.RooArgList()
muVec.add(mux)
muVec.add(mua)


#mu2=ROOT.RooRealVar('mu2','mu2',0,-20,20)
sigma2=ROOT.RooRealVar('sigma2','sigma2',4,-10,10)

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

varset1_1=ROOT.RooArgSet()
varset1_1.add(x)
varset1_1.add(amp_sig)


dataset1=ROOT.RooDataSet('dataset1','dataset1',tuple1,varset1)
dataset1_1=ROOT.RooDataSet('dataset1_1','dataset1_1',tuple1_1,varset1_1)

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

mvg=ROOT.RooMultiVarGaussian("mvg","mvg",xVec,muVec,cov)

start=time.time()
pdf1.fitTo(dataset1)
print '1d fittime',time.time()-start
h1=dataset1.createHistogram(x,amp,1000,1000)
h1_g=gausvar.createHistogram('x',1000,1000)
c1.cd(1)
h1.Draw()
h1_g.Draw('same')

'''
mvg.fitTo(dataset1_1)
cov.Print()
print math.sqrt(cov[0][0]),math.sqrt(cov[1][1])
h1_1=dataset1_1.createHistogram(x,amp_sig,1000,1000)
h1_1g=mvg.createHistogram('x,amp_sig',100,100,100)
c1.cd(2)
h1_1.Draw()
h1_1g.Draw('same cont')
c1.Update()
'''

u=raw_input('next?')
x2=ROOT.RooRealVar('x2','x2',-10,10)
y2=ROOT.RooRealVar('y2','y2',-10,10)
gamp2d=ROOT.RooRealVar('gamp2d','gamp2d',80,0,4000)
amp2=ROOT.RooRealVar('amp2','amp2',0,5000)
mux=ROOT.RooRealVar('mux','mux',2,-5,5)
muy=ROOT.RooRealVar('muy','muy',2,-5,5)
sigx=ROOT.RooRealVar('sigx','sigx',5,0,100)
sigy=ROOT.RooRealVar('sigy','sigy',5,0,100)
rho=ROOT.RooRealVar('rho','rho',0.23,-1,1)
sigma3=ROOT.RooRealVar('sigma3','sigma3',20,0,300)
varset2=ROOT.RooArgSet()
varset2.add(x2)
varset2.add(y2)
varset2.add(amp2)

gausvar2d=ROOT.Roo2DimGaussVar('gauss2d','gauss2d',x2,y2,rho,sigx,sigy,mux,muy,gamp2d)
pdf2=ROOT.RooGaussian('gaus2','gaus2',amp2,gausvar2d,sigma3)

tuple2=ROOT.TNtuple('3d gauss','','x2:y2:amp2')
N=int(ROOT.gRandom.Uniform(5,20))
for i in range(N):
    xpos=ROOT.gRandom.Uniform(-10,10)
    ypos=ROOT.gRandom.Uniform(-10,10)
    x2.setVal(xpos)
    y2.setVal(ypos)
    ampval=gausvar2d.getVal()+20*ROOT.gRandom.Gaus()
    while ampval<0:
        ampval=gausvar2d.getVal()+2*ROOT.gRandom.Gaus()
    tuple2.Fill(xpos,ypos,ampval)
 

'''       
mux.setVal(5)
muy.setVal(-2)
rho.setVal(0)
gamp2d.setVal(80)
sigx.setVal(2)
sigy.setVal(9)
'''
x2.setRange(tuple2.GetMinimum('x2'),tuple2.GetMaximum('x2'))
y2.setRange(tuple2.GetMinimum('y2'),tuple2.GetMaximum('y2'))
#gamp2d.

sigx.setConstant(True)
sigy.setConstant(True)
rho.setConstant(True)
#sigma3.setConstant(True)

dataset2=ROOT.RooDataSet('dataset2','dataset2',tuple2,varset2)

start=time.time()
pdf2.fitTo(dataset2)
print 'runtime:',time.time()-start,'N=',N
h2=ROOT.RooAbsData.createHistogram(dataset2,'h2',x2,ROOT.RooFit.YVar(y2),ROOT.RooFit.ZVar(amp2))
h2_g=gausvar2d.createHistogram('x2,y2',100,100,100)

#c1.cd(2)
#h2.Draw()
#h2_g.Draw('same surf')
#c1.Update()
u=raw_input('done?')