from numpy.random import random as rng # import the rng, random number generator

import numpy as np # making math functions available
from pylab import * # get all pylab functions

num_steps = 1000 # number of steps per walk
num_walks = 100000 # total number of walks

x_final = np.empty(num_walks) # an empty list, will store the final x positions
y_final = np.empty(num_walks) # an empty list, will store the final y positions
displacement = np.empty(num_walks) # an empty list, will store the distance from the origin of the walk

for i in range(num_walks): # create four plots

    samples1 = (rng(num_steps) < 0.5)  # create list of random values for sample 1
    samples2 = (rng(num_steps) < 0.5)  # create list of random values for sample 2

    # convert random values into 1s and -1s
    x_steps = np.cumsum(2 * samples1 - 1)
    y_steps = np.cumsum(2 * samples2 - 1)

    # add the final x, y positions to the list
    x_final[i] = x_steps[-1]
    y_final[i] = y_steps[-1]

    # add the displacement of this walk to the list of all displacements
    displacement[i] = np.sqrt(x_steps[-1] ** 2 + y_steps[-1] ** 2)


# scatter plot of final positions of all the random walks
fig = plt.figure()
axes = fig.gca()

axes.set_xlim([-150,150]) # set domain to -150 to 150
axes.set_ylim([-150,150]) # set range to -150 to 150
axes.set_aspect('equal') # set aspect ratio of graph to equal width and height
grid(True) # include the cartesian dotted grid on the graph

# add the title, x, and y labels
title("Location After 1000 Steps")
xlabel("$\Delta$x")
ylabel("$\Delta$y")

scatter(x_final[:1000], y_final[:1000]) # plot first 1,000 random walks

# begin a new figure for the displacement chart
figure()

# set the x, y, and title labels
xlabel("$\sqrt{\Delta x^2 + \Delta y^2}$")
ylabel("Number of Walks")
title("Displacement After 1000 Steps")
hist(displacement, bins=25)
print(mean(displacement))
print(median(displacement))

# begin a new figure for the square displacement chart
figure()

# set the x, y, and title labels
xlabel("$\Delta x^2 + \Delta y^2$")
ylabel("Number of Walks")
title("Square Displacement After 1000 Steps")

n, bins, patches = hist(displacement ** 2, bins=25) # plot the histogram, and get the bins from it
bins_mean = [0.5 * (bins[i] + bins[i+1]) for i in range(len(n))] # get the middle of the histogram bins

# begin a new figure for the logarithmic square displacement chart
fig = figure()

# set the x, y, and title labels
xlabel("$\Delta x^2 + \Delta y^2$")
ylabel("Number of Walks")
title("Square Displacement After 1000 Steps")

logged = np.log(n) # take the log of the counts

semilogy(bins_mean, n, "+", color='red', markersize=12, mew=3, ms=20) # plot the log scatter plot

z = np.polyfit(bins_mean, logged, 1) # get the slope and y-intercept
p = np.poly1d(z)(bins_mean) # convert slope and y-intercept to y-values

semilogy(bins_mean, exp(p)) # plot the line of best fit, with log scale