from ROOT import TCanvas, TPad, TFile, TPaveText
from ROOT import gBenchmark, gStyle, gROOT

#gROOT.Reset()

c1 = TCanvas('c1','The Ntuple canvas',200,10,700,780)

gBenchmark.Start('ntuple1')

#
# Connect ROOT histogram/ntuple demonstration file
# generated by example hsimple.C.
f1 = TFile('hsimple.root')

#
# Inside this canvas, we create 4 pads
pad1 = TPad('pad1','This is pad1',0.02,0.52,0.48,0.98,21)
pad2 = TPad('pad2','This is pad2',0.52,0.52,0.98,0.98,21)
pad3 = TPad('pad3','This is pad3',0.02,0.02,0.48,0.48,21)
pad4 = TPad('pad4','This is pad4',0.52,0.02,0.98,0.48,1)
pad1.Draw()
pad2.Draw()
pad3.Draw()
pad4.Draw()

#
# Change default style for the statistics box
gStyle.SetStatW(0.30)
gStyle.SetStatH(0.20)
gStyle.SetStatColor(42)

#
# Display a function of one ntuple column imposing a condition
# on another column.
pad1.cd()
pad1.SetGrid()
pad1.SetLogy()
pad1.GetFrame().SetFillColor(15)
ntuple = gROOT.FindObject('ntuple')
ntuple.SetLineColor(1)
ntuple.SetFillStyle(1001)
ntuple.SetFillColor(45)
ntuple.Draw('3*px+2','px**2+py**2>1')
ntuple.SetFillColor(38)
ntuple.Draw('2*px+2','pz>2','same')
ntuple.SetFillColor(5)
ntuple.Draw('1.3*px+2','(px^2+py^2>4) && py>0','same')
c1.Update()

#
# Display the profile of two columns
# The profile histogram produced is saved in the current directory with
# the name hprofs
pad2.cd()
pad2.SetGrid()
pad2.GetFrame().SetFillColor(32)
ntuple.Draw('pz:px>>hprofs','','goffprofs')
hprofs = gROOT.FindObject('hprofs')
hprofs.SetMarkerColor(5)
hprofs.SetMarkerSize(0.7)
hprofs.SetMarkerStyle(21)
hprofs.Fit('pol2')

#
# Get pointer to fitted function and modify its attributes
fpol2 = hprofs.GetFunction('pol2')
fpol2.SetLineWidth(4)
fpol2.SetLineColor(2)
c1.Update()

#
# Display a scatter plot of two columns with a selection.
# Superimpose the result of another cut with a different marker color
pad3.cd()
pad3.GetFrame().SetFillColor(38)
pad3.GetFrame().SetBorderSize(8)
ntuple.SetMarkerColor(1)
ntuple.Draw('py:px','pz>1')
ntuple.SetMarkerColor(2)
ntuple.Draw('py:px','pz<1','same')
c1.Update()

#
# Display a 3-D scatter plot of 3 columns. Superimpose a different selection.
pad4.cd()
ntuple.Draw('pz:py:px','(pz<10 && pz>6)+(pz<4 && pz>3)')
ntuple.SetMarkerColor(4)
ntuple.Draw('pz:py:px','pz<6 && pz>4','same')
ntuple.SetMarkerColor(5)
ntuple.Draw('pz:py:px','pz<4 && pz>3','same')
l4 = TPaveText(-0.9,0.5,0.9,0.95)
l4.SetFillColor(42)
l4.SetTextAlign(12)
l4.AddText('You can interactively rotate this view in 2 ways:')
l4.AddText('  - With the RotateCube in clicking in this pad')
l4.AddText('  - Selecting View with x3d in the View menu')
l4.Draw()

#
# done
c1.cd()
c1.Update()
gStyle.SetStatColor(19)
gBenchmark.Show('ntuple1')