# -*- coding: utf-8 -*-
"""
Created on Sun Apr 18 16:42:22 2021
1D random walk as a CA for heat diffusion
NOTE doing CA in a high level language means it is very!! slow
@author: neil
"""

import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import norm

# Using 'w' random walkers, all starting at the middle of an 'n' long vector
# each walker takes 'steps' numbers of randomized steps
# Run the algorythm 'aver' times to average the results
n = 100        # number of cells that the random walker can occupy
w = 1000       # number of walkers
steps = 100    # number of steps a walker takes
aver  = 20    # number of realizations to average

x = np.zeros(n, dtype = int)  # create a vector to use as our random walk space
xn= np.zeros(n, dtype = int)  # a work space to move walkers from 'x' to new 'xn'
xa= np.zeros(n, dtype = int)  # an averaging space, to average the end results

for i in range(aver):  # doing 'aver' number of realizations
     
    x[n//2] = w    # put a pile of 'walkers' in the central location, 'xn' is zeros
    
    for j in range(steps):  # make all the walkers to walk for 'steps' number of steps
        
        m = np.random.randint(0,3,w) -1 # create enough random numbers for the walkers
        im = 0                          # an index to select the random directions in m
        
        for k in range(1,n-1):    # go to each of the cells (k is the cell number)

            while(x[k]):          # move each walker in the cell separately
                xn[k+m[im]] += 1  # use random "-1,0,1" to move walker to k-1, k, or k+1
                x[k]  -= 1        # moved one walker, proceed to next if there is one
                im    += 1        # increment 'im' to select the next unused random 'm'
        
        x = xn.copy()             # moved all the walkers, go to the next 'step'
        xn= np.zeros(n, dtype = int)  # clear 'xn' for next step or next average
        
    #at this point we have done the 'steps' for this realization
    xa += x   
    x = np.zeros(n, dtype = int)  # clear the 'board' for another realization
    
xa = xa/aver # average over all the realizations

fig = plt.figure()
ax  = fig.add_subplot(111)  
ax.plot(xa,"r-*")  #
ax.set_ylabel("CA Temperature")
ax.set_title("CA random walk, (compared to gaussian)")

# for comparison, plot a matching 'normal' bell curve centered on 50, blue curve
xnorm = np.linspace(0,n)
normal = np.exp(-0.5*((xnorm-50)/8)**2)*49 # scaled normal curve with variance of 8
ax.plot(xnorm,normal,"b-d")