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

  import matplotlib.pyplot as plt
  import numpy as np

  print ( '' )
  print ( 'walk_1d_test():' )
  print ( '  If a walker randomly takes steps of +1 and -1,' )
  print ( '  how many steps will be taken on average to reach' )
  print ( '  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

  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_1d ( 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.xlabel ( '<-- Distance in steps-->' )
  plt.ylabel ( '<-- Average steps taken -->' )
  plt.legend ( [ 'Observed', 'Predicted' ] )
  filename = 'walk_1d.png'
  plt.savefig ( filename )
  print ( '  Graphics saved as "' + filename + '"' )
  plt.close ( )

  return

def walk_1d ( l ):

## walk_1d() takes a random 1D 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 = [ -1, +1 ]
  location = 0
  steps = 0

  while ( -l < location and location < l ):
    steps = steps + 1
    location = location + np.random.choice ( directions )

  return steps

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