#!/usr/bin/env python

# ---------------------------------------------------------------------------------- #
# 
#  Python/ROOT macro for producing Accept-Rejection method illustration,
#  with the function:  f(x) = (1 + sin(x)) / 2pi   for x in [0,2pi]
#
#  Author: Troels C. Petersen (NBI)
#  Email:  petersen@nbi.dk
#  Date:   4th of December 2016
# 
# ---------------------------------------------------------------------------------- #

from ROOT import *        # Import ROOT libraries (to use ROOT)
import math               # Import MATH libraries (for extended MATH functions)

gStyle.SetOptStat("emr")
gStyle.SetOptFit(1111)

SavePlots = True         # Determining if plots are saved or not

r = TRandom3()
r.SetSeed(1)

Npoints = 2000

pi = TMath.Pi()


#----------------------------------------------------------------------------------- #
# Plot with function only:
#----------------------------------------------------------------------------------- #

canvas1 = TCanvas("canvas1","", 80, 80, 1000, 600)
frame1 = canvas1.DrawFrame(0.0, 0.0, 2.0*pi, 1.0/pi)
frame1.SetTitle("Function: f(x) = (1 + sin(x)) / 2#pi");
frame1.GetXaxis().SetTitle("X-axis");
frame1.GetYaxis().SetTitle("Y-axis");

f_1sin = TF1("f_1sin", "(1.0 + sin(x))/(2.0*3.141592653)", 0.0, 2.0*pi)
f_1sin.Draw("same")

canvas1.Update()
if (SavePlots) : canvas1.SaveAs("Plot_1plusSinX.pdf")


#----------------------------------------------------------------------------------- #
# Plot with points:
#----------------------------------------------------------------------------------- #

hist_accept = TH2F("hist_accept", "", 500, 0.0, 2.0*pi, 500, 0.0, 1.0/pi)
hist_reject = TH2F("hist_reject", "", 500, 0.0, 2.0*pi, 500, 0.0, 1.0/pi)

for i in xrange(Npoints) :
    x = 2.0 * pi * r.Uniform()
    y = 1.0 / pi * r.Uniform()
    if ((1.0+sin(x))/(2.0*pi) > y) : hist_accept.Fill(x, y)
    else                           : hist_reject.Fill(x, y)

canvas2 = TCanvas("canvas2","", 100, 100, 1000, 600)
frame2 = canvas2.DrawFrame(0.0, 0.0, 2.0*pi, 1.0/pi)
frame2.SetTitle("Illustration of Accept-Rejection method for PDF: f(x) = (1 + sin(x)) / 2#pi");
frame2.GetXaxis().SetTitle("X-axis");
frame2.GetYaxis().SetTitle("Y-axis");

f_1sin = TF1("f_1sin", "(1.0 + sin(x))/(2.0*3.141592653)", 0.0, 2.0*pi)
f_1sin.Draw("same")

hist_accept.SetMarkerStyle(20)
hist_accept.SetMarkerColor(kBlue)
hist_accept.Draw("same")
hist_reject.SetMarkerStyle(24)
hist_reject.SetMarkerColor(kRed)
hist_reject.Draw("same")

canvas2.Update()
if (SavePlots) : canvas2.SaveAs("Plot_1plusSinX_withPoints.pdf")

eff = float(hist_accept.GetEntries()) / float(Npoints)
eff_error = sqrt(eff * (1.0-eff) / float(Npoints))
print "  Efficiency of Accept-Rejection method in this case: %5.3f +- %5.3f"%(eff, eff_error)


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