#!/usr/bin/env python

# ----------------------------------------------------------------------------------- #
#
#  ROOT macro for illustrating the two methods for generating random numbers according
#  to a known PDF:
#   - 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:   16th of September 2014
#
# ----------------------------------------------------------------------------------- #


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)

verbose = True


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

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

# Make histograms (note how the choice of 120 bins):
Hist_Trans   = TH1F("Hist_Trans",   "Hist_Trans",   120, -0.1, 1.1)
Hist_HitMiss = TH1F("Hist_HitMiss", "Hist_HitMiss", 120, -0.1, 1.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) = x^2, which inverted gives F-1(x) = sqrt(x):
    x_trans = sqrt(r.Uniform())
    Hist_Trans.Fill(x_trans)

    # Hit & Miss method:
    # ------------------
    # Generate two random numbers uniformly distributed in [0,1]x[0,2], until they
    # fulfill the "Hit requirement":
    x_hitmiss = 0.0
    y = 1.0
    while (2.0*x_hitmiss < y) :
        x_hitmiss = 1.0 * r.Uniform()
        y  = 2.0 * r.Uniform()
      
    Hist_HitMiss.Fill(x_hitmiss)



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

c0 = TCanvas("c0", "", 100, 20, 600, 400)

# 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(2)
Hist_HitMiss.SetLineWidth(2)
Hist_HitMiss.SetLineColor(4)

# Fitting histogram with Transformation numbers:
fit_Trans = TF1("fit_Trans","[0] + [1]*x", 0.0, 1.0)
fit_Trans.SetLineColor(4)
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", 0.0, 1.0)
fit_HitMiss.SetLineColor(4)
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()
# c0.SaveAs("Hist_TransVsHitMiss.png")


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)




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.
#     Try also to use the Kolmogorov-Smirnov test (which we will get to shortly!).
#     Are they the same by these tests?
# 
#  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.
# 
# Advanced:
#  4) Try also to apply the run test (Barlow 8.3.2) also known as the Wald-Wolfowitz
#     runs test to either of the two cases. If this is succesful, then try to shift
#     one distribution by a little and see if the runs test is more sensitive than
#     the Chi2 test.
#
#----------------------------------------------------------------------------------