#!/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

pi = TMath.Pi()

print "  Yes, you have to do some calculations and coding, before the correct results appear!"



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

# FIRST:
# Which method should you use to produce random numbers according to this distribution?
# Think about this, argue for one method and possibly discuss it with a fellow student.

# NEXT:
# Do so, and see that it fits, by drawing (not fitting) the function on top!

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

for i in range (Npoints) :
    x = r.Uniform()             # This is of course where you should change things!!!
    Hist_exp.Fill(x)

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

# Ask yourself, if you can figure out the normalisation constant, knowing the number
# of events produced and taking the bin width into account.
f_exp = TF1("f_exp", "1.0 * exp(-x/3.0)", 0, 20.0)
f_exp.SetLineColor(kRed)
f_exp.SetLineWidth(2)
f_exp.Draw("same")

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




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

# FIRST:
# Which method should you use to produce random numbers according to this distribution?

# NEXT:
# Do so, and see that it fits, by drawing (not fitting) the function on top!

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

Nhit = 0
for i in range (Npoints) :
    x = r.Uniform()            # Use these to make a "box" around the range wanted.
    y = r.Uniform()            # Use these to make a "box" around the range wanted.
    if (y < x) :               # Test if x-value is accepted or not (change yourself!)
        Nhit += 1
        Hist_xcos.Fill(x)

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

# Again, think about the normalisation:
f_xcos = TF1("f_xcos", "3.1416*x*cos(x)", 0, pi/2.0)
f_xcos.Draw("same")

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


# ALSO:
# Use the random numbers to calculate an integral and its uncertainty.

# Integral of f(x) in the range [0,pi/2]:
integral = float(Nhit) / float(Npoints) * pi/2.0 * 1.0
eintegral = -1.0                 # Calculate this yourself!
print "  Integral of f(x) = x*cos(x), x in [0,pi/2] is: %7.4f +- %6.4f"%(integral, eintegral)




#----------------------------------------------------------------------------------- #
# Problem 3: Try to combine the two methods 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.

Nhit = 0
for i in range ( Npoints ) :
      
    # Well, the inside of the loop is up to you!
    Nhit += 1


integral  = -1.0     # Well, put your value here
eintegral = -1.0     # ...
print "  Integral of f(x) = exp(-x/3)*cos(x)^2, x in [0,inf] is: %7.4f +- %6.4f"%(integral,eintegral)





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

# ----------------------------------------------------------------------------------- #
# 
# The purpose of this exercise is understand how to produce random numbers both using
# the Hit-and-Miss (Von Neumann) and the Transformation method, possibly 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.
# 
# ----------------------------------------------------------------------------------- #