import ROOT

c1=ROOT.TCanvas()
c2=ROOT.TCanvas()
c3=ROOT.TCanvas()

#generate dataset
tuple1d=ROOT.TNtuple("1dDist",'','x')

for i in range(1,100):
    for j in range(i*i):
        tuple1d.Fill(i)

x=ROOT.RooRealVar('x','x',0,100)
varset1d=ROOT.RooArgSet('1d')
varset1d.add(x)
dataset1d=ROOT.RooDataSet('dataset1d','1d dataset',tuple1d,varset1d)
a=ROOT.RooRealVar('a','a',1,)
b=ROOT.RooRealVar('b','b',1,-1,1)
c=ROOT.RooRealVar('c','c',0,-1,1)
poly1d=ROOT.RooPolynomial('p1','p1',x,ROOT.RooArgList(a,b))
poly1d.fitTo(dataset1d)


plot1d=x.frame()
dataset1d.plotOn(plot1d)
poly1d.plotOn(plot1d)
plot1d.Draw()


y=ROOT.RooRealVar('y','y',0,10000)
tuple2d=ROOT.TNtuple("2dDist",'','x:y')
for i in range(100):
    for j in range(5):
        tuple2d.Fill(i,5*i*i+10*ROOT.gRandom.Gaus())

varset2d=ROOT.RooArgSet('2d')
varset2d.add(x)
varset2d.add(y)

a2d=ROOT.RooRealVar('a2d','a2d',0,0,100)
b2d=ROOT.RooRealVar('b2d','b2d',2,-5,5)
mu=ROOT.RooPolyVar('mu','mu',x,ROOT.RooArgList(ROOT.RooFit.RooConst(1),a2d,b2d))
poly2d=ROOT.RooPolynomial('p2','p2',x,ROOT.RooArgList(a2d,b2d))
sigma=ROOT.RooRealVar('sigma','sigma',200,0,200)
gauss=ROOT.RooGaussian('gaus','gaus',y,mu,sigma)

gauss2dpdf=ROOT.RooProdPdf('gauss2dpdf','gaus*p2',ROOT.RooArgList(gauss,poly2d))
#pdf2=ROOT.RooGenericPdf('pdf2','1+a2d*x+b2d*x*x',ROOT.RooArgSet(a2d,b2d,x)


dataset2d=ROOT.RooDataSet('dateset2d','2d dataset',tuple2d,varset2d)
gauss.fitTo(dataset2d)
#gauss2dpdf.fitTo(dataset2d)

h2d=dataset2d.createHistogram(x,y,100,100)
h2d.Draw()
h2d_mu=mu.createHistogram('x',100,100)
h2d_mu.Draw('same')


x3=ROOT.RooRealVar('x3','x3',0,100)
y3=ROOT.RooRealVar('y3','y3',0,100)
z3=ROOT.RooRealVar('z3','z3',0,20000)
tuple3d=ROOT.TNtuple('3dDist','','x3:y3:z3')
for i in range(100):
    for j in range(100):
        tuple3d.Fill(i,j,1.5*i**2+j**2+10*ROOT.gRandom.Gaus())

varset3d=ROOT.RooArgSet('3d')
varset3d.add(x3)
varset3d.add(y3)
varset3d.add(z3)

a3d1=ROOT.RooRealVar('a3d1','a3d1',0,0,10)
b3d1=ROOT.RooRealVar('b3d1','b3d1',1,0,1000)
mux=ROOT.RooPolyVar('mux','mux',x3,ROOT.RooArgList(ROOT.RooFit.RooConst(1),a3d1,b3d1))

a3d2=ROOT.RooRealVar('a3d2','a3d2',0,0,10)
b3d2=ROOT.RooRealVar('b3d2','b3d2',1,0,1000)
muy=ROOT.RooPolyVar('muy','muy',y3,ROOT.RooArgList(ROOT.RooFit.RooConst(1),a3d2,b3d2))

#muz=ROOT.RooFormulaVar('muz','muz','mux+muy',ROOT.RooArgList(mux,muy))
muz=ROOT.RooFormulaVar('muz','muz','@0+@1*(@2*@2)+@3*(@4*@4)',ROOT.RooArgList(ROOT.RooFit.RooConst(1),a3d1,x3,a3d2,y3))

sigma2=ROOT.RooRealVar('sigma2','sigma2',20000,0,20000000)
gauss2=ROOT.RooGaussian('gaus2','gaus2',z3,muz,sigma2)

dataset3d=ROOT.RooDataSet('dataset3d','3d dataset',tuple3d,varset3d)

gauss2.fitTo(dataset3d)


h3d=ROOT.RooAbsData.createHistogram(dataset3d,'h3d',x3,ROOT.RooFit.YVar(y3),ROOT.RooFit.ZVar(z3))
h3d_muz=muz.createHistogram('x3,y3',100,100,100)
h3d.Draw()
h3d_muz.Draw('same surf')

u=raw_input()