import numpy as np
import scipy.stats as st

np.random.seed(42)

N=100000 #Number of samples
N2=int(N/2) #Number of samples for AV

X0=0 #initial step
K=5 #stopping time
T=10 #maximum stopping time
alpha = 0.05
Z_mc=np.zeros(N)
Z_mc2=np.zeros(N2)
Z_av=np.zeros(N2)
C_alpha=st.norm.ppf(1-alpha/2)
a=1/3

# we define the random walker for the CMC and the AV 
#----------------------------------------------------
def rw(a,T,K,X0=0):
    
    finished_mc=0
    finished_av=0
    x_mc=X0
    x_av=X0
    for i in range(T):        
        u=np.random.random(1)        
        # moves crude MC
        if u<a:
            x_mc=x_mc+1
        else:
            x_mc=x_mc-1
        if x_mc>=K:
            finished_mc=1       
         # moves AV with the same random realization         
        if  1-u<a:
            x_av=x_av+1
        else:
            x_av=x_av-1
        if x_av>=K:
            finished_av=1  

           
    return finished_mc,finished_av
#----------------------------------------------------
    
for i in range(int(N)):    
    Z_mc[i],_=rw(a,T,K,X0)

#computes statistics
mean_cmc=np.mean(Z_mc)
var_cmc=np.var(Z_mc)
ci_cmc=C_alpha*np.sqrt(var_cmc)/np.sqrt(N)

# We now compute the AV estimator 
# We begin by doing a crude MC estimator 
for i in range(int(N2)):    
    Z_mc2[i],Z_av[i]=rw(a,T,K,X0)

Y_av=0.5*(Z_mc2+Z_av)
mean_av=np.mean(Y_av)
C=np.cov(Z_mc2,Z_av)
ci_av=C_alpha*np.sqrt( (np.sum(C))/(2*N)  )

#Print results
print('--------------------------------------------------')
print('Reduction for a='+str(a))
print('mean MC '+str(mean_cmc)+' +- '+str(round(ci_cmc,5)))
print('mean AV '+str(mean_av)+' +- '+str(round(ci_av,5)))
print('Reduction '+str(round(ci_cmc/ci_av,5)))

print('--------------------------------------------------')
print('Sample covariance matrix')
print(C)