import ROOT
import math

ROOT.gROOT.ProcessLine(".x rootlogon.C") 

nevts=8

data_tup=ROOT.TNtuple("data",'','x:y:a')
hist=ROOT.TH1D('hist1','x*a',100,-100,100)
hist2=ROOT.TH2D('hist2','x vs a',100,-100,100,100,0,2000)
hist2_y=ROOT.TH2D('hist2_y','y vs a',100,-100,100,100,0,2000)
hist2_r=ROOT.TH2D('hist2_r','r vs a',100,-100,100,100,0,2000)

cov=ROOT.TMatrixD(2,2)
cov[0][0]=0
cov[1][1]=0
cov[1][0]=cov[0][1]=0
sumamp=0
cogx=0
cogy=0
#generate some events:
for i in range(nevts):
    x=ROOT.gRandom.Gaus(6,20)
    y=ROOT.gRandom.Gaus(-3,10)
    dist=math.sqrt((x-6)*(x-6)+(y+3)*(y+3))
    amp=ROOT.gRandom.Landau(500-3*dist,200-2*dist)
    if amp<0:
        continue
    sumamp+=amp
    cogx+=amp*x
    cogy+=amp*y
    #while amp<0:
    #    amp=ROOT.gRandom.Landau(200-3*abs(x),100-2*abs(x))
        #amp=ROOT.gRandom.Landau(200,abs(x))
    cm=ROOT.TMatrixD(2,1)
    cm[0][0]=x-6
    cm[1][0]=y+3
    cm_t=ROOT.TMatrixD(ROOT.TMatrixD.kTransposed,cm)
    acov=ROOT.TMatrixD(cm,ROOT.TMatrixD.kMult,cm_t)
    acov*=amp
    cov+=acov
    hist.Fill(x,amp)
    hist2.Fill(x,amp)
    hist2_y.Fill(y,amp)
    hist2_r.Fill(dist,amp)
    data_tup.Fill(x,y,amp)
    
cov*=1./sumamp
cov=ROOT.TMatrixDSym(2,cov.GetMatrixArray())

cogx/=sumamp
cogy/=sumamp


x=ROOT.RooRealVar('x','x',0,-100,100)
y=ROOT.RooRealVar('y','y',0,-100,100)
amp=ROOT.RooRealVar('a','a',0,0,2000)
varset=ROOT.RooArgSet('varset')
varset.add(x)
varset.add(y)
varset.add(amp)

g_mean=ROOT.RooRealVar("g_mean","mean of gaus",0,-10,10)
g_sigma=ROOT.RooRealVar('g_sigma','sigma of gaus',16,0,50)
gaus=ROOT.RooGaussian('gaus','gaus',x,g_mean,g_sigma)

xVec=ROOT.RooArgList()
xVec.add(x)
xVec.add(y)
mux=ROOT.RooRealVar('mux','mean in x',0,-10,10)
muy=ROOT.RooRealVar('muy','mean in y',0,-10,10)
muVec=ROOT.RooArgList()
muVec.add(mux)
muVec.add(muy)
mvg=ROOT.RooMultiVarGaussian("mvg","mvg",xVec,muVec,cov) 


#l_mpv=ROOT.RooFormulaVar('l_mpv','500-3*sqrt( (x-{0})*(x-{0}) + (y-{1})*(y-{1}) )'.format(cogx,cogy),ROOT.RooArgList(x,y))
l_mpv=ROOT.RooFormulaVar('l_mpv','500-3*sqrt( (x-mux)*(x-mux) + (y-muy)*(y-muy) )',ROOT.RooArgList(x,y,mux,muy))
#l_sigma=ROOT.RooFormulaVar('l_mpv','200-2*sqrt( (x-{0})*(x-{0}) + (y-{1})*(y-{1}) )'.format(cogx,cogy),ROOT.RooArgList(x,y))
l_sigma=ROOT.RooFormulaVar('l_mpv','200-2*sqrt( (x-mux)*(x-mux) + (y-muy)*(y-muy) )'.format(cogx,cogy),ROOT.RooArgList(x,y,mux,muy))
#landau=ROOT.RooLandau('landau','amp landau',amp,l_mpv,l_sigma)
landau=ROOT.RooAmpLandau('landau','amp landau',amp,x,y,ROOT.RooFit.RooConst(cogx),ROOT.RooFit.RooConst(cogy),ROOT.RooFit.RooConst(0),ROOT.RooFit.RooConst(0),ROOT.RooFit.RooConst(0),ROOT.RooFit.RooConst(0),ROOT.RooFit.RooConst(0),ROOT.RooFit.RooConst(0))
#landau=ROOT.RooAmpLandau('landau','amp landau',amp,x,y,mux,muy,ROOT.RooFit.RooConst(0),ROOT.RooFit.RooConst(0),ROOT.RooFit.RooConst(0),ROOT.RooFit.RooConst(0),ROOT.RooFit.RooConst(0),ROOT.RooFit.RooConst(0))

#l_mpv=ROOT.RooRealVar('l_mpv','mpv of landau',100,0,300)
#l_sigma=ROOT.RooRealVar('l_sigma','sigma of landau',80,0,200)
#landau=ROOT.RooLandau('landau','landau',amp,ROOT.RooFit.RooConst(200),x)

#pdf=ROOT.RooProdPdf('pdf','gaus*landau',ROOT.RooArgSet(gaus),ROOT.RooFit.Conditional(ROOT.RooArgSet(landau),ROOT.RooArgSet(amp)))
pdf=ROOT.RooProdPdf('pdf','gaus*landau',ROOT.RooArgSet(mvg),ROOT.RooFit.Conditional(ROOT.RooArgSet(landau),ROOT.RooArgSet(amp)))
dataset=ROOT.RooDataSet('dataset','dataset',data_tup,varset)
dataset_weight=ROOT.RooDataSet('dataset','dataset',data_tup,varset,'','a')
 
#print 'generate dataset'
#dataset2=pdf.generate(ROOT.RooArgSet(x,amp),2000)
#pdf=gaus 
print '\n\n\n\nfitting pdf:'
pdf.fitTo(dataset)#,ROOT.RooFit.ConditionalObservables(ROOT.RooArgSet(x)))
print '\n\n\n\nfitting mvg:'    
mvg.fitTo(dataset)
print '\n\n\n\nfittinh weighted:'
mvg.fitTo(dataset_weight,ROOT.RooFit.SumW2Error(True))
print 'fitting done'
canvas=ROOT.TCanvas('c1','c1',700,700)
canvas.Divide(2,2)
canvas.cd(1)
hist2.Draw('colz')
canvas.cd(2)
hist2_y.Draw('colz')
canvas.cd(3)
hist3=x.createHistogram('hist3',amp,'')
pdf.fillHistogram(hist3,ROOT.RooArgList(x,amp))
hist3.Draw('colz')
canvas.cd(4)
#hist4=x.createHistogram('hist4',amp,'')
#dataset2.fillHistogram(hist4,ROOT.RooArgList(x,amp))
#hist4.Draw('colz')
u=raw_input('done?')