% stochasticky neoklasicky rustovy model s nabidkou prace

clear all; 
close all; 

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

 
hg   = 6;   % pocet bodu (gridu) pro ridici promenou h
kg   = 11;  % pocet bodu (gridu) pro ridici a stavovou promenou k

% steady state
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;      % kdyz je nepotrebujete, vymazte

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 je citac pro ridici promenou h
% j je citac pro ridici promenou k'
% m je citac pro stavovou promenou k

warning off; % vypnuti varovani log nuly (akorat otravuje)

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;      % zapnuti varovani
% clear c          % vycisteni pameti
% clear i j m n     

v        = zeros(kg,1);

convcrit = 1E-11;  % konvergencni kriterium (male cislo)
diff = 1;         % arbitrarne zvolena cislo (vetsi jak convcrit)
iter = 0;         % citac iteraci

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

% nalezeni rozhodovacich pravidel ze zkonvergovane Bellmanovy rovnice

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, pro vsechna k)
        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

% vykresleni

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')

