clear all;
close all;

% Define the structural parameters of the model,
% i.e. policy invariant preference and technology parameters
% alpha : capital's share of output
% beta  : time discount factor
% delta : depreciation rate
% sigma : risk-aversion parameter, also intertemp. subst. param.

alpha = .35;
beta  = .98;
delta = 1;
sigma = 1;
gamma = 5; 


% Find the steady-state level of capital as a function of 
%   the structural parameters

kstar = ((1/beta - 1 + delta)/(alpha*gamma))^(1/(alpha-1));

% Define the number of discrete values k can take

gk = 101;
k  = linspace(0.90*kstar,1.10*kstar,gk);


% Compute a (gk x gk x gz) dimensional consumption matrix c
%   for all the (gk x gk x gz) values of k_t, k_t+1 and z_t.
for h = 1 : gk
   for i = 1 : gk
          c(h,i) = gamma*k(h)^alpha + (1-delta)*k(h) - k(i);
            if c(h,i) < 0
               c(h,i) = 0;
            end
            % h is the counter for the endogenous state variable k_t
            % i is the counter for the control variable k_t+1
            % j is the counter for the exogenous state variable z_t
   end
end

% Compute a (gk x gk x gz) dimensional consumption matrix u
%   for all the (gk x gk x gz) values of k_t, k_t+1 and z_t.

if sigma == 1
      u = log(c);
         else
      u = (c.^(1-sigma) - 1)/(1-sigma);    
end


% Define the initial matrix v as a matrix of zeros (could be anything)

v = zeros(gk,1);
enere = ones(gk,1); 

% Set parameters for the loop
convcrit = 1E-6;  % chosen convergence criterion
diff = 100;         % arbitrary initial value greater than convcrit
iter = 0;         % iterations counter




while diff > convcrit
    % for each combination of k_t and gamma_t
    %   find the k_t+1 that maximizes the sum of instantenous utility and
    %   discounted continuation utility
       Tv = max(u + beta*enere*v',[],2);
    iter = iter + 1;
    diff = norm(Tv - v);
    v = Tv;
    
    if iter == 300 
        v1 = v;
    elseif iter == 350
        v2 = v;
    elseif iter == 400
        v3 = v;
    elseif iter == 500
        v4 = v;
    end
end
disp(iter)

%%
b = alpha/(1-beta*alpha);
a = 1/(1-beta)*((1+beta*b)*log(gamma/(1+beta*b)) + beta*b*log(beta*b));

trueval = a + b*log(k);

hold on
plot(k,trueval,'r',k,v1,'b',k,v2,'b',k,v3,'b',k,v4,'b')
legend('true vf','approximated vf',4)
xlabel('k')
title('Value functions')


%%

        % Using the [ ] syntax for max does not only give the value, but
        %   also the element chosen.
        [Tv,gridpoint] = max(u + beta*enere*v1',[],2);
        % Find what grid point of the k vector which is the optimal decision
        kgridrule = gridpoint;
        % Find what value for k_t+1 which is the optimal decision
        kdecrule = k(gridpoint);

truekdecrule = beta*alpha*gamma*k.^alpha;

figure
plot(k,truekdecrule,'r',k,kdecrule,'b',k,k,'k')
title('Decision rule for capital')
legend('true dec. rule','approximated dec. rule',4)
xlabel('k_t')
ylabel('k_t+1')