%% Life-cycle Model
% In this section, we'll numerically solve the life-cycle model through backwards induction.
% We begin by assigning the parameters.
% Generally, we assume that real economic activity starts around real-world age 21.  Accordingly, age 1 in the model will correspond to real age 21.  We will assume that maximum age is 80 years old.
% To parameterize the model, we will use real data estimated from Texas residents surveyed in the 2012 Current Polulation Survey.  This vector 'productivity' represents estimated log hourly wage for each age.

warning('off')
optim_options = optimset('Display', 'off', 'TolFun', 1e-5, 'TolX', 1e-5, 'MaxFunEvals', 10e4);

% Assigning parameters
discount_factor = 0.96 ;
crra            = 2    ;
wage            = 1    ;
interest_rate   = 1.04 ;
initial_wealth  = 0    ;
max_age         = 60   ;

% Reading in productivty profile from file
productivity    = xlsread('productivity_profile.xls'); % Estimated 2012 Texas log hourly earnings ages 21-80

% Storing values in a structure
params.discount_factor = discount_factor ;
params.crra            = crra            ;
params.wage            = wage            ;
params.interest_rate   = interest_rate   ;
params.initial_wealth  = initial_wealth  ;
params.max_age         = max_age         ;
params.productivity    = productivity    ;

% Regarding the period T problem, we will simply assign V_{T+1} = 0 and let the computer solve for optimal savings.  We expect this value to be a^*_{T+1} = 0, and we can check this later on.
% As we mentioned in the introduction to dynamic optimization, we need to know the continuation value correpsonding to the value function from period t+1 before we can solve the period t problem.
% That said, we can't simply create a function for each continuation value!  This is possible for the three-period problem, as we showed in the respective code.  However, the complexity of this problem explodes for more than three periods.
% Instead, we can simplify the problem by solving each period t's value function (which will become period t-1's continuation value) at some predetermined set of points, and save those for future use.
% This process is known as "discretizing the state space."
% Then, when we solve the period t value function, we will have to interpolate over the continuation value.
% Let's continue by discretizing the state space.

% Discretizing the state space
asset_upper_bound = 200  ;    % Take this as given now...we'll have more to say about it later.
asset_lower_bound =  -2  ;    % A negative value here means we can borrow.
n_assets          = 100  ;    % Be careful here.  A large value will take a long time to run.
assets            = linspace( asset_lower_bound, asset_upper_bound, n_assets);

% Loading 'assets' in the 'params' structure
params.assets = assets ;

% Ok, that wasn't so bad - we've actually created several discretized grids before.  This time, we're just determining at which points we want to solve the value function.
% Now we can continue by preallocating the arrays where we will store the value functions and optimal choices:

% Preallocating arrays
value_function = zeros( n_assets, max_age + 1 ) ;
savings        = zeros( n_assets, max_age     ) ;

% Ok, that wasn't too bad either - as long as you knew the right size to make the arrays.  Notice that the first dimension has size equal to the number of gridpoints where we want to solve the value function.  In theory, we do not have to store the value function for each age, but at this point, it's easy for us to do it.  Plus, we may need it later in our training.
% Notice that we gave 'value_function' size T+1 in the second dimension.  This was an easy way for us to go ahead and assign V_{T+1}(a_{T+1}) = 0 for every possible value of a_{T+1}.  Because the continuation value at T+1 does not increase with assets, we can begin to build confidence that optimal savings will exceed zero in the last period.
% We need to backwards-solve the problem starting from period T all the way to period 1.  We can do this by instructing the for loop to iterate backwards.
% Within the loop, we will call 'calculate_utility' which will calculate utility for given state variables, choice variables, and parameters.

% Solving the lifecycle model
for age = max_age:-1:1    % Instructing the for loop to iterate by increments of -1
    fprintf( 'Solving age %d.\n', age )
    for ia = 1:n_assets
        continuation_values               = value_function( :, age+1 ) ;  % Continuation values over all possible savings
        asset                             = assets(ia)                 ;  % Solving at gridpoint for given asset
        initial_guess                     = eq(age,max_age)*0 + (1 - eq(age,max_age))*assets(ia);    % Making initial guess
        [optimal_saving, optimal_utility] = fminsearch(@(saving) calculate_utility(...
                                                       saving, asset , age, continuation_values,...
                                                       params ), initial_guess ,optim_options       );
        value_function( ia, age ) = -1*optimal_utility;
        savings       ( ia, age ) = optimal_saving    ;
    end
    
end

%
% At this point, we have solved for the policy function a^*_{t+1}(a_t) as it relates to the state variable.
% How can we use the policy function to simulate life-cycle behavior?
% Start by assuming that initial wealth is zero, i.e., a_1 = 0.
% The vector determined by evaluating 'savings' at age = 1 gives a grid of optimal age 1 savings for different initial wealths.  Accordingly, to get optimal savings when a_1 = 0, all we have to do is interpolate:

initial_wealth  = 0                                                   ;
age             = 1                                                   ;
optimal_saving = interp1( assets, savings( :, age ), initial_wealth ) ;
fprintf('Optimal savings in period 1 when initial wealth is zero is %0.4d.\n', optimal_saving)

%
% Ok, we have successfully used our policy function to determine how much the individual would save in period 1 if initial wealth were zero.
% We could do the same thing for every age.  In fact, we could determine the entire life-cycle savings by iterating over this process:

savings_profile    = zeros(1,max_age) ;    % Preallocating the savings profile
savings_profile(1) = -1.5   ;    % Assigning initial wealth to the first entry
% Solving for optimal savings over the life-cycle, assuming that initial wealth is zero
for age = 1:max_age-1
    savings_profile(age+1) = interp1( assets, savings( :, age ), savings_profile(age), 'spline' );
end
ages = 1:1:max_age;    % Creating an x-axis for the graph
plot( ages, savings_profile )

%

function utility = calculate_utility( saving, asset, age, continuation_values, params )
    
    % Testing constraints
    if age<params.max_age
        minimum_saving = params.assets(1) ;    % Regular debt limit is lower bound of asset grid.
    elseif age==params.max_age
        minimum_saving = 0                ;    % Implementing no-Ponzi condition
    end
    
    % Treating saving constraint violation
    if saving < minimum_saving || saving > params.assets(end)
        utility = inf;
        return
    end
    
    % Solving for consumption
    consumption = params.interest_rate * asset  - saving + params.wage * params.productivity(age);
    
    % Treating consumption constraint violation
    if consumption < 0
        utility = inf;
        return
    end
    
    % Solving for utility at the saving value
    continuation_value = interp1(params.assets, continuation_values, saving, 'linear' );
    
    utility = (1 /( 1-params.crra )) * consumption^(1 - params.crra) + params.discount_factor * continuation_value;
    utility = -1*utility ;
    
    
end

%
%