% solve stochastic HH lifecycle problem on a FINITE GRID of capital
% Problem: 
%    max sum_{t=1}^T beta^t u(c_t)
%      subject to
%    c = (1+r)*Kbeg + y(t) - Kend;
%    Kend >= 0
%    y(t) follows Markov process
% Michael Reiter, IHS, February 2020
% Input arguments:
%   Kmax: maximal capital level
%   nK:   size of capital grid
%   nZ:   size of productivity grid
function [V,indxOptSaving,gridK] = lcstoch(Kmax,nK,nZ)
  % STOCHASTIC and VECTORIZED version of lc.m
  % parameters:
  beta = 0.98;
  r = 0.015;
  T = 60;
  rho = 0.8;
  std = 0.1;
  [Pi,gridlogZ] = markovappr(rho,std,3,nZ);
  gridZ = exp(gridlogZ);

  Kmin = 0;  % no borrowing allowed
  % grid of possible capital states:
  gridK = linspace(Kmin,Kmax,nK)';

  V = NaN(nK,nZ,T+1);  % value function
  V(:,:,T+1) = 0;   % no value of capital after death (no bequest motive)
  indxOptSaving = NaN(nK,T);
  for t=T:-1:1  % iterate backwards in time
    fprintf(1,'now solving period %d\n',t);
    % define VcontExpected(i,j) as expected continuation value, conditional on current K=gridK(i) and current  iZ=j:
    % VcontExpected(i,j) = Exp_t V(i,jNext,t+1) = sum_k V(i,k,t+1)*Pi(j,k) = V(i,:,t+1)*Pi(j,:)'
    %  therefore:
    VcontExpected = V(:,:,t+1) * Pi';  
    for i=1:nK   % optimize at each grid point of capital
      Kbeg = gridK(i);
      for iz=1:nZ  % optimize at each grid point of productivity
        z = gridZ(iz);
        Kend = gridK;% try all levels of saving
        if t<=40
          y = z;  % labor income equals productivity
        else
          y = 0;  % retired
        end
        c = (1+r)*Kbeg + y - Kend; % all consumption levels
        vTry = util(c) + beta*VcontExpected(:,iz); % current objective function depends on expected continuation value
        [V(i,iz,t),indxOptSaving(i,iz,t)] = max(vTry);
        % the slow version is:
        % vTry(j) = util(c) + beta*dot(Pi(iz,:),V(j,:,t+1));   % inefficient, recomputes expected continuation many times
      end
    end
  end

end
function u = util(c);
  u = -Inf(size(c));
  iPos = c>0;
  u(iPos) = log(c(iPos));
end

