#!/usr/bin/env python
from ROOT import *

# C r e a t e   l o w   s t a t s   1 - D   d a t a s e t 
# -------------------------------------------------------

#Weight variables
w  = RooRealVar("w", "w", 0.0, 2.0) #weight variable
fw = RooGaussian("fw","fw",w, RooFit.RooConst(0.01), RooFit.RooConst(0.001))


# Create a 2d toy pdf for sampling
x = RooRealVar("x","x",0,20)
fx = RooPolynomial("fx","fx",x,RooArgList(RooFit.RooConst(0.01),RooFit.RooConst(-0.01),RooFit.RooConst(0.0004) )) ;

# Construct a 2D toy pdf for sampling
y = RooRealVar("y","y",0,20) ;
fy = RooPolynomial("fy","fy",y, RooArgList(RooFit.RooConst(0.01), RooFit.RooConst(0.01), RooFit.RooConst(-0.0004))) ;

pdf2darg = RooArgList(fx,fy,fw)
pxy = RooProdPdf("pxy","pxy", pdf2darg ) ;

argset2d   = RooArgSet(x,y,w)
data2_temp = pxy.generate( argset2d, 1000) ;
data2_temp.Print()

data2 = RooDataSet("data2", "data2", data2_temp, argset2d, "", "w");
data2.Print()


# C r e a t e   2 - D   k e r n e l   e s t i m a t i o n   p d f
# ---------------------------------------------------------------

# Create 2D adaptive kernel estimation pdf with mirroring 
kest4 = RooNDKeysPdf("kest4","kest4", x,y, data2,"am") ;

# Create a histogram of the data
#hh_data = data2.createHistogram("hh_data", x, RooFit.Binning(20), RooFit.YVar(y, RooFit.Binning(20)))
hh_data = data2.createHistogram(x,y, "w", "hist")

# Create histogram of the 2d kernel estimation pdfs
hh_pdf  = kest4.createHistogram("x,y", 20,20)

c = TCanvas("kernelestimation","kernelestimation",1400, 1200) ;
hh_data.Draw("lego E")
hh_pdf.Draw("surf same")

raw_input(' Done dealio ... ')