%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Neoclassical growth model using Bellman Equation iteration   %
% plots non-linear and linearized policy for                   %
% CRRA or CARA and plots time series                           %
%                                                              %
% by Ivan Werning, 2003                                        %
% note: uses f.m and u.m functions                             %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
tic
% housekeeping
clear
global alpha delta beta sigma CARA A

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% parameter values
% suggestion: play with the parameters (especially sigma, alpha and delta)
% and note how it affects the speed of convergence to the steady state
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
alpha = .36;
delta = .08;
beta  = .96;
rho = beta^-1 -1;
A = 1;
sigma =  1; % if you want to use the CARA specification set sigma='exp'
CARA  =  1;

% set capital state space interval as fractions of steady state
klow    = .5; khigh = 1.5;

% grid size and vector of ones
gridn   = 1000;
e       = ones(gridn,1);

% compute steady state values
kss     = ((1/beta - 1 + delta)/alpha)^(-1/(1 - alpha));
css     = f(kss) - kss; % 
rss      = beta^(-1)-1;
kgr     = ((delta)/alpha)^(-1/(1 - alpha)); % compute golden rule capital

% create grid for capital
klow    = kss*klow; khigh   = kss*khigh;
k = (klow: (khigh-klow)/(gridn-2) : khigh)';
k = sort([k ; kss]);        % add ss capital level

% consumption matrix (today's k in column, tomorrow's in rows)
C = e*f(k') - k*e' ; 

% make negative consumption unattractive
if sigma ~= 'exp'
    C = max(C , 0);         % take out negative consumption (or else u(c) would not be defined)
    U = u(C) + ~C*(-1e50);  % add penalty for zero consumption
else
    U = u(C);
end

% initialize iteration
crit    = 1; iter = 1; 
v  = u(f(k) - k)/(1-beta);      % initial guess (value of keeping k and c constant)
%v =0*v + 5;                    % alternative initial guess

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% iterate on Bellman equation  %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
while crit > 1e-20
    
    vnew = max( U + beta*v(:,iter)*e' )';

    % compute normalized criterion (scale is in consumption per period)
    if sigma ~= 'exp'
        critnew = (1-beta)*max(abs(vnew - v(:,iter)))/css^(-sigma+1);
    else
        critnew = (1-beta)*max(abs(vnew - v(:,iter)))/css/exp(-CARA*css);
    end
    
    % display iteration information
    % note: contraction mapping implies that critnew/crit < beta
    display([iter , critnew, critnew/crit ])
    crit = critnew; iter = iter +1; v(: , iter)  = vnew;
    
end

% compute policy function
[vnew, policykindex]= max( U + beta*v(:,iter)*e');
policyk = k(policykindex);
figure(1); plot(k,[k policyk]); grid

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% plot time series of capital evolution starting
% at k0 half the value of the steady state
ktime(1) = min(find(k > kss/2)); T = 50;
for t=1:T-1 ;  ktime(t+1) = policykindex(ktime(t)); end
ktime = k(ktime); figure(2);  plot( 1:T, ktime/kss,'b.'); grid
title('evolution of capital (as fraction of steady state')

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% solves linearized policy rule (based on linearized EE)
ddf = alpha*(alpha-1)*kss^(alpha-2);df  = 1/beta; ddUdU = -sigma/css;
if sigma == 'exp' ;  ddUdU = -CARA; end
a = 1; b = -[ 1 + 1/beta + ddf/df/ddUdU ]; c = 1/beta;

% compute roots
lamda1 = [- b + sqrt(b^2 - 4*a*c)]/a/2;
lamda2 = [- b - sqrt(b^2 - 4*a*c)]/a/2;
lamda  = [lamda1, lamda2]; 
[bla index] = min(abs(lamda));  % check root with lowest ABSOLUTE value
lamda = lamda(index);           % take that root

% compute linear policy rule
policyklinear = (k-kss)*lamda + kss;

% plot linear policy rule on top of actual policy function
figure(1); hold on; plot(k,policyklinear,'r'); grid on; hold off

break

% plot time series for linear policy on top of previous time series
ktimelinear(1)=.5*kss; for t=1:T-1 ;  ktimelinear(t+1) = lamda*(ktimelinear(t)-kss)+kss; end
figure(2);hold on; plot( 1:T, ktimelinear/kss,'r.'); hold off
toc

Vaut = u(f(k) - k)/(1-beta);
rr =A*alpha* k.^(alpha-1) - delta;
if sigma ==1 
Vceq    = exp(v(:,end)*(1-beta));
Vautceq = exp(Vaut*(1-beta));
gain_perc = (Vceq-Vautceq)./Vautceq;
else
    Vceq    = ((1-sigma)*v(:,end)*(1-beta)).^(1/(1-sigma));
    Vautceq = ((1-sigma)*Vaut*(1-beta)).^(1/(1-sigma));
    gain_perc = (Vceq-Vautceq)./Vautceq;
    gains_perc = [ rr , gain_perc];
end

upper = ((k-kss)/beta + f(kss) - f(k));
cc = f(k) - k;
gain_upper = upper./cc;

% figure(3); hold on
% plot(rr-rss,gain_perc,'.r',rr-rss,gain_upper,'.k')
% axis([-.025, 0.01, 0, 0.02])
% grid on
% xlabel('difference in interest rate')
% ylabel('welfare gain')
% 
% figure(4); hold on
% plot(rr-rss,gain_perc,'.r')
% axis([-.025, 0.01, 0, 0.02])
% grid on
% xlabel('difference in interest rate')
% ylabel('welfare gain')
%