#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Notice that this takes quite a bit of time to run. 
"""
import numpy as np
import matplotlib.pyplot as plt

np.random.seed(42)

def latticeRW(n=100,X0=0,N=4,aa=0.25):
	X=np.zeros(n+1)
	X[0]=X0
	t=np.arange(0,n+1)
	for i in range(n):
		xcurr=X[i]
		pup=(1-xcurr/2/N**2)*aa*(np.abs(xcurr)<=2*N**2)
		plow=(1+xcurr/2/N**2)*aa*(np.abs(xcurr)<=2*N**2)
		r=np.random.random(1)
		if (r<=plow):                  #decreses state
			X[i+1]=X[i]-1
		elif( (plow<r)*(r<=plow+pup)): # increses state
			X[i+1]=X[i]+1
		else:
			X[i+1]=X[i]
	return X,t 


#
#   Defines some parameters
#
N=4
x0=0
aa=1/4
N_runs=int(1E3) #this might be quite a large number 
n=500
x1=np.zeros((N_runs,n+1))
x2=np.zeros((N_runs,n+1))
x4=np.zeros((N_runs,n+1))
M=np.zeros((n+1,n+1))
tau=np.linspace(-1,1,n+1)

# runs the chain many times and obtains the quantities needed
for i in range(N_runs):
    x1[i,:]=latticeRW(n,x0,N,aa)[0]
    x2[i,:]=x1[i,:]**2.0
    x4[i,:]=x1[i,:]**4.0

#computes moments  
e1=np.mean(x1,0)    
e2=np.mean(x2,0)**0.5    
e4=np.mean(x4,0)**0.25    
Mx=np.mean(M,0)

plt.plot(e1);
plt.plot(e2);
plt.plot(e4);
plt.grid(True,which='both')
plt.legend([r'$m_1$',r'$m_2$',r'$m_4$'])
plt.show()

# computes the MGF
tauvec=np.linspace(-1,1,n+1)
ntau=n+1
MGF=np.zeros((n+1,ntau))
for r in range(N_runs):
    for i in range(ntau):
        MGF[:,i]=MGF[:,i]+np.exp(tauvec[i]*x1[r,:])
MGF=MGF/N_runs
mu=e1[-1]
si2=e2[-1]**2.0-mu**2
plt.semilogy(tauvec,MGF[-1,:])
plt.semilogy(tauvec,np.exp(tauvec*mu+si2/2*tauvec**2.));
plt.legend(['emp. MGF','MGF of Normal distr.']);
plt.xlabel(r'$t$');    
    
plt.grid(True,which='both')

plt.show()