#  Simple in-line (no module) python random 
#  walk program in 2D.
#  Performs and plots up to 5 random walks.
#  Walkers start at the origin: (x,y) = (0,0)
#  At each step, walker moves only one unit in 
#  a random direction : <up, down, left,right> 
import sys, random
import numpy as np
import matplotlib.pyplot as plt
random.seed(None)        # Seed generator, None => system clock

def rand_angle(): 
    rand_angle = random.uniform(0, 2*np.pi)
    xdelta = np.cos(rand_angle)
    ydelta = np.sin(rand_angle)
    return xdelta, ydelta

def random_walk2d(steps):
    """ Performs 2D random walk with number of steps = steps 
        At each step, walker moves either <up, down, left,right>  """
    xwalk = np.zeros(steps)
    ywalk = np.zeros(steps)
###########################################################################
#    Add your code HERE to complete this function.
#    Perform random walk in each axis (xwalk, ywalk)using rand_angle 
#          (or rand_grid) to get lengths of each step.
#    HINT: Use for loop over the number of steps (steps). Change xwalk and
#      ywalk values based on step lengths (xdelta, ydelta) from rand_angle
#    Return xwalk,ywalk arrays
###########################################################################
    return xwalk,ywalk 
#   Enter number of walks to plot and check entered value
num_walks = int(input('Enter number of walks to do (1-5):'))
if num_walks <= 0 or num_walks >5:
    print ('Number of walks selected = ', num_walks,' - must be between 1 and 5')
    sys.exit(0)
#  Sets number of steps in each random walk    
num_steps = 1000 
#  Allows each walk to use a different color/line type on plot 
line_type = ['r-','b-','k-','g-','m-']
#
#  Perform random walks and plot result for each walk
#   
plt.clf     #  Clear the graph figure
for i in range(num_walks):      
    xwalk1, ywalk1 = random_walk2d(num_steps)
    plt.plot(xwalk1, ywalk1, line_type[i], label = 'walk'+ str(i+1))
    plt.plot(xwalk1[-1],ywalk1[-1], 'k*', ms=16)  # Put large * to mark end of walk
#  Create bold lines on each axis through origin
plt.axhline(0, color='black', lw=1)  
plt.axvline(0, color='black', lw=1) 
#  Set plot title , legend, and grid 
plt.title('Two-dimensional random walks(' + str(num_walks) + ')' )
plt.legend(loc="upper left")
plt.grid(True)
plt.show()