#!/usr/bin/env python

# ----------------------------------------------------------------------------------- #
# 
#   Python script giving example of making a nice ROOT figure.
# 
#   Author: Troels C. Petersen (NBI)
#   Email:  petersen@nbi.dk
#   Date:   11th of August 2014
# 
# ----------------------------------------------------------------------------------- #

from ROOT import *
from array import array
from math import sqrt


# ----------------------------------------------------------------------------------- #
# Example of nice plot:
# ----------------------------------------------------------------------------------- #

gROOT.Reset()

# Setting of general plotting style:
gStyle.SetCanvasColor(0)
gStyle.SetFillColor(0)

# Settings for statistics box:
gStyle.SetOptStat(1111)
gStyle.SetOptFit(1111)

# Statistics and fitting box placing and size:
gStyle.SetStatX(0.92)
gStyle.SetStatY(0.94)
gStyle.SetStatW(0.15);                
gStyle.SetStatH(0.10); 

# Random numbers from the Mersenne-Twister:
r = TRandom3()
r.SetSeed(1)

# Text to write on the plot:
text = TLatex()
text.SetNDC()
text.SetTextFont(1)
text.SetTextColor(1)

SavePlots = True


# ------------------------------------------------------------------ #
# Generate data:
# ------------------------------------------------------------------ #

Npoints = 200
x  = array( 'f', [0.0]*Npoints )   # Create an array of length Npoints with zeros.
ex = array( 'f', [0.0]*Npoints )   # [X]*N is just a python syntax to initialize
y  = array( 'f', [0.0]*Npoints )   # a list of values X.
ey = array( 'f', [0.0]*Npoints )   

alpha0 = 3.6
alpha1 = -2.8
alpha2 = 2.3
alpha3 = -0.5
sigmay = 0.2

Npoints2 = 40
Noffset = 50

# Generate data:
# --------------
for i in range( Npoints ) : 

    sigmay_used = sigmay
    if (Noffset < i < Noffset+Npoints2) : sigmay_used = sigmay * 1.6   # Noisy region!

    x[i]  = float(i+3) / 50.0
    ex[i] = 0.0
    y[i]  = alpha0 + alpha1 * x[i] + alpha2 * x[i]*x[i] + alpha3 * x[i]*x[i]*x[i] + r.Gaus(0.0,sigmay_used)
    ey[i] = sigmay


# ------------------------------------------------------------------ #
# Fit the data and plot the result:
# ------------------------------------------------------------------ #

# For graphs in general, see the tutorial page for many more examples:
#   http://root.cern.ch/root/html/tutorials/graphs/index.html


# Make a white canvas and draw the example fit in:
c0 = TCanvas("c0","",50,50,1200,700)
c0.SetFillColor(0)


# Fit data:
# ---------
graph = TGraphErrors( Npoints,x,y,ex,ey )

# How do you control the axis?
# Read more at: http://root.cern.ch/root/html/TGaxis.html
graph.SetTitle("Fit of a graph")
graph.GetXaxis().SetRangeUser(-0.1, 4.1)
graph.GetYaxis().SetRangeUser(-4.0, 5.4)
graph.GetXaxis().CenterTitle();              # Centers titles
graph.GetXaxis().SetTitle("Deciliters of alcohol")
graph.GetXaxis().SetNdivisions(50510)   
  # The "50510" is a code, see http://root.cern.ch/root/html/TAttAxis.html
  # See also: http://root.cern.ch/root/html/tutorials/graphics/gaxis.C.html
graph.GetYaxis().SetTitle("Spirit of the night")
graph.GetYaxis().SetTitleOffset(0.5)
graph.SetLineWidth(2)
graph.SetMarkerStyle(20)
graph.SetMarkerSize(0.6)
graph.SetMarkerColor(kRed)
graph.Draw("AP")

fit = TF1("fit", "pol3", 0.0, float(Npoints+4)/50.0 )
fit.SetLineColor(kRed)
fit.SetParNames("Constant", "Linear coef.", "Quadratic coef.", "Qubic coef.")
graph.Fit("fit","R")

# In ROOT, you can just ask the fit function for it:
Chi2_fit = fit.GetChisquare()
Ndof_fit = fit.GetNDF()
Prob_fit = fit.GetProb()

# Text:
text.SetTextSize(0.04)
text.DrawLatex(0.16, 0.84, "The fit is dubious: Prob(%5.1f,%3d)=%6.4f"%(Chi2_fit,Ndof_fit,Prob_fit))
text.SetTextSize(0.03)
text.DrawLatex(0.16, 0.81, "Perhaps there is a noisy region?")


# ------------------------------------------------------------------ #
# Make insert figure (done using a "pad"):
# ------------------------------------------------------------------ #

# Copy data into a new graph to be drawn
x2  = array( 'f', [0.0]*Npoints2 )
ex2 = array( 'f', [0.0]*Npoints2 )
y2  = array( 'f', [0.0]*Npoints2 )
ey2 = array( 'f', [0.0]*Npoints2 )   

for i in range( Npoints2 ) : 
    x2[i]  = x[i+Noffset]
    ex2[i] = ex[i+Noffset]
    y2[i]  = y[i+Noffset]
    ey2[i] = ey[i+Noffset]

graph2 = TGraphErrors( Npoints2,x2,y2,ex2,ey2 )


# Make the pad:
pad = TPad("pad", "pad", 0.12, 0.11, 0.65, 0.55, 0, 1, 1);
pad.SetTopMargin(0.0);
pad.SetFillColor(0);
pad.GetFrame().SetFillColor(0);
pad.GetFrame().SetBorderSize(0);
pad.Draw();
pad.cd();            # Go to this pad, such that next "Draw" is here!

xlow  = 1.05
xhigh = 1.85
ylow  = 1.90
yhigh = 4.10

graph2.SetTitle("")
graph2.GetXaxis().SetRangeUser(xlow,xhigh)
graph2.GetYaxis().SetRangeUser(ylow,yhigh)
graph2.SetLineWidth(2)
graph2.SetMarkerStyle(20)
graph2.SetMarkerColor(2)
graph2.Draw("AP")

fit.Draw("same")


# Go back to the large canvas and indicate the cut-out:
c0.cd()
line0 = TLine(xlow, ylow, xhigh, ylow)
line0.Draw()
line1 = TLine(xlow, ylow, xlow, yhigh)
line1.Draw()
line2 = TLine(xlow, yhigh, xhigh, yhigh)
line2.Draw()
line3 = TLine(xhigh, ylow, xhigh, yhigh)
line3.Draw()

line0b = TLine(0.36, 1.34, 1.05, 1.9)
line0b.SetLineStyle(3)
line0b.Draw()

line1b = TLine(1.86, 1.9, 2.52, 1.34)
line1b.SetLineStyle(3)
line1b.Draw()


if (SavePlots) : c0.SaveAs("NiceFigureExample.pdf")



# ---------------------------------------------------------------------------------- #
raw_input( ' ... Press enter to exit ... ' )