#
#  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_grid():#
    xdelta,ydelta = random.choice([ (0,1), (1,0),(-1,0), (0,-1)])   
    return xdelta, ydelta

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) 
    for i in range(steps-1):
        xdelta, ydelta = rand_angle()
        xwalk[i+1] = xwalk[i] + xdelta 
        ywalk[i+1] = ywalk[i] + ydelta 
    return xwalk, ywalk
#
#   Enter number of walks to plot and check entered value
#
num_walks = int(raw_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 walk 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()