%This script performs value function iteration for the baseline model.
%Written by Philip Shaw, Fordham University, 2022.
clear all
global delta eta alpha beta
delta=1; %depreciation rate
eta=1; %utility parameter
alpha=.27; %production parameter
beta=.994; %patience parameter
eps=.01; %tolerance parameter
n=250; %grid size
kstar=((1-beta*(1-delta))/(alpha*beta))^(1/(alpha-1)); %steady-state capital
cstar=kstar^alpha+(1-delta)*kstar-kstar; %steady state consumption
Kl=.75*kstar; %lower bound for capital grid
Ku=1.25*kstar; %upper bound for capital grid
kstep=(Ku-Kl)/(n-1); %step size for capital
kgrid=Kl:kstep:Ku; %capital grid
%v0=zeros(n,1); %initialize value function 
v0=ones(n,1)*(utilchp4(cstar,eta))/(1-beta);%initialize value function
kprime0=ones(n,1)*.33; %initialize policy function
%load vnew0 
%v0=vnew;
dist=1; %initialize distance for the value function
distpol=.01; %initialize distance for the policy function
it=1; %initialize counter
eps2=.000001;
tic %start clock
linint=1; %toggle linear interpolation
while(distpol>eps*(1-beta))
    for i=1:length(kgrid)
        %Step 1
        for j=1:length(kgrid)
            ct=kgrid(i)^alpha+(1-delta)*kgrid(i)-kgrid(j); %consumption
            ct=ct*(ct>0)+eps2; %correct for potential negative consumption
            w(i,j)=utilchp4(ct,eta)+beta*v0(j,1);
        end
        %Step 2
        jind=find(max(w(i,:))==w(i,:)); %find the j that max rhs of Bellman
        %Step 3 
        if linint==1 %linear interpolation
            if jind==1
                [ktild]=gss(kgrid(jind),kgrid(jind+1),@(x)valuelin(kgrid(jind),kgrid(jind+1),v0(jind),v0(jind+1),kgrid(i),x));
                vnew(i,1)=valuelin(kgrid(jind),kgrid(jind+1),v0(jind),v0(jind+1),kgrid(i),ktild);
            else
                [ktild]=gss(kgrid(jind-1),kgrid(jind+1),@(x)valuelin(kgrid(jind-1),kgrid(jind+1),v0(jind-1),v0(jind+1),kgrid(i),x));
                vnew(i,1)=valuelin(kgrid(jind-1),kgrid(jind+1),v0(jind-1),v0(jind+1),kgrid(i),ktild);
            end
            kprime(i,1)=ktild;
            
           
        else
            vnew(i,1)=w(i,jind); %update value function
            h(i,1)=jind; %form your policy function
            kprime(i,1)=kgrid(h(i,1)); %policy function
        end
            
        
            
        
        
    end 
    diststore(it,1)=dist; %store distance
    distpolstore(it,1)=distpol;
    %pause(.01)
   
    %plot(1:1:it,distpolstore,1:1:it,ones(it,1)*eps*(1-beta))
    it=it+1; %update counter
    dist=max(abs(v0-vnew)); %calculate distance between value functions
    distpol=max(abs(kprime-kprime0)); %distance between policy functions
    kprime0=kprime; %update policy function
    v0=vnew; %update value function and repeat steps
end
toc
%plot the policy function below
ktrue=alpha*beta*kgrid.^alpha;
ct=kgrid'.^alpha+(1-delta).*kgrid'-kprime; %optimal consumption based on optimal kt+1
cttrue=kgrid.^alpha+(1-delta).*kgrid-ktrue; %true consumption
figure
plot(kgrid,ct,kgrid,cttrue)
title('Optimal Policy for c_{t}')
xlabel('k_t')
ylabel('c_{t}')
legend('Approx c_{t}','True c_{t}')
figure
plot(kgrid,kprime,kgrid,ktrue) 
title('Optimal Policy for k_{t+1}')
xlabel('k_t')
ylabel('k_{t+1}')
legend('Approx k_{t+1}','True k_{t+1}')
display('Max Abs Deviation between k true and k approx')
absdev=max(abs(ktrue'-kprime))
sim=0; 
if sim==1 
    %simulate the economy below
    T=50;
    k0=kgrid(1); %starting capital
    k0true=k0;
    gprime=kprime; %policy function for capital
    for i=1:T
        ind=find(k0==kgrid); %find where k0 is in the grid of capital
        k1=kprime(ind); %next periods capital
        ktrue=alpha*beta*k0true^alpha; %true policy evolution for capital
        kstoretrue(i,1)=k0true; %store true capital over time
        kstore(i,1)=k0; %store approximation for capital
        cstore(i,1)=k0^alpha+(1-delta)*k0-k1; %approximation for consumption
        cstoretrue(i,1)=k0true^alpha+(1-delta)*k0true-ktrue; %true consumption
        k0=k1; %update approximate capital stock
        k0true=ktrue;
    end
    
    
    %plot evolution of captial over time
    figure
    plot(1:1:T,kstore,1:1:T,kstoretrue)
    xlabel('Time')
    ylabel('k_t')
    legend('Approx k_t','True k_t')
    %plot evolution of consumption over time
    figure
    plot(1:1:T,cstore,1:1:T,cstoretrue)
    xlabel('Time')
    ylabel('c_t')
    legend('Approx c_t','True c_t')
else
end


%save vnew0 vnew