%%%%% Problem set 4, Question 2 Discrete-state value function iteration for
%%%%% stochastic growth model

clear all

tic

%%%%% Parameter values

alpha = 0.33;
beta  = 0.95; 
delta = 0.04; 
sigma = 2.00; 

%%%%% Technology
zl = 0.97;
zh = 1.03; 

%%%%% Transition matrix
P = [0.95,0.05;0.05,0.95];

%%%%% Solve for kbarh 
kbarl  = ((alpha*zl*beta)/(1-beta+delta*beta))^(1/(1-alpha))
kbarh  = ((alpha*zh*beta)/(1-beta+delta*beta))^(1/(1-alpha))

%%%%% Set up discrete state space

NK    = 1000;
Kgrid = linspace(0.01,5*kbarh,NK)';

%%%%% Set up return matrices

Rl  = zeros(NK,NK); 
Rh  = zeros(NK,NK);

Cl  = zeros(NK,NK); 
Ch  = zeros(NK,NK);

for i = 1:NK,
    ki = Kgrid(i);
    for h = 1:NK,
        kh = Kgrid(h);
        
        cl  = zl*ki^alpha+(1-delta)*zl*ki-kh;
        ch  = zh*ki^alpha+(1-delta)*zh*ki-kh;
        
        Cl(i,h) = max(0,cl);  
        Ch(i,h) = max(0,ch); %% Imposes consumption non-negative
    end
end

Rl  = (1/(1-sigma))*(Cl.^(1-sigma)); 
Rh  = (1/(1-sigma))*(Ch.^(1-sigma));


%%%%% Begin value function iteration
iter  = 0;
error = 100;

V0 = zeros(NK,2);

while error>10^-6 & iter<400

Vl   = V0(:,1);
Vh   = V0(:,2);
    
TVl0 = max((Rl+beta*P(1,1)*ones(NK,1)*Vl'+beta*P(1,2)*ones(NK,1)*Vh')');
TVh0 = max((Rh+beta*P(2,1)*ones(NK,1)*Vl'+beta*P(2,2)*ones(NK,1)*Vh')');    %% Bellman operations

TV0 = [TVl0',TVh0'];

error = max(max(abs(V0-TV0))) 

V0    = TV0;   

iter  = iter+1

end

V    = TV0; 
time = toc/60;

if error<10^6 & iter<300
display('Congratulations, you found a fixed point!')
display(['Time taken = ' ,num2str(time),'   minutes'])
end

figure(1)
plot(Kgrid,V(:,1),'b',Kgrid,V(:,2),'r')
legend('V(k,zL)','V(k,zH)',2)
title('Value functions')
xlabel('Capital stock')
ylabel('Lifetime utility')

%%%%% Compute policy function

display('Now computing policy functions')

gl = zeros(NK,1);
gh = zeros(NK,1);

[vl,il] = max((Rl+beta*P(1,1)*ones(NK,1)*V(:,1)'+beta*P(1,2)*ones(NK,1)*V(:,2)')');
[vh,ih] = max((Rh+beta*P(1,1)*ones(NK,1)*V(:,1)'+beta*P(1,2)*ones(NK,1)*V(:,2)')');

for i = 1:NK,

gl(i)  = Kgrid(il(i));
gh(i)  = Kgrid(ih(i));

end

figure(2)
plot(Kgrid,Kgrid,'k--',Kgrid,gl,'b',Kgrid,gh,'r')
legend('45-degree','k_{t+1}=g(k_{t},zL)','k_{t+1}=g(k_{t},zH)',2)
title('Policy functions')
xlabel('Capital stock')
ylabel('Optimal policy')

%break

%%%%% Compute transition matrix for (k,z) pairs

display('Now computing transition matrix')

Igl = zeros(NK,NK); %% indicator functions
Igh = zeros(NK,NK);

for i = 1:NK,
    Igl(i,il(i)) = 1; %% if z=zl,k=ki, then Igl(i) = 1 for best response, zero o/w
    Igh(i,ih(i)) = 1; %% if z=zh,k=ki, then Igh(i) = 1 for best response, zero o/w
end    

QLL = P(1,1)*Igl;
QHL = P(1,2)*Igl;
QLH = P(2,1)*Igh;
QHH = P(2,2)*Igh;

Q   = [QLL,QHL;QLH,QHH];

display('Now computing stationary distribution...')

[VV,DD] = eig(Q');
dd      = diag(DD); %% check how many unit eigenvalues there are!!
                    %% if kmin = 0, there are 2 (one trivial, only relevant if initial capital is zero)
                    %% if kmin > 0, there is  1 
    

[smallest,index] = min(abs(dd-1)); %% find a unit eigenvalue
vv      = VV(:,index);

Qbar  = (vv/sum(vv))';

Qbarl = Qbar(1:NK);
Qbarh = Qbar(NK+1:2*NK);

display('...done!')

figure(3)
plot(Kgrid,Qbarl,'b',Kgrid,Qbarh,'r')
legend('conditional on zL','conditional on zH',2)
title('Stationary distribution')
xlabel('Capital stock')
ylabel('Probability mass')

figure(4)
plot(Kgrid,cumsum(Qbarl),'b',Kgrid,cumsum(Qbarh),'r')
legend('conditional on zL','conditional on zH',2)
title('Cumulative stationary distribution')
xlabel('Capital stock')
ylabel('Cumulative probability')