import matplotlib.pyplot as plt
import numpy as np
import scipy.integrate as integrate
import random

def Gaussian(x, mu, sigma):
    return (1.0/(sigma*np.sqrt(2*np.pi)))*np.exp(-(x-mu)**2/(2*sigma**2))

# Problem 1
x = np.linspace(-5,5,1000)
f1 = Gaussian(x,0.0,1.0)
f2 = Gaussian(x,1.0,1.0)
f3 = Gaussian(x,0.0,2.0)

plt.plot(x, f1, label = "mu = 0, sigma = 1")
plt.plot(x, f2, label = "mu = 1, sigma = 1")
plt.plot(x, f3, label = "mu = 0, sigma = 2")
plt.legend(loc="upper left")

plt.savefig('E1_Problem_1.pdf')
plt.clf()

print("Probability to produce the new particle with a mass of 205 GeV or more is " + str(round(100*(integrate.quad(lambda x: Gaussian(x,200.0, 2.0), 205.0, +np.inf))[0],2)) + "%.")
#print("Probability to produce the new particle with a mass with mass between 199 and 201 GeV " + str(100*(integrate.quad(lambda x: Gaussian(x,200.0, 2.0), 199.0, 201.0))[0]) + "%.")
print("Probability to produce the two particles with masses above 203 GeV is " + str(round(100*(integrate.quad(lambda x: Gaussian(x,200.0, 2.0), 203.0, +np.inf))[0]**2,2)) + "%.")


# Problem 2
x_histo_3 = []
x_histo_30 = []
x_histo_60 = []
x_histo_100 = []

for i in range(10000):
    x = 0.0
    N = 3
    for j in range(N):
        x = x + np.random.exponential(scale= j + 1.0, size=None)
    x_histo_3.append(x/N)

    x = 0.0
    N = 30
    for j in range(N):
        x = x + np.random.exponential(scale= j + 1.0, size=None)
    x_histo_30.append(x/N)

    x = 0.0
    N = 60
    for j in range(N):
        x = x + np.random.exponential(scale= j + 1.0, size=None)
    x_histo_60.append(x/N)

    x = 0.0
    N = 100
    for j in range(N):
        x = x + np.random.exponential(scale= j + 1.0, size=None)
    x_histo_100.append(x/N)

# Initialise the subplot function using number of rows and columns
figure, axis = plt.subplots(2, 2, constrained_layout=True)

axis[0, 0].hist(x_histo_3, bins=20)
axis[0, 0].set_title("N = 3")
axis[0, 0].set_xlabel("x")

axis[0, 1].hist(x_histo_30, bins=20)
axis[0, 1].set_title("N = 30")
axis[0, 1].set_xlabel("x")

axis[1, 0].hist(x_histo_60, bins=20)
axis[1, 0].set_title("N = 60")
axis[1, 0].set_xlabel("x")

axis[1, 1].hist(x_histo_100, bins=20)
axis[1, 1].set_title("N = 100")
axis[1, 1].set_xlabel("x")

plt.savefig('E1_Problem_2.pdf')
plt.clf()

# Problem 3 - 4
def pdf_1(x):
    return Gaussian(x, 0.0, 1.0)

def random_pdf(pdf, N, x_min, x_max):
    x_rand = []
    n_tries = 0
    while (len(x_rand) < N):
        n_tries += 2
        (x,y)= (random.uniform(x_min,x_max), random.uniform(0.0,1.0))
        if(pdf(x) > y):
            x_rand.append(x)
    print("I had to generate " + str(n_tries) + " random numbers to draw " + str(N) + " random numbers according to a given PDF.")
    return x_rand

N = 10000
x_rand = random_pdf(pdf_1, N, -5.0, 5.0)
x = np.linspace(-5,5,1000)

plt.title("Acceptance-rejection method")
plt.hist(x_rand, bins = 20, label="MC toys")
plt.plot(x, (N/2.5)*pdf_1(x), label="PDF")
plt.legend(loc="upper left")
plt.savefig('E1_Problem_34.pdf')
plt.clf()

# Problem 5*
def CDF(pdf, x_min = -5.0, x_max = 5.0, steps = 1000):
    cdf = []
    for i in np.linspace(x_min,x_max,steps):
        cdf.append((integrate.quad(lambda x: pdf(x), -np.inf, i))[0])
    return cdf

x = np.linspace(-5,5,1000)
cdf = CDF(pdf_1)

plt.plot(x, cdf, label = "CDF(x)")
plt.legend(loc="upper left")
plt.savefig('E1_Problem_5_1.pdf')
plt.clf()

def random_pdf_fast(pdf, N, x_min, x_max):
    x_rand = []
    x = np.linspace(x_min,x_max,1000)
    cdf = CDF(pdf, x_min, x_max)
    n_tries = 0
    while (len(x_rand) < N):
        n_tries += 1
        u = random.random()
        index = np.argmin(np.abs(np.array(cdf)-u))
        x_rand.append(x[index])
    print("I had to generate " + str(n_tries) + " random numbers to draw " + str(N) + " random numbers according to a given PDF.")
    return x_rand

x_rand = random_pdf_fast(pdf_1, 10000, -5.0, 5.0)

plt.title("Inversion method")
plt.hist(x_rand, bins = 20, label="MC toys")
plt.plot(x, (N/2.5)*pdf_1(x), label="PDF")
plt.legend(loc="upper left")
plt.savefig('E1_Problem_5_2.pdf')