#! /usr/bin/env python3
#
def walk_3d_test ( ):

  import matplotlib.pyplot as plt
  import numpy as np

  print ( '' )
  print ( 'walk_3d_test():' )
  print ( '  If a walker randomly takes steps in 3D' )
  print ( '  on a regular grid, how many steps will be taken on average' )
  print ( '  to reach a distance of L?' )

  print ( '' )
  print ( '  To answer this question, consider each value of L' )
  print ( '  from 0 to 10, take many random walks, and determine' )
  print ( '  S, the average number of steps taken.' )
  print ( '' )
  print ( '  Theory says S = L^2' )

  n = 10000
  n = 1000

  lvec = np.linspace ( 0, 10, 11 )
  step_average = np.zeros ( 11 )
  for l in range ( 0, 11 ):
    step_sum = 0
    for i in range ( 0, n ):
      step_sum = step_sum + walk_3d ( l )
    step_average[l] = step_sum / n
#
#  Plot distance L versus steps S.
#
  plt.clf ( )
  plt.plot ( lvec, step_average, 'b-' )
  plt.plot ( lvec, lvec**2, 'r.', markersize = 20 )
  plt.grid ( True )
  plt.title ( 'Distance versus steps in 3D' )
  plt.xlabel ( '<-- Distance in steps-->' )
  plt.ylabel ( '<-- Average steps taken -->' )
  plt.legend ( [ 'Observed', 'Predicted' ] )
  filename = 'walk_3d.png'
  plt.savefig ( filename )
  print ( '  Graphics saved as "' + filename + '"' )
  plt.close ( )

  return

def walk_3d ( l ):

## walk_3d() takes a random 3D walk until reaching a distance of L units.
#
#  Discussion:
#
#    Each step is of unit size.
#
#  Input:
#
#    integer l: the distance at which the walk will terminate.
#
#  Output:
#
#    integer steps: the number of steps taken to reach a distance of L.
#
  import numpy as np

  directions = [ 'N', 'S', 'E', 'W', 'U', 'D' ]
  location = np.array ( [ 0, 0, 0 ] )
  steps = 0

  while ( np.linalg.norm ( location ) < l ):
    steps = steps + 1
    d = np.random.choice ( directions )
    if ( d == 'E' ):
      location[0] = location[0] + 1
    elif ( d == 'W' ):
      location[0] = location[0] - 1
    elif ( d == 'N' ):
      location[1] = location[1] + 1
    elif ( d == 'S' ):
      location[1] = location[1] - 1
    elif ( d == 'U' ):
      location[2] = location[2] + 1
    elif ( d == 'D' ):
      location[2] = location[2] - 1

  return steps

if ( __name__ == "__main__" ):
  walk_3d_test ( )