% solve HH infinite horizon saving problem on finite grid, using acceleration steps
% and linear interpolation of the value function
% labor income follows Markov process
% Michael Reiter, IHS, February 2017
% last update: Feburary 2020
function Sol = consinterp(Vcont,Kmax,nAccel)

  [nK,nZ] = size(Vcont);
  Kmin = 0;
  gridK = logspaceshift(Kmin,Kmax,nK,Kmax/20); % finite grid of capital
  %gridK = linspace(Kmin,Kmax,nK)'; % finite grid of capital

  r = 0.02;
  beta = 0.975;
  m = 2.5;   %width of grid in terms of standard deviations
  rho = 0.9;
  sigma = 0.2;
  [TransZ,gridlogZ,probst]=markovappr(rho,sigma,m,nZ);
  gridZ = exp(gridlogZ);
  TransZTransp = TransZ';

  Vend = zeros(nK,nZ);
  Vtry = NaN(nK,1);
  OptCons = zeros(nK,nZ);
  IndxOptSaving = zeros(nK,nZ);
  indxChoice = zeros(nK,nZ);
  pHiOptSaving = zeros(nK,nZ);
  for iter = 0:10000  % maximum number of steps = 10000
    Vend = beta*Vcont*TransZTransp;
    if(mod(iter,nAccel)==0)
      for j=1:nZ % solve problem on each exogenous grid point 
        for i=1:nK % solve problem on each endogenous grid point 
          Kbeg = gridK(i);
          % try all possible levels of saving at once:
          [cMin,cMax] = cbounds(Kbeg,gridZ(j),gridZ,r,Kmin,Kmax);
          OptCons(i,j) = golden(@obj2max,cMin,cMax,Kbeg,j,gridK,gridZ,Vend,r);
          [V(i,j),IndxOptSaving(i,j), pHiOptSaving(i,j)] = ...
            obj2max(OptCons(i,j),Kbeg,j,gridK,gridZ,Vend,r);
        end
      end
      dist = max(max(abs(V-Vcont)));
      fprintf(1,'iter = %d; dist = %e\n',iter,dist);
      if(dist<1e-8)
        break;
      end
    else
      for j=1:nZ % solve problem on each grid point (state)
        for i=1:nK % solve problem on each grid point (state)
          Kbeg = gridK(i);
          % use best policy of earlier iteration:
          iOpt = IndxOptSaving(i,j);
          pHi = pHiOptSaving(i,j);
          Kend = gridK(iOpt);
          c = OptCons(i,j);
          V(i,j) = util(c) + (1-pHi).*Vend(iOpt,j) + pHi.*Vend(iOpt+1,j);
        end
      end
    end
    Vcont = V;
  end
  Sol = struct( 'V',V,'OptCons',OptCons,'IndxOpt',IndxOptSaving,'pHiOpt',pHiOptSaving,'indxChoice',indxChoice,...
    'gridEndog',gridK,'gridExog',gridZ,'TransExog',TransZ);
end

% objective function
function [val,indx,pHi] = obj2max(c,Kbeg,j,gridK,gridZ,Vend,r)
  y = gridZ(j);
  Kend = (1+r)*Kbeg + y - c; 
  [indx,pHi] = lookupbisect(gridK,Kend); %pHi is defined as (Kend-gridK(indx))/(gridK(indx+1)-gridK(indx));
  % linear interpolation of Vend at Kend:
  val = util(c) + (1-pHi).*Vend(indx,j) + pHi.*Vend(indx+1,j);
end

function u = util(c)
  u = log(c);
end
function u = utilCheck(c)
  u = -Inf(size(c));
  indx = c>0;
  u(indx) = util(c(indx));
end
function [cMin,cMax] = cbounds(Kbeg,y,gridY,r,Kmin,Kmax)
  cMin = max(1e-100,(1+r)*Kbeg + y - Kmax); %cannot consume less
  cMax = (1+r)*Kbeg + y - Kmin;
end

