import ROOT
import rootnotes #for plotting with iPython
c1 = rootnotes.default_canvas()

w = ROOT.RooWorkspace()
w.factory('Gaussian::g(x[-5,5],mu[-3,3],sigma[1])')
w.factory('Exponential::e(x,tau[-.5,-3,0])')
w.factory('SUM::model(s[50,0,100]*g,b[100,0,1000]*e)')
w.Print() #this isn't displaying in iPython

x = w.var('x')
pdf = w.pdf('model')
frame = x.frame()
data = pdf.generate(ROOT.RooArgSet(x))
data.plotOn(frame)
fitResult = pdf.fitTo(data,ROOT.RooFit.Save())
pdf.plotOn(frame)
frame.Draw()
c1

mu = w.var('mu')
print 'best fit value of mean for Gaussian is', mu.getVal(), '±', mu.getError()

s=w.var('s')
sframe=s.frame()
lnL = pdf.createNLL(data)
lnProfileL = lnL.createProfile(ROOT.RooArgSet(s))
lnProfileL.plotOn(sframe)
sframe.Draw()
c1

corrMatrixHist = fitResult.correlationHist()
corrMatrixHist.SetStats(False)
corrMatrixHist.Draw('colz text')
c1

% pylab inline
import networkx as nx

pdf.graphVizTree('expressionTree.dot')
G = nx.DiGraph(nx.read_dot('expressionTree.dot'))
nx.draw_graphviz(G,'dot')

import pydot
g = pydot.graph_from_dot_file('expressionTree.dot')
G2 = nx.from_pydot(g)
nx.draw_graphviz(G2,'dot')

"""
print "inspect networkx graph"
#iterate over networkx graph
print G2.node
for n in G2.nodes():
    print n, G2.out_degree(n), G2.successors(n)

print "inspect pydot graph"
#iterate over pydot graph
for n in g.get_node_list():
    print n.get_label(), n.get_attributes(), n.get_layer()
"""
g.write_png('test.png')

!cat expressionTree.dot

print "inspect networkx graph"
#iterate over networkx graph
print G2.node
for n in G2.nodes():
    print n, G2.out_degree(n), G2.successors(n)
G2