import sys
import math
import random
import time
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as anim
import matplotlib.cm as cm
from scipy import stats

N = 100000				# number of walkers
ns = 10000				# number of steps to take

if len(sys.argv) > 1: N = int(sys.argv[1])
if len(sys.argv) > 2: ns = int(sys.argv[2])

x = np.zeros(N)
np.random.seed(12345)

deltax = 0.01			# step size
print(N, 'walkers,', ns, 'steps of', deltax)

t0 = time.time()
for i in range(ns):
    x += deltax*(2*np.random.randint(2, size=(N))-1)
    #x += np.random.uniform(-deltax, deltax, N)
    if i > 0 and i%(ns/10) == 0: print('step', i, 't =', time.time() - t0)

lim = 5.*np.sqrt(1.e-4*ns)
plt.xlim(-lim, lim)

# Display the data.

nb = 45
plt.hist(x, nb, density=True)	# display as a normalized probability density

# Display the normalized fundamental solution.

x = np.linspace(-lim, lim, 1000)
sig2 = deltax**2*ns
prob = np.exp(-x**2/(2*sig2))/math.sqrt(2*math.pi*sig2)
plt.plot(x, prob, 'r-')

plt.show()