% Code adapted from the work by Mattias Almgren (IIES)
clear all; 
clc;


%% Parameters
cbar    = 1;        % satiation point
r       = 0.05;     % exogeneous interest rate
d_init  = 0.1;      % initial level of real debt
Np      = 7;        % number periods in simulation

A1 = 0.9;           % initial productivitiy, Case 1
A2 = 1;             % initial productivitiy, Case 2

% Construct vector of shocks: shock affecting in period 4
Avec1    = [ones(floor(Np/2),1)*A1; ones(round(Np/2),1)*1.1*A1 ]; %productivity vector, Case 1
Avec2    = [ones(floor(Np/2),1)*A2; ones(round(Np/2),1)*1.1*A2 ]; %productivity vector, Case 2
Amat     = [Avec1 Avec2]; %Combine matrices

%% Simulation
% Preallocate vectors for storing results
y       = NaN(Np,2); 
c       = NaN(Np,2);
h       = NaN(Np,2);
tb      = NaN(Np,2);
ca      = NaN(Np,2);
d       = NaN(Np,2);
t       = NaN(Np,2);    % time

for jj = 1:2
    %period 1
    c(1,jj)    = (1/(1+Amat(1,jj)^2))  * (Amat(1,jj)^2 * cbar - r * d_init);  
    h(1,jj)    = (Amat(1,jj)/(1+Amat(1,jj)^2)) * (cbar+r * d_init);
    y(1,jj)    = Amat(1,jj)*h(1,jj);
    tb(1,jj)   = y(1,jj)-c(1,jj);
    ca(1,jj)   = 0;
    d(1,jj)    = d_init;
    t(1,jj)    = 1-round(Np/2);

    %periods 2 through Np
    for tt = 2:Np
        t(tt,jj)   = tt-round(Np/2);
        c(tt,jj)   = (1/(1+Amat(tt,jj)^2)) * (Amat(tt,jj)^2*cbar-r*d(tt-1,jj));
        h(tt,jj)   = (Amat(tt,jj)/(1+Amat(tt,jj)^2)) * (cbar + r*d(tt-1,jj));
        y(tt,jj)   = Amat(tt,jj)*h(tt,jj);
        tb(tt,jj)  = y(tt,jj)-c(tt,jj);
        ca(tt,jj)  = tb(tt,jj) - r*d(tt-1,jj);
        d(tt,jj)   = d(tt-1,jj);
    end
    
end

%plot results
subplot(3,1,1)
plot(t,y(:,1),'-.b*') 
hold on
plot(t,y(:,2),'-.r*') 
title('output')
subplot(3,1,2)
plot(t,c(:,1),'-.b*')  
hold on
plot(t,c(:,2),'-.r*') 
title('consumption')
subplot(3,1,3)
plot(t,h(:,1),'-.b*')  
hold on
plot(t,h(:,2),'-.r*') 
title('labor')
saveas(gcf,'shock.png')