% Solving the standard RBC Model from Hansen 1985 by Linear-Quadratic
% Approximation -- social planer's problem
%
%Using the elimination method from Hansen and Prescot 1995


clear all
close all

% Parameters

theta = 0.36;
beta = 0.99;
gamma = 0.95;
delta = 0.025;
A = 2; 

% Steady state

lambda_bar = 1;
BB = (1-beta*(1-delta))/(beta*theta);
h_bar = ((1-theta)*BB)/(A*(BB-delta) + (1-theta)*BB);
K_bar = BB^(1/(theta-1))*h_bar;
I_bar = delta*K_bar;
Y_bar = lambda_bar*K_bar^theta*h_bar^(1-theta);
C_bar = Y_bar - I_bar;

nz = 1;     % number of exogenous state variables
ns = 1;     % number of endogenous state variables
nd = 2;     % number of control variables

% Computing the Jacobian and the Hessian
R_bar = log(C_bar) + A*log(1-h_bar);    % the return function evaluated at the steady state

Jz = Y_bar/C_bar;
Jk = (theta*Y_bar/K_bar)/C_bar;
Ji = -1/C_bar;
Jh = (1-theta)*(Y_bar/h_bar)/C_bar - A/(1-h_bar);

Hzz = -(Y_bar*Y_bar)/(C_bar*C_bar);
Hkk = theta*((theta-1)*C_bar - theta*Y_bar)*Y_bar/(K_bar*K_bar*C_bar*C_bar);
Hii = -1/(C_bar*C_bar);
Hhh = ((1-theta)*Y_bar*(-theta*C_bar - (1-theta)*Y_bar))/(C_bar*C_bar*h_bar*h_bar) - A/((1-h_bar)*(1-h_bar));

Hzk = ((C_bar - Y_bar)*theta*Y_bar/K_bar)/(C_bar*C_bar);
Hzi = Y_bar/(C_bar*C_bar);
Hzh = ((C_bar - Y_bar)*(1-theta)*Y_bar/h_bar)/(C_bar*C_bar);

Hki = (theta*Y_bar/K_bar)/(C_bar*C_bar);
Hkh = ((C_bar - Y_bar)*theta*(1-theta)*Y_bar/(K_bar*h_bar))/(C_bar*C_bar);

Hih = ((1-theta)*Y_bar/h_bar)/(C_bar*C_bar);

% Declaring the matrices
J_bar = [Jz;Jk;Ji;Jh];

H_bar = [Hzz Hzk Hzi Hzh
        Hzk Hkk Hki Hkh
        Hzi Hki Hii Hih
        Hzh Hkh Hih Hhh];

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%    
        
% theta = sym('theta');
% lambda = sym('lambda');
% h = sym('h');
% k = sym('k');
% i = sym('i');
% zero = sym('0');
% vars = [lambda,k,i,h];
% 
% obfun = log(lambda*k^(theta)*h^(1-theta) - i) + A*log(1-h);
% Jz_sym = (k^(theta)*h^(1-theta))/(lambda*k^(theta)*h^(1-theta) - i);
% Jk_sym = (lambda*theta*k^(theta-1)*h^(1-theta))/(lambda*k^(theta)*h^(1-theta) - i);
% Ji_sym = -1/(lambda*k^(theta)*h^(1-theta) - i);
% Jh_sym = (lambda*(1-theta)*k^theta*h^(-theta))/(lambda*k^(theta)*h^(1-theta) - i) - A/(1-h);
% 
% JJ = jacobian(obfun,vars)';
% JJ_bar = subs(JJ,{theta,lambda,k,i,h},[ttheta,lambda_bar,K_bar,I_bar,h_bar]);
% 
% Hz = jacobian(Jz_sym,vars);
% Hz_bar = subs(Hz,{theta,lambda,k,i,h},[ttheta,lambda_bar,K_bar,I_bar,h_bar]);
% 
% Hk = jacobian(Jk_sym,vars);
% Hk_bar = subs(Hk,{theta,lambda,k,i,h},[ttheta,lambda_bar,K_bar,I_bar,h_bar]);
% 
% Hi = jacobian(Ji_sym,vars);
% Hi_bar = subs(Hi,{theta,lambda,k,i,h},[ttheta,lambda_bar,K_bar,I_bar,h_bar]);
% 
% Hh = jacobian(Jh_sym,vars);
% Hh_bar = subs(Hh,{theta,lambda,k,i,h},[ttheta,lambda_bar,K_bar,I_bar,h_bar]);
% 
% HH = [Hz;Hk;Hi;Hh];
% HH_bar = [Hz_bar;Hk_bar;Hi_bar;Hh_bar];

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%        
        
W_bar = [lambda_bar;K_bar;I_bar;h_bar];
F = [1;lambda_bar;K_bar];

Q11 = R_bar - W_bar'*J_bar + 0.5*W_bar'*H_bar*W_bar;
Q12 = 0.5*(J_bar - H_bar*W_bar);
Q22 = 0.5*H_bar;

Q = [Q11 Q12'
     Q12 Q22];

% Imposing the laws of motion
B = [1 0 0 0 0 0 0 0                  % the last three columns are for tomorrow's states
     1-gamma gamma 0 0 0 0 0 0        % Don't forget 1-gamma here!
     0 0 1-delta 1 0 0 0 0];

% initial guess for P
TP = -0.01*eye(1+nz+ns);       % the iterated P
P = ones(1+nz+ns);
     
% Create the matrix R
RR = zeros(1+nz+ns+nd+1+nz+ns);
RR(1:1+nz+ns+nd,1:1+nz+ns+nd) = Q;
RR(1+nz+ns+nd+1:1+nz+ns+nd+1+nz+ns,1+nz+ns+nd+1:1+nz+ns+nd+1+nz+ns) = beta*P;

% iterate on Bellman's equation
norm_conv = 0.0000001;
step = 0;
step_max = 5000;

while norm(TP - P)/norm(P) > norm_conv &  step < step_max
    P = TP;
    R = RR;

R(1+nz+ns+nd+1:1+nz+ns+nd+1+nz+ns,1+nz+ns+nd+1:1+nz+ns+nd+1+nz+ns) = beta*P;

% Eliminating tomorrow's variables
for j = 1+nz+ns+nd+1+nz+ns:-1:1+nz+ns+nd+1
    S = [eye(j-1);B(j-(1+nz+ns+nd),1:j-1)];
    R = S'*R*S;
end

% Eliminating decision variables
coefs = zeros(nd,nz+ns+nd);
for j = 1+nz+ns+nd:-1:1+nz+ns+1
    coefs(j-(1+nz+ns),1:j-1) = -1/R(j,j)*R(j,1:j-1);
    S = [eye(j-1);coefs(j-(1+nz+ns),1:j-1)];
    R = S'*R*S;
end

TP = R;
     
    step = step + 1;
end


JJ1 = coefs(1,1:1+nz+ns);
if nd > 1
    JJ2 = coefs(2,1:1+nz+ns) + coefs(2,1+nz+ns+1)*JJ1;
end
JJ = [JJ1;JJ2];

if nd > 2
    for jj = 3:nd
        JJ = [JJ;(coefs(jj,:) + coefs(jj,1+nz+ns+1:1+nz+ns+(jj-1))*JJ)];
    end
end

dpol = JJ*F;

V = F'*P*F;

JJ
P
ipol = dpol(1)
hpol = dpol(2)
V


 
    
