import numpy as np
import matplotlib.pyplot as plt
from matplotlib import rc

np.random.seed(42)

#-------------------------------------
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
})

#-----------------------------------------------
def chain1(n):
    xn=np.zeros(n)
    for i in range(n-1):
        u=np.random.random(1)
        if u <p:
            a=1
        else:
            a=-1
        xn[i+1]= np.max((xn[i] + a,0))
    return xn

def chain2(n):
    zn=np.zeros(n)
    for i in range(n-1):
        u=np.random.random(1)
        if u <p:
            a=1
        else:
            a=-1
        zn[i+1]= np.abs(zn[i] + a)
    return zn

#computes empirical cdf    
def ecdf(data):
    """ Compute ECDF """
    x = np.sort(data)
    n = x.size
    y = np.arange(1, n+1) / n
    return(x,y)

#computes pi hat
def true_cdf_1(K,p):
    a0=(1-2*p)/(1-p) #this is the normalization constant
    pi_hat=(p/(1-p))**np.arange(K)*a0
    return pi_hat

#computes pi bar
def true_cdf_2(K,p):
    r=p/(1-p)
    b0=1/(1+(1/p)*(r/(1-r))) #this is the normalization constant
    print(b0)
    print(1/(1+1/(1-2*p)))
    pi_hat=(1/p)*(p/(1-p))**np.arange(1,K)*b0
    return np.concatenate([[b0],pi_hat])

n=100
p=1/8

N=10000
m1=np.zeros(N)
m2=np.zeros(N)
m3=np.zeros(N)
m4=np.zeros(N)
for i in range(N):
    m1[i]=chain1(n)[-1]
    m2[i]=chain1(n+1)[-1]
    m3[i]=chain2(n)[-1]
    m4[i]=chain2(n+1)[-1]

# computes empirical distribution 
x1,e1=ecdf(m1)    
x2,e2=ecdf(m2)    
x3,e3=ecdf(m3)    
x4,e4=ecdf(m4)    

# computes the CDF for Xn
K1=np.max(x1) #highest value obtained
pi_hat=true_cdf_1(K1,p) #obtaines values of pi hat
cdf_pi_hat=np.cumsum(pi_hat) #computes cdf
cdf_pi_hat=np.concatenate([ [0],cdf_pi_hat]) # adds 0 for plotting

# Computes the CDF for Zn
K2=np.max([np.max(x3),np.max(x4)]) #highest value obtained
pi_bar=true_cdf_2(K2,p) #obtaines values of pi bar
cdf_pi_bar=np.cumsum(pi_bar) #computes cdf
cdf_pi_bar=np.concatenate([ [0],cdf_pi_bar]) # adds 0 for plotting

#fixes parameter for plotting
fig, axes = plt.subplots(nrows = 1, ncols = 2, figsize = (12,4))

# plots cdf of Xn
axes[0].plot(x1,e1);
axes[0].plot(x2,e2,'--');
axes[0].step(np.arange(K1+1),cdf_pi_hat,'-*')
axes[0].set_title(r'$X_n$')
axes[0].legend([r'$X_{n}$',r'$X_{n+1}$',r'$\hat{\pi}$'])

# plots cdf of Zn
axes[1].plot(x3,e3);
axes[1].plot(x4,e4,'--');
axes[1].step(np.arange(K2+1),cdf_pi_bar,'-*')
axes[1].set_title(r'$Z_n$')
axes[1].legend([r'$Z_{n}$',r'$Z_{n+1}$',r'$\bar{\pi}$'])
plt.savefig('./convergence_distribution.png',bbox_inches='tight')