%HJB EQUATION FOR NEOCLASSICAL GROWTH MODEL, IMPLICIT METHOD
clear all; clc;
tic;
s = 2;
a = 0.3;
d = 0.05;
r = 0.05;
A = 1;

kss = (a*A/(r+d))^(1/(1-a));
css = A.*kss.^a - d.*kss;
dVss = css.^(-s);

I=1001; %NEEDS TO BE AN ODD NUMBER FOR kss TO BE AT CENTER OF GRID
kmin = 0.5*kss;
kmax = 1.5*kss;
k = linspace(kmin,kmax,I)';
dk = (kmax-kmin)/(I-1);
iss = (I+1)/2;
check = k(iss)-kss; %CHECK THAT MID-POINT OF GRID IS kss
disp(check)

Delta = 5; %STEP SIZE: CAN BE LARGE BECAUSE WE'RE USING AN IMPLICIT FD METHOD
maxit=100;
crit = 10^(-10);

dV = zeros(I,1);
c = zeros(I,1);
B = zeros(I,I);

%INITIAL GUESS
v(:,1)= (A.*k.^a - d.*k).^(1-s)/(1-s)/r;

for n=1:maxit
V = v(:,n);

%AT STEADY STATE: USE BOUNDARY CONDITION
dV(iss) = dVss;
%BELOW STEADY STATE: USE FORWARD DIFFERENCE
for i=1:iss-1
    dV(i)=(V(i+1)-V(i))/dk;
end
%ABOVE STEADY STATE: USE BACKWARD DIFFERENCE
for i=iss+1:I
    dV(i)=(V(i)-V(i-1))/dk;
end

c = dV.^(-1/s); %FOC
U = c.^(1-s)/(1-s);

g = A.*k.^a - d.*k - c;

%CONSTRUCT MATRIX
for i=1:iss-1
    x(i) = 0;
    y(i) = r + g(i)/dk + 1/Delta;
    z(i) = -g(i)/dk;
end
for i=iss+1:I
    x(i) = g(i)/dk;
    y(i) = r - g(i)/dk + 1/Delta;
    z(i) = 0;
end
x(iss) = 0;
y(iss) = r+1/Delta;
z(iss) = 0;

for i=2:I-1
    B(i,i-1) = x(i);
    B(i,i) = y(i);
    B(i,i+1) = z(i);
end
B(1,1)=y(1); B(1,2) = z(1);
B(I,I)=y(I); B(I,I-1) = x(I);

b = U + V/Delta;
b(iss) = U(iss)+V(iss)/Delta + dVss*g(iss);

v(:,n+1) = B\b; %SOLVE SYSTEM OF EQUATIONS
Vchange(:,n) = (v(:,n+1)-v(:,n));
dist(n) = max(abs(Vchange(:,n)));
if dist(n)<crit
    disp('Value Function Converged, Iteration = ')
    disp(n)
    break
end
end
toc;
set(gca,'FontSize',14)
plot(dist,'LineWidth',2)
grid
xlabel('Iteration')
ylabel('||V^{n+1} - V^n||')

Verr = c.^(1-s)/(1-s) + dV.*(A.*k.^a - d.*k - c) - r.*V;

set(gca,'FontSize',14)
plot(k,Verr,'LineWidth',2)
grid
xlabel('k')
ylabel('Error in HJB Equation')
xlim([kmin kmax])

kdot = A.*k.^a - d.*k - c;

set(gca,'FontSize',12)
plot(k,V,'LineWidth',2)
grid
xlabel('k')
ylabel('V(k)')
xlim([kmin kmax])

set(gca,'FontSize',14)
plot(k,c,'LineWidth',2)
grid
xlabel('k')
ylabel('c(k)')
xlim([kmin kmax])

set(gca,'FontSize',14)
plot(k,kdot,k,zeros(1,I),'--','LineWidth',2)
grid
xlabel('k')
ylabel('kdot(k)')
xlim([kmin kmax])