import numpy as np 
import scipy.stats as st 
import matplotlib.pyplot as plt 
from matplotlib import rc
#graphical parameters
############################################################################### 
plt.style.use('ggplot')
rc('font',**{'family':'serif','serif':['Computer Modern Roman'],
     'size' : '12'})
rc('text', usetex=True)
rc('lines', linewidth=2)
plt.rcParams['axes.facecolor']='w'
###############################################################################

def Algo1(N0, lamda, alpha, eps):
    M = N0
    X = lamda(M)
    sig = np.std(X)

    while 2*(1.-st.norm.cdf(np.sqrt(M)*eps / sig)) > alpha:

            M = 2 * M
            X = lamda(M)
            sig = np.std(X)

    barX = np.mean(lamda(M))

    return barX, M

def Algo2(N0, lamda, alpha, eps):
    M = N0
    X = lamda(M)
    barX = np.mean(X)
    sig = np.std(X)

    while 2*(1.-st.norm.cdf(np.sqrt(M)*eps / sig)) > alpha:
        Xnew = lamda(1)
        barX = M*barX/(M+1) + 1*Xnew/(M+1)
        sig = np.sqrt(  (M-1)*sig**2/M + (Xnew-barX)**2/(M+1)  )
        M = M +1

    barX = np.mean(lamda(M))
    return barX, M

# Variable introduction
alpha = 10 ** (-1.5)
epsilon = 1./10
N0 = np.arange(10, 51, 10)

pNPareto = np.zeros(len(N0))
randPareto = lambda n: (1 - st.uniform.rvs(size = (n,))) ** (-1/ 3.1)
meanPareto = 3.1/2.1

pNLognormal = np.zeros(len(N0))
randLognorm = lambda n: np.exp(st.norm.rvs(size = (n,) ))
meanLognormal = np.exp(1./2)

pNUniform = np.zeros(len(N0))
randUni = lambda n: (2 * st.uniform.rvs(size = (n,)) - 1.)
meanUniform = 0

# Number of times to run each simulation for a given N0
N = int(20/alpha)

# Algorithm 1
for i in range(len(N0)):
    for j in range(N):
        print('Running simulation ',j,'/',N, 'for N0 = ', N0[i], ' for Algorithm 1')
        barX,M = Algo1(N0[i], randPareto, alpha, epsilon)
        pNPareto[i] = pNPareto[i] + float(np.abs(barX - meanPareto) > epsilon) / N

        barX,M = Algo1(N0[i], randLognorm, alpha, epsilon)
        pNLognormal[i] = pNLognormal[i] + float(np.abs(barX - meanLognormal) > epsilon) / N

        barX,M = Algo1(N0[i], randUni, alpha, epsilon)
        pNUniform[i] = pNUniform[i] + float(np.abs(barX - meanUniform) > epsilon) / N 

plt.figure()
eBars = 2 * np.sqrt(pNPareto * (1 - pNPareto) / N)
plt.errorbar(N0, pNPareto, yerr = eBars, label = 'Pareto')

eBars = 2 * np.sqrt(pNLognormal * (1 - pNLognormal) / N)
plt.errorbar(N0, pNLognormal, yerr = eBars, fmt = '-s', label = 'Lognormal')

eBars = 2 * np.sqrt(pNUniform * (1 - pNUniform) / N)
plt.errorbar(N0, pNUniform, eBars, fmt = '--.', label = 'Uniform')

plt.plot(N0, alpha * np.ones(len(N0)), 'r-o', label = r'$\alpha$')
plt.xlabel(r'$N_0$')
plt.ylabel(r'$\bar{p}_N$')
plt.legend()
plt.show()

# Algorithm 2
for i in range(len(N0)):
    for j in range(N):
        print('Running simulation ',j,'/',N, 'for N0 = ', N0[i], ' for Algorithm 2')
        barX,M = Algo2(N0[i], randPareto, alpha, epsilon)
        pNPareto[i] = pNPareto[i] + float(np.abs(barX - meanPareto) > epsilon) / N

        barX,M = Algo2(N0[i], randLognorm, alpha, epsilon)
        pNLognormal[i] = pNLognormal[i] + float(np.abs(barX - meanLognormal) > epsilon) / N

        barX,M = Algo2(N0[i], randUni, alpha, epsilon)
        pNUniform[i] = pNUniform[i] + float(np.abs(barX - meanUniform) > epsilon) / N 

plt.figure()
eBars = 2 * np.sqrt(pNPareto * (1 - pNPareto) / N)
plt.errorbar(N0, pNPareto, yerr = eBars, label = 'Pareto')

eBars = 2 * np.sqrt(pNLognormal * (1 - pNLognormal) / N)
plt.errorbar(N0, pNLognormal, yerr = eBars, fmt = '-s', label = 'Lognormal')

eBars = 2 * np.sqrt(pNUniform * (1 - pNUniform) / N)
plt.errorbar(N0, pNUniform, eBars, fmt = '--.', label = 'Uniform')

plt.plot(N0, alpha * np.ones(len(N0)), 'r-o', label = r'$\alpha$')
plt.xlabel(r'$N_0$')
plt.ylabel(r'$\bar{p}_N$')
plt.legend()
plt.show()