clear all
close all

alpha=.3;
beta=.6;

k = [.04; .08; .12; .16; .2];
v = [0; 0; 0; 0; 0];

y = k.^alpha;
c = y*ones(1,5) - ones(5,1)*k';
u = log(c);


% dva kroky mechanicky
% prvni vypocet hodnotove funkce
v1 = u + beta*ones(5,1)*v';

% maximalizace (ktera hodnota k_t+1 dava max hodnotovou funkci)
[v1_max,p1] = max(v1,[],2);
% nej k_t+1 je prvni prvek vektoru, k' = 0.04


% druhy krok, vypocet hodnotove funkce
v2 = u + beta*ones(5,1)*v1_max';
% nej k_t+1 je k' = 0.08 pro (k = 0.04, 0.08. 0.12. 0.16.) a k' = 0.012 pro (k= 0.2)
[v2_max,p2] = max(v2,[],2);


% nastaveni paremetru pro cyklus
convcrit = 1E-6;  % konv. kriterium
diff = 1;         % pocatecni hodnota (vetsi nez convcrit)
iter = 0;         % citac iteraci

while diff > convcrit
    %  najdi k_t+1 ktere maximalizuje RHS
    Tv = max(u + beta*ones(5,1)*v',[],2);
    iter = iter + 1;
    diff = abs(Tv - v);
    v = Tv;
   
    plot(k,v)
    xlabel('k')
    hold on
    iter = iter+1;
    
end


title('Iteration of value function')
iter


% decision rule
[Tv,i] = max((u + beta*ones(5,1)*v'),[],2);
 k_dr = k(i);  


%% analyticke reseni

A = 1/(1-beta)*((alpha*beta)/(1-alpha*beta)*log(alpha*beta) + log(1-alpha*beta));
B = alpha/(1-alpha*beta);

true_vf = A + B*log(k);
true_dr = alpha*beta*k.^alpha;

figure
plot(k,v,'r',k,true_vf)
legend('numerical approximation','true value function','Location','best')
title('Value function')
xlabel('k')

figure
plot(k,k_dr,'r.-',k,true_dr,'b.-')
title('Decision rule')
legend('numerical approximation','true decision rule','Location','best')
axis([0.04 0.2 0.06 0.14])
xlabel('k_t')
ylabel('k_{t+1}')