import numpy as np
import matplotlib.pyplot as plt

# Exercice 1:
beta = 1. / 2.
tau = 10.
gamma = 1. / 12.
mu = 1. / 50.
C0 = 1. / 2.

def func(scir):
    S, C, I, R = scir
    return np.array([- beta * S * I, beta * S * I - C / tau, C / tau - (gamma + mu) * I, gamma * I])
    

h = 0.01 # h should be much smaller than 1/beta, tau, 1/mu, and 1/\gamma (the characteristic times of the system), so here h << 2. 
tf = 100.
time = np.arange(0, tf + h, h)
scir = np.array([1. - C0, C0, 0., 0.])
scir_all = np.zeros((time.size, 4))
scir_all[0,:] = scir
for i in range(1, time.size):
    k1 = func(scir) * h
    k2 = func(scir + k1 / 2.) * h
    k3 = func(scir + k2 / 2.) * h
    k4 = func(scir + k3) * h
    scir += (k1 + 2. * k2 + 2. * k3 + k4) / 6.
    scir_all[i,:] = scir

plt.figure()
plt.plot(time, scir_all[:,0], 'g', label='S')
plt.plot(time, scir_all[:,1], c='orange', label='C')
plt.plot(time, scir_all[:,2], 'r', label='I')
plt.plot(time, scir_all[:,3], 'b', label='R')
plt.plot(time, np.sum(scir_all, axis=1), c='purple', label='A')
plt.plot(time, 1. - np.sum(scir_all, axis=1), 'k', label='D')
plt.xlim(0., np.amax(time))
plt.grid()
plt.xlabel('Time')
plt.ylabel('Populations')
plt.legend()
plt.show()

h = 0.01
t = 0.
i_before = - np.inf
i = 0.
scir = np.array([1. - C0, C0, 0., 0.])
while i > i_before: # Integration until the maximum
    i_before = i
    k1 = func(scir) * h
    k2 = func(scir + k1 / 2.) * h
    k3 = func(scir + k2 / 2.) * h
    k4 = func(scir + k3) * h
    scir += (k1 + 2. * k2 + 2. * k3 + k4) / 6.
    i = scir[2]
    t += h
imax = i_before
while i > 0.01 * imax: # Integration until 0.01I_{max}
    k1 = func(scir) * h
    k2 = func(scir + k1 / 2.) * h
    k3 = func(scir + k2 / 2.) * h
    k4 = func(scir + k3) * h
    scir += (k1 + 2. * k2 + 2. * k3 + k4) / 6.
    i = scir[2]
    t += h
print('t_{1%}=', t)

# Exercice 2
def crossing_proba(ell, N):
    theta = np.random.uniform(0., 2. * np.pi, N)
    x = np.random.uniform(0., 10., N)
    crossings = np.floor(x - ell / 2 * np.cos(theta)) != np.floor(x + ell / 2 * np.cos(theta))
    return float(np.sum(crossings)) / float(N)
    
def crossing_proba_th(z):
    z_arr = np.array(z)
    p = np.zeros_like(z_arr)
    mask = z_arr < 1.
    p[mask] = 2 * z_arr[mask] / np.pi
    mask = z_arr >= 1.
    p[mask] = 2. / np.pi * (z_arr[mask] - np.sqrt(z_arr[mask] ** 2 - 1) + np.acos(1. / z_arr[mask]))
    if type(z) == float:
        return float(p)
    else:
        return p

ell = np.arange(0., 2.1, 0.1)
proba_Buffon = np.array([crossing_proba(l, 100000) for l in ell])
proba_th = crossing_proba_th(ell)

plt.figure()
plt.plot(ell, proba_th, '-k', label='Exact')
plt.scatter(ell, proba_Buffon, c='r', label='Monte Carlo')
plt.xlabel('Needle length')
plt.ylabel('Crossing probability')
plt.legend()
plt.grid()
plt.show()

# Decay as N^{-1/2} (except for Monte Carlo integration).
Nval = np.logspace(2, 8, 7, dtype=int)
pi_est = np.zeros(Nval.size)
ell = 1.
for i, N in enumerate(Nval):
    pi_est[i] = 2. * ell / crossing_proba(ell, N)
plt.figure()
plt.loglog(Nval, np.absolute(pi_est / np.pi - 1.), '-o')
plt.grid()
plt.xlabel('Number of needles')
plt.ylabel(r'Relative error on the estimate of $\pi$')
plt.show()