import numpy as np
import scipy.stats as st
import matplotlib.pyplot as plt
import time

np.random.seed(42)

#defines number of samples
N = 10000

#generates samples X and Y
X = st.norm.rvs(size = (N,))
Y = st.norm.rvs(size = (N,))

#computes R and theta
R = np.sqrt(X**2 + Y**2)
theta = np.zeros(X.shape) #preallocates
#Does the actual computation
for i in range(X.shape[0]):
    angle = np.arctan(Y[i]/X[i])
    theta[i] = (X[i] > 0) * (Y[i] >= 0) * (angle) + (X[i] < 0) * (angle + np.pi)\
             + (X[i] > 0) * (Y[i] <= 0) * (angle + 2 * np.pi)

x = np.linspace(0, 10, 1001)

#fig = plt.figure(figsize = (9,4))
#ax1 = fig.add_subplot(121)
#ax1.plot(x, st.expon.pdf(x, scale = 2.), linewidth = 1,)
#ax1.hist(R**2, bins = 100, density = True)
#ax2 = fig.add_subplot(122)
#ax2.hist(theta, bins = 100, density = True)

#plt.show()

N = [10, 100, 1000, 10000]

times_AR = np.zeros(4)
times_BM = np.zeros(4)
fig = plt.figure(figsize = (8,8))
for n in N:
    # ---- Acc-Rej method

    i = 0
    j = 0
    C = np.exp(1./2) * np.sqrt(2./np.pi)
    X_acc = np.zeros(n)
    start = time.time()
    while i < n:
        B = st.bernoulli.rvs(0.5)
        Y = st.expon(scale = 2).rvs()
        X_prop = (2 * B - 1.) * Y
        U = st.uniform.rvs()
        if U < st.norm.pdf(X_prop) / (C * np.exp(-np.abs(X_prop) / 2.)):
            X_acc[i] = X_prop.copy()
            i = i + 1
            j = j + 1
    end_acc = time.time()

	# ---- Box-Muller method
    # this is a very inefficient implementation of the BM algorithm
    #as it not vectoized. However, we chose to do so in order to have a 
    # more fair comparison between A.R and B.M For a more efficient
    # implementation, please uncomment bellow.
    #	U = st.uniform.rvs(size = (n,))
    #	V = st.uniform.rvs(size = (n,))
    #	rho = np.sqrt(- 2 * np.log(U))
    #	theta = 2 * np.pi * V
    #	X_bm = rho * np.cos(theta)
    #	Y_bm = rho * np.sin(theta)
    #inefficient implementation
    U=np.zeros(n)
    V=np.zeros(n)
    X_bm=np.zeros(n)
    Y_bm=np.zeros(n)

    for i in range(n):
        U[i] = st.uniform.rvs(size = (1,))
        V[i] = st.uniform.rvs(size = (1,))
        rho=np.sqrt(-2*np.log(U[i]))
        theta=2*np.pi*V[i]
        X_bm[i] = rho * np.cos(theta)
        Y_bm[i] = rho * np.sin(theta)

    end_bm = time.time()

    time_AR = end_acc - start
    time_BM = end_bm - end_acc

    # print end_acc, end_bm
    # print(time_AR, time_BM)

    x = np.linspace(-4, 4, 1001)
    if n == 10:
        i0 = 1
    elif n == 100:
        i0 = 2
    elif n == 1000:
        i0 = 3
    else:
        i0 = 4

    #plots results
    times_AR[i0-1] = time_AR
    times_BM[i0-1] = time_BM
    ax = fig.add_subplot(2,2,i0)
    ax.plot(x, st.norm.pdf(x), linewidth = 1.)
    ax.hist(X_acc, bins = int(np.min([n**1/3,100])), density = True, alpha = 0.5)
    ax.hist(X_bm, bins = int(np.min([n**1/3,100])), density = True, alpha = 0.5)
    ax.set_title('n = ' + str(n))
    # plt.legend(['Normal density','AR','BM'])

handles, labels = ax.get_legend_handles_labels()
fig.legend(['Normal density', 'AR', 'BM'],
           loc='lower center',
           ncol=3,
           bbox_to_anchor=(0.5, 0.02))
plt.tight_layout(rect=[0, 0.05, 1, 1])
plt.show()

plt.loglog(N, times_AR, '-*', label = 'Acc-Rej')
plt.loglog(N, times_BM, '-x', label = 'Box-Muller')
plt.loglog(N,0.002* np.array(N),color='r',linestyle= '--', label = 'N')
plt.legend()
plt.ylabel('CPU time (s)')
plt.xlabel('samples (n)')
plt.show()