clear all
bet=.96;alf=.36;thet=2;del=.08; NUM_ITER = 1000;
kss=(1/(bet*alf)-(1-del)/alf)^(1/(alf-1)); % calculating the kss
kd  = .01*kss;
ku  = 1.5*kss;
k = [kd:.005:ku]';
n=size(k,1); 
k1=k'; 

%## consumption ##
PPF = kron(ones(1,n),(1-del)*k+k.^alf);
c = PPF-kron(ones(n,1),k1);
c_policy = zeros(length(k),NUM_ITER);
v_iter   = zeros(length(k),NUM_ITER);
v_iter(:,1) = zeros(n,1);

%## first RHS of the Bellman equ ##
% Values of k_{t} are in columns (down), k_{(t+1) in rows (right)}
c = max(c,1E-8);
Ut = (c.^(1-thet)-ones(n))/(1-thet);
RHS = Ut+bet*kron(ones(n,1),v_iter(:,1)');

% Numerical value function iterations
for i=2:NUM_ITER
    [v_iter(:,i),ind_k_tilda] = max(RHS,[],2); %gives the maximum in each row of RHS, 
                        % saves values in v1 and index of k_tilda in ind_k_tilda
                        % capital next period, they are column vectors
    k_tilda  = k(ind_k_tilda);
    c_policy(:,i) = k.^alf +(1-del)*k - k_tilda;
    % checking the norm of the difference between successive iterations:
    % this is convergence criterion
    err = norm(v_iter(:,i) - v_iter(:,i-1));
    if err < 10^-6
        break;
    end
    RHS = Ut + bet*kron(ones(n,1),v_iter(:,i)');
end

close all
figure(1)
subplot(121);
plot(k,k,k,k_tilda,'--');
axis([0 8 0.2 8.0])
title('Optimal policy upon convergence: k_ tilda vs k')

subplot(122);
plot(k,c_policy(:,i))
title('Optimal policy upon convergence: c vs k')

figure(2)
plot(k,v_iter(:,1),'*',k,v_iter(:,2),'r-',...
    k,v_iter(:,3),'g-',k,v_iter(:,i),'k-')
title('Value function iterations: value function')
legend('1N','2N','3N','Last N')

figure(3)
plot(k,c_policy(:,2),'*',k,c_policy(:,3),'r-', ...
    k,c_policy(:,4),'g-',k,c_policy(:,5),'m-', ...
    k,c_policy(:,i),'k-')
title('Value function iterations: consumption policy function')
legend('1N','2N','3N','4N','Last N')
