%% Two-Period Problem
% First, we will solve the simple two-period problem, then we will solve the three-period problem.
% To do this, we will keep in mind that we need the indirect utility funciton from the two-period problem to solve the period 0 problem through backwards induction.
% Accordingly, we will separate initial wealth from the other parameters, since we intend to convert it into a variable.

% Assigning parameters

parameters.discount_factor = 0.99 ;
parameters.interest_rate   = 1.04 ;
initial_wealth             = 10   ;


% Optimizing utility
initial_guess                      = initial_wealth / 2                                                                            ;
[optimal_savings, optimal_utility] = fminsearch(@(savings) calculate_utility(savings,...
                                                           initial_wealth, parameters), initial_guess ) ;
optimal_utility                    = -1*optimal_utility                                                                            ;

fprintf('Optimal savings is %0.4f \n', optimal_savings )
fprintf('Optimal utility is %0.4f \n', optimal_utility )


% Solving indirect utility over initial wealth grid
n_wealths          = 100                                                           ;
wealth_lower_bound = 0.001                                                         ;
wealth_upper_bound = 100                                                           ;
initial_wealths    = linspace( wealth_lower_bound, wealth_upper_bound, n_wealths ) ;
indirect_utilities = zeros( n_wealths, 1 )                                         ;

for iw = 1:n_wealths
    
    indirect_utilities( iw ) = calculate_indirect_utility( initial_wealths(iw), parameters );
    
end

plot( initial_wealths, indirect_utilities, 'LineWidth', 4 )


% Packaging continuation value for future use
continuation_values.initial_wealths    = initial_wealths    ;
continuation_values.indirect_utilities = indirect_utilities ;

%% Three-period problem
% Now we want to solve the three-period problem by using the continuation value from the two period problem into the three-period problem.
% We could just call 'calculate_indirect_utility' directly from the three-period problem.  However, for the purpose of learning useful things, we will just assume that we know it at the points 'indirect_utilities' corresponding to points 'initial_wealths.'  In other words, we will interpolate over these values to approximate the continuation value.


% Solving for optimal lifetime utility
initial_guess = 3 ;
[optimal_savings, optimal_lifetime_utility] = fminsearch(@(savings) calculate_lifetime_utility( savings,...
                                                                    continuation_values, initial_wealth, parameters), initial_guess );
optimal_lifetime_utility = -1*optimal_lifetime_utility;                                                            

fprintf('Optimal savings is %0.4f.\n'         , optimal_savings          )
fprintf('Optimal lifetime utility is %0.4f.\n', optimal_lifetime_utility )

%

function utility = calculate_utility(savings, initial_wealth, parameters)
    
    % Preallocating consumption vector
    consumption = zeros(2,1) ;
    
    % Solving consumption in each period
    consumption(1) = initial_wealth - savings           ;
    consumption(2) = parameters.interest_rate * savings ;
    
    % Calculating utility
    utility = log( consumption(1) ) + parameters.discount_factor * log( consumption(2) ) ;
    utility = -1*utility;
    
end



% Solving indirect utility

function indirect_utility = calculate_indirect_utility(initial_wealth, parameters)
    
    % Optimizing utility
    initial_guess                      = initial_wealth / 2                                                 ;
    [optimal_savings, optimal_utility] = fminsearch(@(savings) calculate_utility(savings,...
                                                               initial_wealth, parameters), initial_guess ) ;
    indirect_utility                   = -1*optimal_utility                                                 ;
    
end


% Calculating lifetime utility 

function lifetime_utility = calculate_lifetime_utility( savings, continuation_values, initial_wealth, parameters)
    
    consumption = initial_wealth - savings ;
    
    continuation_value = interp1( continuation_values.initial_wealths, continuation_values.indirect_utilities, savings, 'linear' );
    
    lifetime_utility = log( consumption ) + parameters.discount_factor*continuation_value ;
    
    lifetime_utility = -1*lifetime_utility ;
    
    
end