%Performs simple value function iteration for Ramsey model.  Written by
%Philip Shaw, Fordham University.  

clear all;
tic;
alpha=.36;                   %Production function power
beta=.98;    %Discount factor
kstar=(alpha*beta)^(1/(1-alpha)); %steady state value
step=.01; %distance the grid takes between values
kgridi=[.1:step:.30]';         %Discretizing the state space (grid for k)
kgridj=[.1:step:.30]';         %Grid for k prime
dif=1; %difference between value function iterations
dif1=1; %difference between policy function iterations
v0=ones(length(kgridj),1)*(log(kstar^(alpha)-kstar)/(1-beta));  %Initial value for the bellman equation
g0=zeros(length(kgridi),1); %policy function guess
epsilon=.0001;    %Stop criterion 
it=1; %initial value for counter
while abs(dif1)>epsilon       %Loop for value function iteration
    for i=1:length(kgridi);
        for j=1:length(kgridj);
            v(i,j)=log(kgridi(i)^(alpha)-kgridj(j))+beta*v0(j);
        end
        v1(i)=max(v(i,:));
        jpoint(i)=find(v(i,:)==v1(i));
        g(i,1)=kgridj(jpoint(i)); %policy function
    end
    dif1(it)=max(abs(g-g0)); %difference between guess on policy function and new policy function
    dif(it)=max(abs(v1'-v0)); %difference between guess on value function and new value function
    g0=g; %update policy function
    v0=v1'; %updating the value function
    it=it+1;
end; 
for ij=1:length(kgridi)
    kprime(ij)=alpha*beta*kgridi(ij)^alpha;
end
figure;plot(kgridi,g,kgridi,kprime)
legend('Approximate Policy Function','True Policy Function')
title('Approximate Policy Function vs True Policy Function')
xlabel('K')
ylabel('Kprime')



figure;plot(1:1:it-1,dif,1:1:it-1,dif1)
legend('Difference on Value Function','Difference on Policy Function')
title('Policy Function Iteration vs Value Function Iteration')
xlabel('iterations')
ylabel('Difference')

%finding the evolution of capital stock over time
k0=.1; %number of apples we start with
kinit=k0;
ik=1; %counter for finding capital stock evolution
difk=1; %distance between time t and t+1 capital stock

while difk>0
    KPRIME=g(find(kgridi==k0)); %finding kprime for k
    difk=KPRIME-k0; %distance between kprime and k over time
    k0=KPRIME; %updating k0
    capital(ik)=k0; %stores the evolution of capital stock over time
    ik=ik+1; %update our counter
    
end
figure;plot(1:1:ik,[kinit capital])
legend('Capital Stock Over Time')
title('Evolution of Capital Stock')
xlabel('Time')
ylabel('Capital')

toc