#!/usr/bin/env python

# ---------------------------------------------------------------------------------- #
# 
#  Python/ROOT macro for testing which fitting proceedure is likely to give the "best"
#  results. The three cases in question are:
#   * Beta distribution
#   * Double exponential
#   * Cauchy distribution on constant background
#
#  Consider each case/distribution/statistics/etc. and argue/discuss which fitting
#  method should be used. In all cases, a Chisquare fit has been applied.
#
#  Author: Troels C. Petersen (NBI)
#  Email:  petersen@nbi.dk
#  Date:   26th 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)



#----------------------------------------------------------------------------------- #
# CASE 1: Beta distribution:
#         f1(x) = (1 + a*x + b*x^2) / (2 + 2b/3) for x in [-1,1]
# Imagine that we are interested in both a and b.
#----------------------------------------------------------------------------------- #
#   Note that in order to be a PDF, the function is normalised.
#   However, it also need to be strictly positive, which in turn requires:
#   1 + a + b > 0  =>  b+a > -1  and
#   1 - a + b > 0  =>  b-a > -1  
#----------------------------------------------------------------------------------- #
def beta(x,a,b) :
    if (abs(x) <= 1.0) : return (1.0 + a*x + b*x*x) / (2.0 + 2.0*b/3.0)
    else               : return 0.0
#----------------------------------------------------------------------------------- #

Npoints_beta = 10000

# Choose the constants a and b, remembering the two constraints:
# b+a > -1  and  b-a > -1
a = 0.5
b = 0.5

Hist_beta = TH1F("Hist_beta", ";x;Frequency", 100, -1.0, 1.0)
binwidth_beta = (1.0 + 1.0) / 100

# Accept-Reject method:
Nhit = 0
while (Nhit < Npoints_beta) :
    x = -1.0 + 2.0 * r.Uniform()
    y = (1.0+a+b) * r.Uniform()
    if (y < beta(x,a,b)) :
        Nhit += 1
        Hist_beta.Fill(x)

canvas_beta = TCanvas("canvas_beta","", 50, 50, 600, 400)
Hist_beta.SetLineColor(kBlue)
Hist_beta.SetLineWidth(2)
Hist_beta.SetMinimum(0.0)

f_beta = TF1("f_beta", "[0] * (1.0 + [1]*x + [2]*x*x) / (2.0 + 2.0*[2]/3.0)", -1.0, 1.0)
f_beta.SetParameters(float(Npoints_beta*binwidth_beta), a, b)
Hist_beta.Fit("f_beta", "RL")
Hist_beta.Draw("e")

canvas_beta.Update()
if (SavePlots) : canvas_beta.SaveAs("Hist_beta.pdf")





#----------------------------------------------------------------------------------- #
# CASE 2: Double exponential distribution:
#         f2(x) = N * (f/r1*exp(-t/r1) + (1-f)/r2*exp(-t/r2)) for x in [0,inf]
# Imagine that in this case, we are most interested in r1.
#----------------------------------------------------------------------------------- #
# NOTE: The parameters r1 and r2 need to be positive, and f in [0,1],
#       in order for this to be a PDF.
#----------------------------------------------------------------------------------- #

Npoints_2exp = 2000
frac = 0.8   # Fraction that "belows to" first exponential
r1 = 6.0
r2 = 0.5

Hist_2exp = TH1F("Hist_2exp", ";x;Frequency", 100, 0.0, 20.0)
binwidth_2exp = (20.0 - 0.0) / 100

# Transformation method:
for i in xrange(Npoints_2exp) :
    if (r.Uniform() < frac) : t = r.Exp(r1)
    else :                    t = r.Exp(r2)
    Hist_2exp.Fill(t)

canvas_2exp = TCanvas("canvas_2exp","", 80, 80, 600, 400)
Hist_2exp.SetLineColor(kBlue)
Hist_2exp.SetLineWidth(2)
Hist_2exp.SetMinimum(0.0)

f_2exp = TF1("f_2exp", "[0] * ([1]/[2]*exp(-x/[2]) + (1.0-[1])/[3]*exp(-x/[3]))", 0.0, 20.0)
f_2exp.SetParameters(float(Npoints_2exp*binwidth_2exp),frac,r1,r2)
Hist_2exp.Fit("f_2exp", "R")
Hist_2exp.Draw("e")

canvas_2exp.Update()
if (SavePlots) : canvas_2exp.SaveAs("Hist_2exp.pdf")




#----------------------------------------------------------------------------------- #
# CASE 3: Cauchy distribution (also known as Breit-Wigner):
#         f2(x) = N / ((x-m)**2 + L**2) for x in [-inf,inf]
# In this case, we want m and L.
#----------------------------------------------------------------------------------- #
# NOTE: The Cauchy distribution is NOT a PDF, but nevertheless a function one can
#       fit with, and one that is relevant in several fields of physics.
#----------------------------------------------------------------------------------- #

Npoints_bw = 100
m = 2.0
L = 10.2

Hist_bw = TH1F("Hist_bw", ";x;Frequency", 100, -20.0, 20.0)
binwidth_bw = (20.0 + 20.0) / 100

# Transformation method:
for i in xrange(Npoints_bw) :
    x = r.BreitWigner(m,L)
    Hist_bw.Fill(x)

canvas_bw = TCanvas("canvas_bw","", 110, 110, 600, 400)
Hist_bw.SetLineColor(kBlue)
Hist_bw.SetLineWidth(2)
Hist_bw.SetMinimum(0.0)

f_bw = TF1("f_bw", "[0] * 2.0 / 3.1416 / ((x-[1])**2 + [2]**2)", -20.0, 20.0)
f_bw.SetParameters(float(Npoints_bw*binwidth_bw), m, L)
Hist_bw.Fit("f_bw", "R")
Hist_bw.Draw("e")

canvas_bw.Update()
if (SavePlots) : canvas_bw.SaveAs("Hist_bw.pdf")





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