'''
In this solution we do not save intermediate results, hence we need special arrangement
to backtrack the solution
'''
b = 0.8 # discount factor
import math
def f(n, S, cStar = [0]):
    if n == 1: # Consume whatever is left
      cStar[0] = S
      return math.log(S)
    else:
      best = -1
      for c in range(1,S - (n - 1) + 1 ): # Save minimum n-1 to subsequent stages
        test = math.log(c) + b * f(n - 1, S - c)
        if test > best:
          best = test
          cStar[0] = c
      return best

# Main program:
initialCapital = 10
nPeriods = 5
b = 0.8 # Note <b> is a global variable in this module!
'''

In the beginning, we could just run our f() function to get the optimal
utility, i.e., no backtracking:

'''
print("Total utility: " , f(nPeriods, initialCapital))
print("\n*********************")
print("Backtracking:")

'''
But this is not enough, we need to backtrack the solution for each stage
in order to find the spending c each period

We start by solving the problem for stage N = nPeriods

Our recursive function also returns the optimal spending
in stage N as an optinal parameter needed to backtrack the best solutions

Then we know the sloution for stage N in terms of what to spend
in that stage

For stage N-1, we now call the function with state S = initialCapital
minus spending in stage N
We proceed in this manner to backtrack the spending in each stage
'''

S = initialCapital
n = nPeriods
c =[0] 
while n > 0:
  utility = f(n, S, c)
  print("Stage =" , n , "c =" , c[0] , " f =" , utility)
  S -= c[0]
  n -= 1