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

np.random.seed(42)

##### Generate circle plots to illustrate Monte Carlo 
Ns = [10,100,1000,10000]
markersizes=[2,1,0.5,0.25]
n = 2
I = np.pi/4 # Area of circle divided by area of square
fig, ax = plt.subplots(1,len(Ns))

for i in range(len(Ns)):
    circle = plt.Circle((0,0),1,color='black',fill=False)
    N = Ns[i]
    ms = markersizes[i]
    print('Running for N =',N)

    # Generate N samples of n-dimension vectors
    X = st.uniform.rvs(loc=-1,scale=2,size=(N,n))

    # Compute Z, I_N and relative error
    Z = [int(np.linalg.norm(x)<1) for x in X]
    I_N = np.mean(Z)

    # Collect points that are inside and outside
    Xin = [x for (x,z) in zip(X,Z) if z==1]
    Xout = [x for (x,z) in zip(X,Z) if z==0]

    # Plotting
    ax[i].plot([x[0] for x in Xin], [x[1] for x in Xin],'+',markersize=ms)
    ax[i].plot([x[0] for x in Xout], [x[1] for x in Xout],'+',markersize=ms)
    ax[i].set_xlim((-1,1))
    ax[i].set_ylim((-1,1))
    titleString = 'N = %d' %N
    ax[i].set_title(titleString,fontsize=8)
    ax[i].set(aspect='equal')
    ax[i].add_patch(circle)
    ax[i].tick_params(axis='x',labelsize=5)
    ax[i].tick_params(axis='y',labelsize=5)

plt.tight_layout()
plt.show() 

# Generate additional data for the convergence graph
Ns_all = list(range(10,10001))
Xs_all = st.uniform.rvs(loc=-1,scale=2,size=(Ns_all[-1],n))
Zs_all = [int(np.linalg.norm(x)<1) for x in Xs_all]
mean = np.mean(Zs_all[:9])
var = np.var(Zs_all[:9])
means = [] 
variances = []
true_errors = []

for N in Ns_all:
    mean_new = N/(N+1)*mean + Zs_all[N-1]/(N+1)
    var_new = (N-1)/N*var + (Zs_all[N-1]-mean)**2/(N+1)
    means.append(mean_new)
    variances.append(var_new)
    true_errors.append(abs(mean_new-I)/I)

    mean = mean_new
    var = var_new

plt.figure()
plt.loglog(Ns_all, true_errors,color='black',label='Error')
plt.loglog(Ns_all, [N**-0.5 for N in Ns_all],'--',color='red',label=r'$N^{-1/2}$')
plt.xlabel(r'$N$')
plt.ylabel(r'$|I-I_N|/I$')
plt.legend()
plt.grid(which='both')
plt.gca().set_aspect(0.2)
plt.show()

# Plot the confidence intervals
plt.figure()
plt.semilogx([Ns_all[0], Ns_all[-1]], [I,I], label=r'True Mean')
plt.semilogx(Ns_all, means, label=r'Estimated Mean')
upperlim = [(m + st.norm.ppf(0.975)*np.sqrt(v/N)) for m,v,N in zip(means, variances, Ns_all)]
lowerlim = [(m - st.norm.ppf(0.975)*np.sqrt(v/N)) for m,v,N in zip(means, variances, Ns_all)]
plt.semilogx(Ns_all, upperlim,'--',color='red',label=r'Estimated $95\%$ Confidence Interval')
plt.semilogx(Ns_all, lowerlim,'--',color='red')
plt.xlabel(r'N')
plt.ylabel(r'I')
plt.legend()
plt.grid(which='both')
plt.gca().set_aspect(3.0)
plt.show()