import numpy as np
import matplotlib.pyplot as plt
import scipy.stats as st
from matplotlib import rc
#
#   imports graphical paramters
#
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'
import matplotlib
latex_preamble = r'\usepackage{amsmath} \usepackage{amssymb}'
matplotlib.rcParams.update({
    'text.usetex': True,
    'text.latex.preamble': latex_preamble
})

np.random.seed(42)

# functions
def genWalk(prob,X0,Nstop,T):
    X=X0;
    n=0;
    while((X[n]<Nstop)*(n<=T)):
        X=np.append(X, X[n]-1+2*np.random.binomial(1,prob))
        n=n+1
    t=np.arange(0,n+1)
    return X,t
    
def weightfun(prob,Y):
    dY=np.diff(Y)
    # rvec is a vector where only one of the two terms in 
    # the sum below appears in each component of the vector
    rvec=1/(2*prob)*(dY==1) + 1/(2*(1-prob))*(dY==-1) 
    rval=np.prod(rvec)
    return rval

# problem set-up
alist = np.linspace(0.33,0.95,16)
na = np.size(alist);
Nstop = 4;
T = 10;
NMC = int(1e3);
X0 = np.zeros(1);
# for confidence intervals
aa = 1-0.95;
cval = st.norm.ppf(1-aa/2);
# let's get started
varis = np.zeros(na);
varmc = np.zeros(na);

for j in range(na):
    alpha = alist[j];
    print(1,'a=  ',alpha);
    Y = np.zeros(int(NMC),);
    tauNaive = np.zeros(int(NMC),);
    for i in range(NMC):
        # CMC estimator
        Xnaive,tnaive, = genWalk(0.5,X0,Nstop,T);
        tauNaive[i] = tnaive[-1]<T;
        # IS estimator
        X,t = genWalk(alpha,X0,Nstop,T);
        Y[i] = (t[-1]<T)*weightfun(alpha,X);
        
    varmc[j]=np.var(tauNaive)  
    varis[j] = np.var(Y);
    print('MC confidence interval (aa=',0.05,', NM=',NMC,'): ',
          round(np.mean(tauNaive),4),' +/- ',round(cval*np.std(tauNaive)/np.sqrt(NMC),4));
    print('IS confidence interval (aa=',0.05,', NM=',NMC,'): ',
          round(np.mean(Y),4),' +/- ',round(cval*np.std(Y)/np.sqrt(NMC),4));
    print("----------------------------------------------------------------")
np.var(tauNaive)

plt.plot(alist,varmc/varis);
plt.xlabel(r'$a$')
plt.ylabel(r'$V_{MC}[Z]/V_{IS}[Z]$')
plt.title(r'Relative gain vs $a$')
plt.show()