#
#  Simple python random walk in 1D
#  Performs and plots up to 5 random walks
#
import random
import numpy as np
import matplotlib.pyplot as plt

def random_walk(steps):
    """  Performs 1D random walk with number steps = steps """
    walk = np.zeros(steps+1)
    for i in range(steps):
        if random.uniform(0, 1) < .5:
            step =  1  
        else:
            step = -1
        walk[i+1] = walk[i] + step
    return walk
#
#   Enter number of walks to plot 
#
num_walks = int(input('Enter number of walks to do (0-5):  '))
num_steps = 1000 
x = range(num_steps+1)
#  Allows each walk to use a different color/line type on plot 
line_type = ['r-','b-','k-','g-','m-']
#
#  Perform random walk and plot result for each walk
# 
####################################################
#
#   YOUR CODE GOES HERE !!
#   Hint: Use for loop to execute and plot each walk 
#  plot function is:
#   plt.plot(x, walk, line_type[i], label= "walk"+ str(i) ) 
#
#####################################################
#
#  Set plot title , legend, and grid 
#
plt.title('One-dimensional random walks(' + str(num_walks) + ')' )
plt.legend(loc="upper left")
plt.grid(True)
plt.show()