%
%  Program 17.1: Three-period consumption savings model
%
% Ben J. Heijdra, Foundations of Modern Macroeconomics, Third Ed.
% Oxford University Press
%
% Chapter 17
%
% Ben J. Heijdra
% Groningen, November 12, 2015


clear
clf
clc
close all

% Parameters
delta_a      = 0.06;                       % Annual depreciation rate
delta        = 1 - (1 - delta_a)^25;
alpha        = 0.3;                        % Efficiency parameter capital

% Targets
r_a0         = 0.04;                       % Annual interest rate
r            = (1 + r_a0)^25 - 1;
rho_a        = 0.03;
rho          = (1 + rho_a)^25 - 1;
beta         = 1/(1+rho)
ystar0       = 1;                          % Normalized output
kstar0       = alpha / (r + delta);
w            = 1 - alpha;

% Optimal plans -- deterministic

H_1 = w * (1 + 1/(1+r) + 1/(1+r)^2 );
C_1 = H_1 / (1 + beta + beta^2);
C_2 = beta *(1+r)* C_1;
C_3 = beta *(1+r)* C_2;

A_2 = w - C_1;
A_3 = (1+r)*A_2 + w - C_2;
A_4 = (1+r)*A_3 + w - C_3;

% Optimal plans -- stochastic

e_1 = 0.75;
e_2 = 1.00;
e_3 = 1.25;

p_11 = 0.50
p_12 = 1 - p_11;
p_13 = 0;
p_21 = 0.25 ;
p_22 = 0.50 ;
p_23 = 1 - p_21 - p_22;
p_33 = 0.50;
p_32 = 1 - p_33;
p_31 = 0;

pi_11 = 0;
pi_12 = 1;
pi_13 = 0;

pi_21 = p_21;
pi_22 = p_22;
pi_23 = p_23;

pi_31 = p_21 * (p_11 + p_22);
pi_32 = p_21 * (1 - p_11) + p_22^2 + p_23 * (1 - p_33) ;
pi_33 = p_23 * (p_22 + p_33) ;

ExpH_1 = w + (w/(1+r)) * (pi_21 * e_1 + pi_22 * e_2 + pi_23 * e_3) ...
    + w * (1/(1+r))^2 * (pi_31 * e_1 + pi_32 * e_2 + pi_33 * e_3) ;

% Compute A_2 and A_3 chosen in period t=1 
% Load PAR structure

PAR.r         = r;
PAR.w         = w;
PAR.beta      = beta;

PAR.e_1       = e_1;
PAR.e_2       = e_2;
PAR.e_3       = e_3;

PAR.pi_21     = pi_21;
PAR.pi_22     = pi_22;
PAR.pi_23     = pi_23;

PAR.pi_31     = pi_31;
PAR.pi_32     = pi_32;
PAR.pi_33     = pi_33;


x = zeros(2,1);
[x(:,1),fval,exitflag,output] = fsolve(@Function17_01,[A_2 A_3]'...
    ,optimset('MaxFunEvals',1e6,'MaxIter',1000,'TolX',1.0e-8,'TolFun',1.0e-8),PAR);
                                     
fval;
output;
x;

A_2stoch  = x(1,1);
C_1stoch  = w - A_2stoch ;
A_3stoch  = x(2,1);

% Compute A_3 chosen in period t=2 for eta_2 = e_1

FONC2e_1  = @(x)  - 1 / ((1+r) * A_2stoch + e_1 * w - x) ...
   + beta*(1+r)* (p_11 / ((1+r)*x + e_1 * w) ...
   + p_12 / ((1+r)*x + e_2 * w) + p_13 / ((1+r)*x + e_3 * w));
A_3stoch_e_1  = fsolve(FONC2e_1,A_3stoch,optimset('MaxFunEvals',1e6,'MaxIter',...
    1000,'TolX',1.0e-6,'TolFun',1.0e-6));
C_2stoch_e_1  = (1+r) * A_2stoch + e_1 * w -  A_3stoch_e_1;
C_3stoch_e_1_e_1 = (1+r) * A_3stoch_e_1 + e_1 * w;
C_3stoch_e_1_e_2 = (1+r) * A_3stoch_e_1 + e_2 * w;

% Compute A_3 chosen in period t=2 for eta_2 = e_2

FONC2e_2  = @(x)  - 1 / ((1+r) * A_2stoch + e_2 * w - x) ...
   + beta*(1+r)* (p_21 / ((1+r)*x + e_1 * w) ...
   + p_22 / ((1+r)*x + e_2 * w) + p_23 / ((1+r)*x + e_3 * w));
A_3stoch_e_2  = fsolve(FONC2e_2,A_3stoch,optimset('MaxFunEvals',1e6,'MaxIter',...
    1000,'TolX',1.0e-6,'TolFun',1.0e-6));
C_2stoch_e_2  = (1+r) * A_2stoch + e_2 * w -  A_3stoch_e_2;
C_3stoch_e_2_e_1 = (1+r) * A_3stoch_e_2 + e_1 * w;
C_3stoch_e_2_e_2 = (1+r) * A_3stoch_e_2 + e_2 * w;
C_3stoch_e_2_e_3 = (1+r) * A_3stoch_e_2 + e_3 * w;

% Compute A_3 chosen in period t=2 for eta_2 = e_3

FONC2e_3  = @(x)  - 1 / ((1+r) * A_2stoch + e_3 * w - x) ...
   + beta*(1+r)* (p_31 / ((1+r)*x + e_1 * w) ...
   + p_32 / ((1+r)*x + e_2 * w) + p_33 / ((1+r)*x + e_3 * w));
A_3stoch_e_3  = fsolve(FONC2e_3,A_3stoch,optimset('MaxFunEvals',1e6,'MaxIter',...
    1000,'TolX',1.0e-6,'TolFun',1.0e-6));
C_2stoch_e_3  = (1+r) * A_2stoch + e_3 * w -  A_3stoch_e_3;
C_3stoch_e_3_e_2 = (1+r) * A_3stoch_e_3 + e_2 * w;
C_3stoch_e_3_e_3 = (1+r) * A_3stoch_e_3 + e_3 * w;





% Value functions and policy fuctions -- deterministic

avals = linspace(0,0.20,201);

V_1a = -(1+beta+beta^2) * log(1+beta+beta^2) ...
    +  beta * (1+2*beta) * log(beta) + (1+beta+beta^2) * ...
log((1+r)*avals + w *(1 + 1/(1+r) + 1/(1+r)^2 ) ) ...
    + beta*(1+2*beta) * log(1+r);

C_1a =  (1/(1+beta+beta^2)) * ((1+r)*avals  ...
    + w *(1 + 1/(1+r) + 1/(1+r)^2 ) ) ;
Aplus_1a = (1+r)*avals +w - C_1a;


V_2a = beta * log(beta) - (1+beta) * log(1+beta) ...
    + (1+beta) * log( (1+r)*avals + w *(1 + 1/(1+r)) ) ...
    + beta * log(1+r);

C_2a = (1/(1+beta)) * ((1+r)*avals  ...
    + w *(1 + 1/(1+r)) ) ;
Aplus_2a = (1+r)*avals +w - C_2a;

V_3a = log((1+r)*avals + w);

C_3a = (1+r)*avals + w ;
Aplus_3a = 0 + 0*avals;

figure(1)
plot(avals,V_1a,'-','LineWidth',3)
hold on;
plot(avals,V_2a,'--','LineWidth',3)
plot(avals,V_3a,':','LineWidth',3)
hold off;
legend('V_1(a)','V_2(a)','V_3(a)','Location','Best')
box off
xlabel('asset levels (a)')
print('-deps2','-loose','F17_01_V')
pause(1)

figure(2)
plot(avals,C_1a,'-','LineWidth',3)
hold on;
plot(avals,C_2a,'--','LineWidth',3)
plot(avals,C_3a,':','LineWidth',3)
hold off;
legend('C_1(a)','C_2(a)','C_3(a)','Location','Best')
box off
xlabel('asset levels (a)')
print('-deps2','-loose','F17_01_C')
pause(1)

figure(3)
plot(avals,Aplus_1a,'-','LineWidth',3)
hold on;
plot(avals,Aplus_2a,'--','LineWidth',3)
plot(avals,Aplus_3a,':','LineWidth',3)
hold off;
legend('A^{+}_1(a)','A^{+}_2(a)','A^{+}_3(a)','Location','Best')
box off
ylim ([-0.05 0.3])
xlabel('asset levels (a)')
print('-deps2','-loose','F17_01_Aplus')
pause(1)


% Value functions and policy functions -- stochastic

V_3_e_1 = log((1+r)*avals + e_1 * w);
V_3_e_2 = log((1+r)*avals + e_2 * w);
V_3_e_3 = log((1+r)*avals + e_3 * w);

C_3_e_1 = (1+r)*avals + e_1 * w;
C_3_e_2 = (1+r)*avals + e_2 * w;
C_3_e_3 = (1+r)*avals + e_3 * w;

% Compute policy function C_2(a,e_i)
tic

for i=1:length(avals)
    
FONC2C_2e_1  = @(x) -1/x + beta * (p_11 / ((1+r)* avals(i) ...
    + e_1 * w - x + e_1 * w / (1+r)) ...
    + p_12 / ((1+r)* avals(i) ...
    + e_1 * w - x + e_2 * w / (1+r)) ...
    + p_13 / ((1+r)* avals(i) ...
    + e_1 * w - x + e_3 * w / (1+r)) );

C_2_e_1(i)  = fsolve(FONC2C_2e_1,C_2stoch_e_1,optimset('MaxFunEvals',1e6,'MaxIter',...
    1000,'TolX',1.0e-6,'TolFun',1.0e-6));

FONC2C_2e_2  = @(x) -1/x + beta * (p_21 / ((1+r)* avals(i) ...
    + e_2 * w - x + e_1 * w / (1+r)) ...
    + p_22 / ((1+r)* avals(i) ...
    + e_2 * w - x + e_2 * w / (1+r)) ...
    + p_23 / ((1+r)* avals(i) ...
    + e_2 * w - x + e_3 * w / (1+r)) );

C_2_e_2(i)  = fsolve(FONC2C_2e_2,C_2stoch_e_2,optimset('MaxFunEvals',1e6,'MaxIter',...
    1000,'TolX',1.0e-6,'TolFun',1.0e-6));

FONC2C_2e_3  = @(x) -1/x + beta * (p_31 / ((1+r)* avals(i) ...
    + e_3 * w - x + e_1 * w / (1+r)) ...
    + p_32 / ((1+r)* avals(i) ...
    + e_3 * w - x + e_2 * w / (1+r)) ...
    + p_33 / ((1+r)* avals(i) ...
    + e_3 * w - x + e_3 * w / (1+r)) );

C_2_e_3(i)  = fsolve(FONC2C_2e_3,C_2stoch_e_3,optimset('MaxFunEvals',1e6,'MaxIter',...
    1000,'TolX',1.0e-6,'TolFun',1.0e-6));

end

Aplus_2_e_1 = (1+r) * avals + e_1 * w - C_2_e_1 ;
Aplus_2_e_2 = (1+r) * avals + e_2 * w - C_2_e_2 ;
Aplus_2_e_3 = (1+r) * avals + e_3 * w - C_2_e_3 ;

V_2_e_1 = log(C_2_e_1) + beta * ( ...
      p_11 * log((1+r) * Aplus_2_e_1 + e_1 * w) ...
    + p_12 * log((1+r) * Aplus_2_e_1 + e_2 * w) ...
    + p_13 * log((1+r) * Aplus_2_e_1 + e_3 * w) );

V_2_e_2 = log(C_2_e_2) + beta * ( ...
      p_21 * log((1+r) * Aplus_2_e_2 + e_1 * w) ...
    + p_22 * log((1+r) * Aplus_2_e_2 + e_2 * w) ...
    + p_23 * log((1+r) * Aplus_2_e_2 + e_3 * w) );

V_2_e_3 = log(C_2_e_3) + beta * ( ...
      p_31 * log((1+r) * Aplus_2_e_3 + e_1 * w) ...
    + p_32 * log((1+r) * Aplus_2_e_3 + e_2 * w) ...
    + p_33 * log((1+r) * Aplus_2_e_3 + e_3 * w) );

toc

figure(4)
%plot(avals,V_3a,':','LineWidth',2)
hold on;
plot(avals,V_3_e_1,':','LineWidth',3)
plot(avals,V_3_e_2,'--','LineWidth',3)
plot(avals,V_3_e_3,'-','LineWidth',3)
hold off;
%legend('V_3(a)','V_3(a,e_1)','V_3(a,e_2)','V_3(a,e_3)','Location','Best')
legend('V_3(a,e_1)','V_3(a,e_2)','V_3(a,e_3)','Location','Best')
box off
xlabel('asset levels (a)')
print('-deps2','-loose','F17_04_V3')
pause(1)

figure(5)
%plot(avals,V_2a,':','LineWidth',2)
hold on;
plot(avals,V_2_e_1,':','LineWidth',3)
plot(avals,V_2_e_2,'--','LineWidth',3)
plot(avals,V_2_e_3,'-','LineWidth',3)
hold off;
%legend('V_2(a)','V_2(a,e_1)','V_2(a,e_2)','V_2(a,e_3)','Location','Best')
legend('V_2(a,e_1)','V_2(a,e_2)','V_2(a,e_3)','Location','Best')
box off
xlabel('asset levels (a)')
print('-deps2','-loose','F17_04_V2')
pause(1)

% Interpolate the V_2_e_i functions

V_2_e_1_interp  = @(x) interp1(avals,V_2_e_1,x,'spline');
V_2_e_2_interp  = @(x) interp1(avals,V_2_e_2,x,'spline');
V_2_e_3_interp  = @(x) interp1(avals,V_2_e_3,x,'spline');


a_1 = 0;
tic

for i=1:length(avals)
    
V_1_e_2 = @(x) -log(x) - beta * ...
    (p_21 * V_2_e_1_interp((1+r)* avals(i) + e_2 * w - x) ...
   + p_22 * V_2_e_2_interp((1+r)* avals(i) + e_2 * w - x) ...
   + p_23 * V_2_e_3_interp((1+r)* avals(i) + e_2 * w - x) ) ;

options = optimset('Display','iter','TolX',1.0e-8,'TolFun',1.0e-8);
[x,fval] = fminsearch(V_1_e_2,0.62,options);

C_1_e_2(i)      = x;
Value(i)        = -fval;

end

toc

Aplus_1_e_2     = (1 + r) * avals + w -  C_1_e_2 ;

V_1_e_2 = Value;


figure(6)
%plot(avals,V_1a,':','LineWidth',2)
hold on;
plot(avals,V_1_e_2,'-','LineWidth',3)
hold off;
%legend('V_1(a)','V_1(a,e_1)','Location','Best')
legend('V_1(a,e_2)','Location','Best')
box off
xlabel('asset levels (a)')
print('-deps2','-loose','F17_04_V1')
pause(1)

figure(7)
%plot(avals,C_1a,':','LineWidth',2)
hold on;
plot(avals,C_1_e_2,'-','LineWidth',3)
hold off;
%legend('C_1(a)','C_1(a,e_2)','Location','Best')
legend('C_1(a,e_2)','Location','Best')
box off
xlabel('asset levels (a)')
print('-deps2','-loose','F17_05_C1')
pause(1)

figure(8)
%plot(avals,C_2a,':','LineWidth',2)
hold on;
plot(avals,C_2_e_1,':','LineWidth',3)
plot(avals,C_2_e_2,'--','LineWidth',3)
plot(avals,C_2_e_3,'-','LineWidth',3)
plot([0.0835 0.0835],[0 0.9411],':','LineWidth',1)
text(0.08,1.00,'A','FontSize',14)
text(0.085,0.75,'B','FontSize',14)
text(0.085,0.61,'C','FontSize',14)
hold off;
%legend('C_2(a)','C_2(a,e_1)','C_2(a,e_2)','C_2(a,e_3)','Location','Best')
legend('C_2(a,e_1)','C_2(a,e_2)','C_2(a,e_3)','Location','Best')
box off
xlabel('asset levels (a)')
print('-deps2','-loose','F17_05_C2')
pause(1)

figure(9)
%plot(avals,C_3a,':','LineWidth',2)
hold on;
plot(avals,C_3_e_1,':','LineWidth',3)
plot(avals,C_3_e_2,'--','LineWidth',3)
plot(avals,C_3_e_3,'-','LineWidth',3)
plot([0.0884 0.0884],[0 0.9357],':','LineWidth',1)
text(0.085,1.00,'A','FontSize',14)
text(0.090,0.70,'B','FontSize',14)
plot([0.1243 0.1243],[0 1.2064],':','LineWidth',1)
text(0.12,1.26,'C','FontSize',14)
text(0.127,0.96,'D','FontSize',14)
text(0.127,0.79,'E','FontSize',14)
plot([0.1564 0.1564],[0 1.2921],':','LineWidth',1)
text(0.156,1.36,'F','FontSize',14)
text(0.160,1.07,'G','FontSize',14)
hold off;
%legend('C_3(a)','C_3(a,e_1)','C_3(a,e_2)','C_3(a,e_3)','Location','Best')
legend('C_3(a,e_1)','C_3(a,e_2)','C_3(a,e_3)','Location','Best')
box off
xlabel('asset levels (a)')
print('-deps2','-loose','F17_05_C3')
pause(1)

figure(10)
plot(avals,Aplus_1_e_2,'-','LineWidth',3)
legend('A^{+}_1(a,e_{2})','Location','Best')
box off
xlabel('asset levels (a)')
print('-deps2','-loose','F17_05_A1plus')
pause(1)

figure(11)
plot(avals,Aplus_2_e_1,':','LineWidth',3)
hold on;
plot(avals,Aplus_2_e_2,'--','LineWidth',3)
plot(avals,Aplus_2_e_3,'-','LineWidth',3)
plot([0.0835 0.0835],[0 0.1564],':','LineWidth',1)
text(0.08,0.17,'A','FontSize',14)
text(0.0850,0.11,'B','FontSize',14)
text(0.085,0.075,'C','FontSize',14)
hold off;
legend('A^{+}_2(a,e_{1})','A^{+}_2(a,e_{2})','A^{+}_2(a,e_{3})','Location','Best')
box off
xlabel('asset levels (a)')
print('-deps2','-loose','F17_05_A2plus')
pause(1)


% Model parameters
    disp(' ')
    disp('Exogenous variables')
    disp(sprintf('      initial assets A_1:                            %8.4f',0))
    disp(sprintf('      initial human wealth H_1:                      %8.4f',H_1))
    disp(sprintf('      initial expected human wealth E_1H_1:          %8.4f',ExpH_1))
    disp(sprintf('      wage w_t:                                      %8.4f',w))
    disp(sprintf('      interest rate r_t:                             %8.4f',r))
        
    disp('Deterministic choices')
    disp(sprintf('      consumption C_1:                               %8.4f',C_1))
    disp(sprintf('      consumption C_2:                               %8.4f',C_2))
    disp(sprintf('      consumption C_3:                               %8.4f',C_3))
    disp(sprintf('      assets A_2:                                    %8.4f',A_2))
    disp(sprintf('      assets A_3:                                    %8.4f',A_3))
    disp(sprintf('      assets A_4:                                    %8.4f',A_4))
    
    disp('Stochastic choices - Traditional method')
    disp(sprintf('      consumption C_1:                               %8.4f',C_1stoch))
    disp(sprintf('      consumption C_2(e_1):                          %8.4f',C_2stoch_e_1))
    disp(sprintf('      consumption C_2(e_2):                          %8.4f',C_2stoch_e_2))
    disp(sprintf('      consumption C_2(e_3):                          %8.4f',C_2stoch_e_3))
    disp(sprintf('      assets A_2:                                    %8.4f',A_2stoch))
    disp(sprintf('      assets A_3(e_1):                               %8.4f',A_3stoch_e_1))
    disp(sprintf('      assets A_3(e_2):                               %8.4f',A_3stoch_e_2))
    disp(sprintf('      assets A_3(e_3):                               %8.4f',A_3stoch_e_3))
    disp(sprintf('      consumption C_3(e_1,e_1):                      %8.4f',C_3stoch_e_1_e_1))
    disp(sprintf('      consumption C_3(e_1,e_2):                      %8.4f',C_3stoch_e_1_e_2))
    disp(sprintf('      consumption C_3(e_2,e_1):                      %8.4f',C_3stoch_e_2_e_1))
    disp(sprintf('      consumption C_3(e_2,e_2):                      %8.4f',C_3stoch_e_2_e_2))
    disp(sprintf('      consumption C_3(e_2,e_3):                      %8.4f',C_3stoch_e_2_e_3))
    disp(sprintf('      consumption C_3(e_3,e_2):                      %8.4f',C_3stoch_e_3_e_2))
    disp(sprintf('      consumption C_3(e_3,e_3):                      %8.4f',C_3stoch_e_3_e_3))

    disp('Stochastic choices - DP method')
    disp(sprintf('      consumption C_1:                               %8.4f',C_1_e_2(1)))
    disp(sprintf('      assets A_2:                                    %8.4f',Aplus_1_e_2(1)))

