#! /usr/bin/env python3
#
def walker_build ( n, x ):

#****************************************************************************80
#
#  Purpose:
#
#    WALKER_BUILD sets up the data for a Walker sampler.
#
#  Discussion:
#
#    From the probability vector X, this function creates the Y and A 
#    vectors used by the Walker sampler.
#
#  Modified:
#
#    20 February 2016
#
#  Author:
#
#    Original C version by Warren Smith.
#    Python version by John Burkardt.
#
#  Reference:
#
#    Warren Smith,
#    How to sample from a probability distribution,
#    April 18, 2002.
#
#  Parameters:
#
#    Input, unsigned int N, indicates the size of X.
#
#    Input, double X[N+2], contains in X[1] through X[N] the
#    probabilities of outcomes 1 through N.
#
#    Output, double Y[N+2], the Walker threshold vector.
#
#    Output, unsigned int A[N+2], the Walker index vector.
#
  import numpy as np
#
#  Initialize A.
#
  a = np.zeros ( n + 2, dtype = np.int32 )
  for i in range ( 0, n + 2 ):
    a[i] = i
#
#  Initialize B to the "stay here" value, and set sentinel values at the ends.
#
  b = np.zeros ( n + 2, dtype = np.int32 )

  for i in range ( 0, n + 2 ):
    b[i] = a[i]
#
#  Copy X into Y.
#  Scale the probability vector and set sentinel values at the ends.
#
  y = np.zeros ( n + 2, dtype = np.float64 )

  y[0] = 0.0
  for i in range ( 1, n + 1 ):
    y[i] = x[i] * float ( n )
  y[n+1] = 2.0

  i = 0
  j = n + 1

  while ( True ):
#
#  Find i so Y[B[i]] needs more.
#
    while ( True ):
      i = i + 1
      if ( 1.0 <= y[b[i]] ):
        break
#
#  Find j so Y[B[j]] wants less.
#
    while ( True ):
      j = j - 1
      if ( y[b[j]] < 1.0 ):
        break

    if ( j <= i ):
      break
#
#  Swap B[i] and B[j].
#
    k    = b[i]
    b[i] = b[j]
    b[j] = k

  i = j
  j = j + 1

  while ( 0 < i ):
#
#  Find J such that Y[B[j]] needs more.
#
    while ( y[b[j]] <= 1.0 ):
      j = j + 1
#
#  Meanwhile, Y[B[i]] wants less.
#
    if ( n < j ):
      break
#
#  B[i] will donate to B[j] to fix up.
#
    y[b[j]] = y[b[j]] - ( 1.0 - y[b[i]] )     
    a[b[i]] = b[j]             
# 
#  Y[B[j]] now wants less so readjust ordering.
#
    if ( y[b[j]] < 1.0 ):
      k    = b[i]
      b[i] = b[j]
      b[j] = k
      j = j + 1
    else:
      i = i - 1

  return y, a

def walker_build_test ( ):

#*****************************************************************************80
#
## WALKER_BUILD_TEST tests WALKER_BUILD.
#
#  Licensing:
#
#    This code is distributed under the GNU LGPL license.
#
#  Modified:
#
#    20 February 2016
#
#  Author:
#
#    John Burkardt
#
  import numpy as np
  import platform
  from i4_choose import i4_choose
  from r8vec_print import r8vec_print

  print ( '' )
  print ( 'WALKER_BUILD_TEST' )
  print ( '  Python version: %s' % ( platform.python_version ( ) ) )
  print ( '  WALKER_BUILD builds the Walker sampler data vectors Y and A,' )
  print ( '  given a probability vector X.' )

  n = 5
  x = np.zeros ( n + 2, dtype = np.float64 )
  for i in range ( 1, n + 1 ):
    x[i] = float ( i4_choose ( n - 1, i - 1 ) ) / float ( 2 ** ( n - 1 ) )
  r8vec_print ( n + 2, x, '  Binomial PDF (ignore first and last entries):' )

  y, a = walker_build ( n, x )
  print ( '' )
  print ( '   I    A[I]    Y[i] (ignore first and last entries)' )
  print ( '' )
  for i in range ( 0, n + 2 ):
    print ( '  %2d  %2d  %10.4g' % ( i, a[i], y[i] ) )
#
#  Terminate.
#
  print ( '' )
  print ( 'WALKER_BUILD_TEST' )
  print ( '  Normal end of execution.' )
  return

if ( __name__ == '__main__' ):
  from timestamp import timestamp
  timestamp ( )
  walker_build_test ( )
  timestamp ( )