%HJB EQUATION FOR RBC MODEL WITH ORNSTEIN-UHLENBECK PROCESS FOR LOG PRODUCTIVITY
%THANKS TO TOMMASO PORZIO FOR OPTIMIZING FOR SPEED
clear all; 
clc;

s = 2;
a = 0.3;
d = 0.05;
r = 0.05;


%ORNSTEIN-UHLENBECK PROCESS dlog(A) = -the*log(A)dt + sig2*dW
%STATIONARY DISTRIBUTION IS log(A) ~ N(0,Var) WHERE Var = sig2/(2*the)
Var = 0.13;
Amean = exp(Var/2); %MEAN OF LOG-NORMAL DISTRIBUTION N(0,Var)
Corr = 0.9;
% Corr = 0.7;
the = -log(Corr);
sig2 = 2*the*Var;

J = 40;
Amin = Amean*0.6;
Amax = Amean*1.4;
A = linspace(Amin,Amax,J)';
dA = (Amax-Amin)/(J-1);
dA2 = dA^2;
plot(A,lognpdf(A,0,Var)) %PLOT STATIONARY DISTRIBUTION AND CHECK THAT CHOICE OF GRID DOESN'T CUT OFF TOO MUCH OF TAILS
mu = (-the*log(A) + sig2/2).*A; %DRIFT (FROM ITO'S LEMMA)
s2 = sig2.*A.^2; %VARIANCE (FROM ITO'S LEMMA)

kss = (a*Amean/(r+d))^(1/(1-a)); %DETERMINISTIC STEADY STATE

I=151;
kmin = 0.5*kss;
kmax = 1.5*kss;
k = linspace(kmin,kmax,I)';
dk = (kmax-kmin)/(I-1);

% Define size of initial Delta step: first one is exp(tom)
s_start = 1.4;
% NOTE: if you increase tom above 1.4 (on my pc - it may change across pcs), then you get the error for
% too high Delta. Of course as tom increases then speed and number of
% iterations are reduced. 

%Delta = 0.1*dk; %CANNOT BE TOO LARGE RELATIVE TO dk and dA, SEE CANDLER p.181 (DISCUSSION OF CFL NUMBER)
maxit= 300000;
%spline = linspace(tom^(1/8),1,maxit).^8;
%spline = linspace(tom,1,maxit);
spline = exp(linspace(s_start,0,maxit));
Delta = 0.002*spline;
crit = 10^(-6);

Vk = zeros(I,J);
c = zeros(I,J);
VA = zeros(I,J);
VAA = zeros(I,J);

kk = k*ones(1,J);
AA = ones(I,1)*A';
muA = ones(I,1)*mu';
s2A = ones(I,1)*s2';

v = zeros(I,J);
Vchange = v;
dist = zeros(maxit,1);

% NATURAL INITIAL GUESS
%v(:,:,1) = (AA.*kk.^a - d.*kk).^(1-s)/(1-s)/r;
%PROBLEM: THIS INITIAL GUESS CAN BE DOWNWARD-SLOPING SO CHANGE SLIGHTLY
dalt = 0.01;
v(:,:) = (AA.*kk.^a - dalt*kk).^(1-s)/(1-s)/r;


plot(k,v(:,:))
tic

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

%CENTRAL DIFFERENCE APPROXIMATION
Vk(2:I-1,:) = (V(3:I,:)-V(1:I-2,:))/(2*dk);
Vk(1,:) = (V(2,:)-V(1,:))/dk;
Vk(I,:) = (V(I,:)-V(I-1,:))/dk;
%Vk(I,:) = (V(I,:)-V(I-1,:))/(2*dk); %ALTERNATIVE: REFLECTING BARRIER AT k(I+1): (V(I+1)-V(I))/dk = 0;

VA(:,2:J-1) = (V(:,3:J)-V(:,1:J-2))/(2*dA);
VAA(:,2:J-1) =  (V(:,3:J)-2*V(:,2:J-1)+V(:,1:J-2))/dA2;

VA(:,1) = 0; %REFLECTING BARRIER AT Amin
VA(:,J) = 0; %REFLECTING BARRIER AT Amax
VAA(:,1) = (V(:,2)-2*V(:,1) + V(:,2))/dA2; %USING V(:,2)=V(:,0)
VAA(:,J) = (V(:,J-1)-2*V(:,J) + V(:,J-1))/dA2; %USING V(:,J+1)=V(:,J-1)


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

if isreal(c)==0
    disp('Error: Imaginary number')
    disp('Iteration:')
    disp(n)
    break
end


U = c.^(1-s)/(1-s);
Vchange(:,:) = U + Vk.*(AA.*kk.^a - d.*kk - c) + VA.*muA + (1/2).*VAA.*s2A - r.*V;
v1(:,:) = v(:,:) + Delta(n)*Vchange(:,:);
dist(n) = max(max(abs(Vchange(:,:))));
if dist(n)<crit
    disp('Value Function Converged, Iteration = ')
    disp(n)
    break
end
%if isint(n/500)==1
%   display('Iteration Number')
%   display(n)
%   display('Distance*e6')
%   display(dist(n)*10^(6))
%end

v = v1;
end
toc

% set(gca,'FontSize',14)
% plot(dist,'LineWidth',2)
% grid
% xlabel('Iteration')
% ylabel('||V^{n+1} - V^n||')

set(gca,'FontSize',14)
plot(dist(n-100:n),'LineWidth',2)
grid
xlabel('Iteration')
ylabel('||V^{n+1} - V^n||')

Verr = c.^(1-s)/(1-s) + Vk.*(AA.*kk.^a - d.*kk - c) + VA.*muA + (1/2).*VAA.*s2A - r.*V;

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


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

surf(A,k,v(:,:))
ylabel('k')
xlabel('A')
zlabel('V(k,A)')

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

surf(A,k,c)
ylabel('k')
xlabel('A')
zlabel('c(k,A)')


kdot = AA.*kk.^a - d.*kk - c;

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

surf(A,k,kdot)
ylabel('k')
xlabel('A')
zlabel('kdot(k,A)')