#!/usr/bin/env python

# ----------------------------------------------------------------------------------- #
#
#  ROOT macro for illustrating the two methods for generating random numbers according
#  to a known PDF, starting with random uniformly distributed numbers:
#   - Transformation method (if function can be integrated and then inverted).
#   - Hit & Miss (or Accept-Reject) method (by Ulam Stanislav and John Von Neumann).
#
#  The function used for illustration is: f(x) = 2x, in the interval [0,1].
#
#  For more information see:
#    G. Cowan: Chapter 3
#    P. R. Bevington: page 81-84
#    
#  Author: Troels C. Petersen (NBI)
#  Email:  petersen@nbi.dk
#  Date:   25th of November 2016
#
# ----------------------------------------------------------------------------------- #


from ROOT import *
from array import array


# ----------------------------------------------------------------------------------- #
def sqr(a) : return a*a
# ----------------------------------------------------------------------------------- #

# Setting what to be shown in statistics box:
gStyle.SetOptStat("emr")
gStyle.SetOptFit(1111)

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

SavePlots = False


# ----------------------------------------------------------------------------------- #
# 
# ----------------------------------------------------------------------------------- #

 # Set parameters:
Nnumbers = 10000       # Number of random numbers produced.

# Constant c, for f(x) = x^3 in [1,c] to be normalised:
c = 5**(1.0/4.0)

# Make histograms (note how the choice of 120 bins):
Hist_Trans   = TH1F("Hist_Trans",   "Hist_Trans",   120, 1-0.1, c+0.1)
Hist_HitMiss = TH1F("Hist_HitMiss", "Hist_HitMiss", 120, 1-0.1, c+0.1)


#----------------------------------------------------------------------------------
# Loop over process (generate numbers according to PDF f(x) = 2x in [0,1]):
#----------------------------------------------------------------------------------

for i in range (Nnumbers) :

    # Transformation method:
    # ----------------------
    # Integration gives the function F(x) = 1/4 x^4 + 1/4, which inverted gives F^-1(y) = (4y+1)*1/4:
    x_trans = (4.0*r.Uniform()+1)**(1.0/4.0)
    Hist_Trans.Fill(x_trans)

    # Hit & Miss method:
    # ------------------
    # Generate two random numbers uniformly distributed in [1,c]x[0,c^3], until they
    # fulfill the "Hit requirement" (and start with numbers, that DON'T full it!):
    x_hitmiss = 0.0
    y = 1.0
    while (x_hitmiss**3 < y) :      # ...so keep making new numbers, until this is fulfilled!
        x_hitmiss = r.Uniform(1,c)
        y  = r.Uniform(0,c**3)
      
    Hist_HitMiss.Fill(x_hitmiss)



#---------------------------------------------------------------------------------- 
# Plot histograms on screen:
#---------------------------------------------------------------------------------- 

c0 = TCanvas("c0", "", 20, 20, 1000, 600)

# Setting line width and color and axis titles of histograms:
Hist_Trans.GetXaxis().SetTitle("Random number")
Hist_Trans.GetYaxis().SetTitle("Frequency")
Hist_Trans.SetLineWidth(2)
Hist_Trans.SetLineColor(kBlue)
Hist_HitMiss.SetLineWidth(2)
Hist_HitMiss.SetLineColor(kRed)

# Fitting histogram with Transformation numbers:
fit_Trans = TF1("fit_Trans","[0]+[1]*x**3", 1, c)
fit_Trans.SetLineColor(kBlue)
Hist_Trans.Fit("fit_Trans", "RL")
Hist_Trans.Draw()
gPad.Update()
st1 = Hist_Trans.GetListOfFunctions().FindObject("stats")
st1.SetX1NDC(0.12)
st1.SetX2NDC(0.40)
st1.SetY1NDC(0.89)
st1.SetY2NDC(0.70)
st1.SetLineColor(4)

# Fitting histogram with Hit/Miss numbers:
fit_HitMiss = TF1("fit_HitMiss","[0]+[1]*x**3", 1, c)
fit_HitMiss.SetLineColor(kRed)
Hist_HitMiss.Fit("fit_HitMiss", "RL")
Hist_HitMiss.Draw()
gPad.Update()
st2 = Hist_HitMiss.GetListOfFunctions().FindObject("stats")
st2.SetX1NDC(0.12)
st2.SetX2NDC(0.40)
st2.SetY1NDC(0.69)
st2.SetY2NDC(0.50)
st2.SetLineColor(2)

# Drawing histograms with errors in the same plot:
Hist_Trans.Draw("e")
Hist_HitMiss.Draw("e same")

# Legend:
legend = TLegend(0.60, 0.18, 0.85, 0.33)
legend.SetLineColor(1)
legend.SetFillColor(0)
legend.AddEntry(Hist_HitMiss, "Hit & Miss", "LE")
legend.AddEntry(Hist_Trans, "Transformation", "LE")
legend.Draw()

c0.Update()
if (SavePlots) : c0.SaveAs("Hist_TransVsHitMiss.pdf")


# Calculate the probability that these two distributions are the same:
pChi2 = Hist_HitMiss.Chi2Test(Hist_Trans, "UU")
pKolm = Hist_HitMiss.KolmogorovTest(Hist_Trans, "UU")
print "  The probability of the two distributions being the same is:"
print "  Chi-Square:  p = %6.4f      Kolmogorov:  p = %6.4f"%(pChi2, pKolm)
# NOTE: We will get to what the "Kolmogorov Test" is...



raw_input('Press Enter to exit')


#---------------------------------------------------------------------------------- 
# 
# Questions:
# ----------
#  1) Make sure that you understand how the two methods works! What is the
#     efficiency of each of the methods? 
# 
#  2) Compare the two histograms - are they from the same distribution? Answer by
#     first comparing the fit parameters, and see if they overlap. Afterwards,
#     calculate the Chi2 between the two histograms (i.e. the difference divided
#     by the uncertainty bin by bin), and compare this to the number of degrees
#     of freedom.
# 
#  3) Try to repeat the exercise with the function f(x) = x^3 in the range [1,C],
#     (perhaps in a new file, so that you save this macro), where C is a number,
#     which ensures that the PDF f(x) is normalized. Think about what changes.
#     Again compare the two histograms.
# 
#----------------------------------------------------------------------------------