#!/usr/bin/env python

# ---------------------------------------------------------------------------------- #
# 
#  Python/ROOT macro for illustrating how to generate random numbers following a
#  specific PDF using uniformly distributed random numbers. Both the Accept-Reject
#  (Von Neumann) and transformation methods are used.
#
#  Author: Troels C. Petersen (NBI)
#  Email:  petersen@nbi.dk
#  Date:   27th of November 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 = False         # Determining if plots are saved or not

r = TRandom3()
r.SetSeed(1)

Npoints = 10000



#----------------------------------------------------------------------------------- #
# Problem 1: Produce random points following f(x) ~ exp(-x/3), x in [0,inf]
#----------------------------------------------------------------------------------- #

Hist_exp = TH1F("Hist_exp", ";x (exponentially distributed);Frequency", 100, 0.0, 20.0)

for i in range (Npoints) :
    x = -3.0 * log(r.Uniform())
    Hist_exp.Fill(x)

canvas1 = TCanvas("canvas1","", 50, 50, 600, 400)
Hist_exp.SetLineColor(kBlue)
Hist_exp.SetLineWidth(2)
Hist_exp.Draw()

# Normalization is Npoints * binwidth = 2000:
f_exp = TF1("f_exp", "2000 * 1.0/3.0 * exp(-x/3.0)", 0, 20.0)
f_exp.Draw("same")

if (SavePlots) : canvas1.SaveAs("Hist_exp.pdf")




#----------------------------------------------------------------------------------- #
# Problem 2: Produce random points following f(x) ~ x*cos(x), x in [0,pi/2]
#----------------------------------------------------------------------------------- #

pi = TMath.Pi()

Hist_xcos = TH1F("Hist_xcos", ";x (exponentially distributed);Frequency", 100, 0.0, pi/2.0)

# A bit annoyingly, there is no do-while loop in Python, which explains the below "funny" while:
Ntry = 0
for i in range (Npoints) :
    while True :
        Ntry += 1                    # Count the number of tries, to get efficiency/integral
        x = pi/2.0 * r.Uniform()     # Range that f(x) is defined/wanted in
        y = 1.0 * r.Uniform()        # Upper bound for function values (a better bound exists!)
        if (y < x*cos(x)) : break
    Hist_xcos.Fill(x)

eff = float(Npoints) / float(Ntry)                # Efficiency
eff_error = sqrt(eff * (1.0-eff) / float(Ntry))   # Error on efficiency (binomial!)
integral =  eff       * pi/2.0 * 1.0              # Integral
eintegral = eff_error * pi/2.0 * 1.0              # Error on integral
print "  Integral of f(x) = x*cos(x), x in [0,pi/2] is: %7.4f +- %6.4f"%(integral, eintegral)
# Since the integral of x*cos in [0,pi/2] is around 0.57, the proper normalisation is 1/0.57

canvas2 = TCanvas("canvas2","", 80, 80, 600, 400)
Hist_xcos.SetLineColor(kBlue)
Hist_xcos.SetLineWidth(2)
Hist_xcos.Draw()

# Further normalization here is Npoints * binwidth = 10000 * pi/2 / 100 = pi * 50:
f_xcos = TF1("f_xcos", "3.1416 * 50 * 1.0/0.570 * x*cos(x)", 0, pi/2.0)
f_xcos.Draw("same")

if (SavePlots) : canvas2.SaveAs("Hist_xcos.pdf")



#----------------------------------------------------------------------------------- #
# Problem for you: Try to combine the two yourself:
#----------------------------------------------------------------------------------- #
#
# Estimate the integral of exp(-x/3.0)*cos(x)^2 in the interval [0,inf] using a
# combination of Transformation and Hit-and-Miss method.

# Make histograms for illustration (but of course not to infinity!!!):
Hist_expcos = TH2F("Hist_expcos", "Hist_expcos", 200, 0.0, 10.0, 200, 0.0, 1.0)

Nhit = 0
for i in range ( Npoints ) :
      
    # Get an exponentially distributed number:
    x = r.Exp(3.0)

    # Get a uniformly distributed number y and test if it is below cos(x)^2:
    y = r.Uniform()
    if (y < cos(x)*cos(x)) :
        Nhit += 1
        Hist_expcos.Fill(x,y*exp(-x/3.0))


# Calculate the fraction of points within the circle and its error:
f  = float(Nhit) / float(Npoints)
ef = sqrt(f * (1.0-f) / float(Npoints))

integral  = 3.0 * f       # Multiply integral of exponential alone with fraction
eintegral = 3.0 * ef      # Same for error
print "  Integral of f(x) = exp(-x/3)*cos(x)^2, x in [0,inf] is: %7.4f +- %6.4f"%(integral,eintegral)


# Distribution of (x,y) points:
canvas3 = TCanvas("canvas3","", 130, 130, 600, 400)

Hist_expcos.SetMarkerColor(kRed)
Hist_expcos.SetMarkerStyle(20)
Hist_expcos.SetMarkerSize(0.4)
Hist_expcos.Draw()

if (SavePlots) : canvas3.SaveAs("Hist_expcos.pdf")




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

# ----------------------------------------------------------------------------------- #
# 
# The purpose of this exercise is to combine Transformation and Hit-and-Miss methods
# to evaluate an integral (here only in 1D, to make it easy/illustrative).
# 
# Questions:
# ----------
#  1) Find the integral of exp(-x/3.0)*cos(x)^2 in the interval [0,inf] using a
#     combination of the Transformation and Hit-and-Miss methods.
# 
# ----------------------------------------------------------------------------------- #