% Solving the standard RBC Model from Hansen 1985 by Linear-Quadratic
% Approximation -- competitive equilibrium
%
%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);   % just Y_bar/K_bar
YK_bar = BB;
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);
Yh_bar = Y_bar/h_bar;
C_bar = Y_bar - I_bar;

nz = 1;     % number of exogenous state variables
nS = 1;     % number of aggregate endogenous state variables
ns = 1;     % number of individual endogenous state variables
nD = 2;     % number of aggregate control variables
nd = 2;     % number of individual control variables

% Order of the variables: [Z,K,k,I,i,H,h] -> ex. states, agg. states, ind.
% states, agg. decision vars., ind. decision vars.
% --------------------------------------------------------------------

% 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 = 0;
Jk = (theta*YK_bar)/C_bar;
JI = 0;
JH = 0;
Ji = -1/C_bar;
Jh = (1-theta)*Yh_bar/C_bar - A/(1-h_bar);

Hzz = -(Y_bar*Y_bar)/(C_bar*C_bar);
HKK = (-theta*(theta-1)*YK_bar)/(K_bar*C_bar);
Hkk = -(theta*YK_bar)*(theta*YK_bar)/(C_bar*C_bar);
HII = 0;
HHH = (theta*(1-theta)*Yh_bar)/(h_bar*C_bar);
Hii = -1/(C_bar*C_bar);
Hhh = -(((1-theta)*Yh_bar)*((1-theta)*Yh_bar))/(C_bar*C_bar) - A/((1-h_bar)*(1-h_bar));

HzK = 0;
Hzk = (theta*YK_bar*(C_bar - Y_bar))/(C_bar*C_bar);
HzI = 0;
HzH = 0;
Hzi = Y_bar/(C_bar*C_bar);
Hzh = (1-theta)*Yh_bar*(C_bar - Y_bar)/(C_bar*C_bar);


HKk = theta*(theta-1)*YK_bar/(K_bar*C_bar);
HKI = 0;
HKH = -(theta*(1-theta)*Yh_bar)/(K_bar*C_bar);
HKi = 0;
HKh = theta*(1-theta)*Yh_bar/(K_bar*C_bar);

HkI = 0;
HkH = theta*(1-theta)*Yh_bar/(K_bar*C_bar);
Hki = theta*YK_bar/(C_bar*C_bar);
Hkh = -theta*(1-theta)*YK_bar*Yh_bar/(C_bar*C_bar);

HIH = 0;
HIi = 0;
HIh = 0;

HHi = 0;
HHh = -theta*(1-theta)*Yh_bar/(h_bar*C_bar);

Hih = (1-theta)*Yh_bar/(C_bar*C_bar);

% % Computing the derivatives numerically (using symbolic toolbox)
% alfa = sym('alfa');
% Z = sym('Z');
% K = sym('K');
% k = sym('k');
% I = sym('I');
% H = sym('H');
% i = sym('i');
% h = sym('h');
% zero = sym('0');
% vars = [Z,K,k,I,i];
% 
% obfun = log(alfa*Z*K^(alfa-1)*H^(1-alfa)*k + (1-alfa)*Z*K^alfa*H^(-alfa)*h - i) + A*log(1-h);
% Jz_sym = (alfa*K^(alfa-1)*H^(1-alfa)*k + (1-alfa)*K^alfa*H^(-alfa)*h)/...
%     (alfa*Z*K^(alfa-1)*H^(1-alfa)*k + (1-alfa)*Z*K^alfa*H^(-alfa)*h - i);
% JK_sym = (alfa*(alfa-1)*Z*K^(alfa-2)*H^(1-alfa)*k + alfa*(1-alfa)*Z*K^(alfa-1)*H^(-alfa)*h)/...
%     (alfa*Z*K^(alfa-1)*H^(1-alfa)*k + (1-alfa)*Z*K^alfa*H^(-alfa)*h - i);
% Jk_sym = (alfa*Z*K^(alfa-1)*H^(1-alfa))/...
%     (alfa*Z*K^(alfa-1)*H^(1-alfa)*k + (1-alfa)*Z*K^alfa*H^(-alfa)*h - i);
% JI_sym = zero;
% JH_sym = (alfa*(1-alfa)*Z*K^(alfa-1)*H^(-alfa)*k + (-alfa)*(1-alfa)*Z*K^(alfa)*H^(-alfa-1)*h)/...
%     (alfa*Z*K^(alfa-1)*H^(1-alfa)*k + (1-alfa)*Z*K^alfa*H^(-alfa)*h - i);
% Ji_sym = -1/(alfa*Z*K^(alfa-1)*H^(1-alfa)*k + (1-alfa)*Z*K^alfa*H^(-alfa)*h - i);
% Jh_sym = ((1-alfa)*Z*K^(alfa)*H^(-alfa))/...
%     (alfa*Z*K^(alfa-1)*H^(1-alfa)*k + (1-alfa)*Z*K^alfa*H^(-alfa)*h - i) - A/(1-h);
% 
% 
% JJ = jacobian(obfun,[Z,K,k,I,H,i,h])';
% JJ_bar = eval(subs(JJ,{alfa,Z,K,k,I,H,i,h},[theta,lambda_bar,K_bar,K_bar,I_bar,h_bar,I_bar,h_bar]));
% 
% Hz = jacobian(Jz_sym,[Z,K,k,I,H,i,h]);
% Hz_bar = eval(subs(Hz,{alfa,Z,K,k,I,H,i,h},[theta,lambda_bar,K_bar,K_bar,I_bar,h_bar,I_bar,h_bar]));
% 
% HK = jacobian(JK_sym,[Z,K,k,I,H,i,h]);
% HK_bar = eval(subs(HK,{alfa,Z,K,k,I,H,i,h},[theta,lambda_bar,K_bar,K_bar,I_bar,h_bar,I_bar,h_bar]));
% 
% Hk = jacobian(Jk_sym,[Z,K,k,I,H,i,h]);
% Hk_bar = eval(subs(Hk,{alfa,Z,K,k,I,H,i,h},[theta,lambda_bar,K_bar,K_bar,I_bar,h_bar,I_bar,h_bar]));
% 
% HI = jacobian(JI_sym,[Z,K,k,I,H,i,h]);
% HI_bar = subs(HI,{alfa,Z,K,k,I,H,i,h},[theta,lambda_bar,K_bar,K_bar,I_bar,h_bar,I_bar,h_bar]);
% 
% H_H = jacobian(JH_sym,[Z,K,k,I,H,i,h]);
% H_H_bar = eval(subs(H_H,{alfa,Z,K,k,I,H,i,h},[theta,lambda_bar,K_bar,K_bar,I_bar,h_bar,I_bar,h_bar]));
% 
% Hi = jacobian(Ji_sym,[Z,K,k,I,H,i,h]);
% Hi_bar = eval(subs(Hi,{alfa,Z,K,k,I,H,i,h},[theta,lambda_bar,K_bar,K_bar,I_bar,h_bar,I_bar,h_bar]));
% 
% Hh = jacobian(Jh_sym,[Z,K,k,I,H,i,h]);
% Hh_bar = eval(subs(Hh,{alfa,Z,K,k,I,H,i,h},[theta,lambda_bar,K_bar,K_bar,I_bar,h_bar,I_bar,h_bar]));
% 
% HH = [Hz;HK;Hk;HI;H_H;Hi;Hh];
% HH_bar = [Hz_bar;HK_bar;Hk_bar;HI_bar;H_H_bar;Hi_bar;Hh_bar];


% for l = 1:length(JJ_bar)
%     for ll = l:length(JJ_bar)
%         HH(l,ll) = diff(JJ(l),vars(ll));
%         HH(ll,l) = diff(JJ(l),vars(ll));
%         HH_bar(l,ll) = subs(HH(l,ll),{alfa,Z,K,k,I,i},[alpha,Z_bar,K_bar,K_bar,I_bar,I_bar]);
%         HH_bar(ll,l) = subs(HH(l,ll),{alfa,Z,K,k,I,i},[alpha,Z_bar,K_bar,K_bar,I_bar,I_bar]);
%     end
% end
% HH = jacobian(JJ,[Z,K,k,I,i]);
% alfa = alpha;Z = Z_bar;K = K_bar;k = K_bar;I = I_bar;i = I_bar;
% HH_bar = subs(HH);
% HH_bar = eval(subs(HH,{alfa,Z,K,k,I,i},[alpha,Z_bar,K_bar,K_bar,I_bar,I_bar]));


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

H_bar = [Hzz HzK Hzk HzI Hzi HzH Hzh
         HzK HKK HKk HKI HKi HKH HKh
         Hzk HKk Hkk HkI Hki HkH Hkh
         HzI HKI HkI HII HIi HIH HIh
         Hzi HKi Hki HIi Hii HHi Hih
         HzH HKH HkH HIH HHi HHH HHh
         Hzh HKh Hkh HIh Hih HHh Hhh];
        
% H_bar = HH_bar;         
         
W_bar = [lambda_bar;K_bar;K_bar;I_bar;I_bar;h_bar;h_bar];
F = [1;lambda_bar;K_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
% const. z K k I i H h
B = [1 0 0 0 0 0 0 0 0 0 0 0             % the last four columns are for tomorrow's states
     1-gamma gamma 0 0 0 0 0 0 0 0 0 0     % Don't forget 1-gamma here!
     0 0 1-delta 0 1 0 0 0 0 0 0 0
     0 0 0 1-delta 0 1 0 0 0 0 0 0];
 
% initial guess for P
TP = -0.01*eye(1+nz+nS+ns);       % the iterated P
P = ones(1+nz+nS+ns);
     
% Create the matrix R
RR = zeros(1+nz+nS+ns+nD+nd+1+nz+nS+ns);
RR(1:1+nz+nS+ns+nD+nd,1:1+nz+nS+ns+nD+nd) = Q;
RR(1+nz+nS+ns+nD+nd+1:1+nz+nS+ns+nD+nd+1+nz+nS+ns,1+nz+nS+ns+nD+nd+1:1+nz+nS+ns+nD+nd+1+nz+nS+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+ns+nD+nd+1:1+nz+nS+ns+nD+nd+1+nz+nS+ns,1+nz+nS+ns+nD+nd+1:1+nz+nS+ns+nD+nd+1+nz+nS+ns) = beta*P;

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

% Eliminating decision variables
index = false(1,1+nz+nS+ns+nD+nd);
z_index = index;
z_index(1:1+nz) = true;
nS_index = index;
nS_index(1+nz+1:1+nz+nS) = true;
ns_index = index;
ns_index(1+nz+nS+1:1+nz+nS+ns) = true;
nD_index = index;
nD_index(1+nz+nS+ns+1:2:1+nz+nS+ns) = true;

coefs = zeros(nD,1+nz+nS+ns+nD+nd);
for jj = nd:-1:1
    J = 1+nz+nS+ns+2*jj;
    index = false(1,J-2);
    z_index = index;
    z_index(1:1+nz) = true;
    S_index = index;
    S_index(1+nz+1:1+nz+nS) = true;
    s_index = index;
    s_index(1+nz+nS+1:1+nz+nS+ns) = true;
    
    % Eliminating aggregate decision variables
    denom = -(R(J,J) + R(J,J-1));
    D_coefs = zeros(1,J-2);
    D_coefs(z_index) = R(J,z_index);
    D_coefs(S_index) = R(J,S_index) + R(J,s_index);
    D_coefs(s_index) = 0;
    if jj>1
        D_index = index;
        D_index(1+nz+nS+ns+1:2:J-3) = true;
        d_index = index;
        d_index(1+nz+nS+ns+2:2:J-2) = true;
        
        D_coefs(D_index) = R(J,D_index) + R(J,d_index);
        D_coefs(d_index) = 0;
    end
    
    D_coefs = D_coefs/denom;
    
    % Eliminating the individual decision variables
    d_coefs = zeros(1,J-1);
    d_coefs = -1/R(J,J)*R(J,1:J-1);
    
    S = [eye(J-1);d_coefs];
    R = S'*R*S;
    
    % and again the aggregate one
    S = [eye(J-2);D_coefs];
    R = S'*R*S;
    
    coefs(jj,1:J-2) = D_coefs;
end

TP = R;
     
    step = step + 1;
end


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

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

Dpol = JJ*F;

V = F'*P*F;

JJ
P
Ipol = Dpol(1)
Hpol = Dpol(2)
V


 

