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.
% gamma : unconditional expectation of the technology parameter

alpha = .35;
beta  = .98;
delta = .025;
sigma = 2;
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.95*kstar,1.05*kstar,gk);


% Compute a (gk x gk) dimensional consumption matrix c
%   for all the (gk x gk) values of k_t, k_t+1 
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
   end
end

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

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

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

enere = ones(gk,1); 

while diff > convcrit
    % for each k_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;
end

disp(iter)

% Find what gridpoint of the k_t+1 vector which gave the optimal value

[Tv,gridpoint] =  max((u + beta*enere*v'),[],2);

kgridrule = gridpoint;


% Find what element of k_t+1 which gave the optimal value
kdecrule = k(kgridrule);      

figure
plot(k,kdecrule,k,k);
xlabel('k_t')
ylabel('k_{t+1}')
%print -djpeg kdecrule.jpg


% Compute the optimal decision rule for c as a function of the state
% variable

cdecrule = gamma*k.^alpha + (1-delta)*k - kdecrule;


figure
plot(k,cdecrule)
xlabel('k_t')
ylabel('c_t')
%print -djpeg cdecrule.jpg