% Solving a stochastic neoclassical growth model with elastic labor supply
% ECON 4310, University of Oslo
% by Espen Henriksen

clear all; % clears all variables from the memory
close all; % closes all figures

tic % starts the "stop watch"

alpha   = .3;
beta    = .95;
delta   = .05;
psi     = 4.12;
sigma   = 1;

 
hg   = 51;   % Number of grid points for the control variable h
kg   = 101;  % Number of grid points for both the control var and state var k

khratio = ((1/beta - 1 + delta)/alpha)^(1/(alpha-1));
chratio = khratio^alpha - delta*khratio;
hstar   = ((1-alpha)*khratio^alpha)/(psi*chratio);
kstar   = khratio*hstar;
cstar   = chratio*hstar;

clear khratio chratio;      % if you won't need 'em delete 'em

kgrid = linspace(.85*kstar,1.15*kstar,kg)';
hgrid = linspace(.8*hstar,1.2*hstar,hg)';

c = zeros(hg,kg,kg);
u = zeros(hg,kg,kg);

% i is a counter for the control variable h
% j is a counter for the control variable k'
% n is a counter for the endogenous state variable k

warning off; % disable warning for taking log of zero (just annoying)

for i = 1 : hg
    for j = 1 : kg
        for n = 1 : kg
                c(i,j,n) = (kgrid(n)^alpha)*hgrid(i)^(1-alpha) + (1-delta)*kgrid(n) - kgrid(j);

                    if c(i,j,n) < 0
                          c(i,j,n) = 0;
                    end
                 u(i,j,n) = log(c(i,j,n)) - psi*hgrid(i);
        end
    end
end


% warning on;      % turn warnings on again
% clear c          % free up memory
% clear i j n      % clean up

v        = zeros(kg,1);
convcrit = 1E-11;   % chosen convergence criterion
diff = 1;           % arbitrary initial value greate
iter = 0;           % iterations counter

while diff > convcrit
        diff = 0;
                for n = 1 : kg
                    objfn(:,:,n) = u(:,:,n) + beta*ones(hg,1)*v';
                Tv(n,1) = max(max(objfn(:,:,n)));
    
                end
        diff = norm(v-Tv);
        v = Tv;
        iter = iter + 1;
end



for n = 1 : kg
        objfn(:,:,n) = u(:,:,n) + beta*ones(hg,1)*v';
        [tmp1,x1] = max(objfn(:,:,n),[],1);     % vyber max h (pro vsechna k')
        [tmp2,x2] = max(tmp1,[],2);             % vyber max k'(pro opt. h)
        kgridrule(n) = x2;
        hgridrule(n) = x1(x2);
        kdecrule(n) = kgrid(kgridrule(n));
        hdecrule(n) = hgrid(hgridrule(n));
        cdecrule(n) = (kgrid(n)^alpha)*hdecrule(n)^(1-alpha) + (1-delta)*kgrid(n) - kdecrule(n);
end


 

%If you won't need 'em -- delete 'em
% clear tmp1 tmp2 x1 x2
% clear diff convcrit iter m n
% clear objfn Tv
% clear u

figure;
plot(kgrid,v);
title('Value function')
xlabel('k_t')

figure;
plot(kgrid,kgrid,kgrid,kdecrule);
title('Decision rules for capital');
xlabel('k_t')
ylabel('k_{t+1}')

figure;
plot(kgrid,cdecrule)
title('Decision rules for consumption');
xlabel('k_t')
ylabel('c_t')

figure;
plot(kgrid,hdecrule)
title('Decision rules for labor supply');
xlabel('k_t')
ylabel('h_t')

