clear all;

%%% Parameters %%%
beta=0.96; %discount factor%
delta=.05; %depreciation rate%
alpha=.3; %capital share%
p=.6; %probability of remaining on a state%
P=[p 1-p; 1-p p]; %transition matrix%

zl=0.9; %low productivity state%
zh=1.1; %high productivity state%

Z=[zl zh];

%%% Non-stochastic steady state, z=1 %%%

kss=((1/alpha)*(1/beta-(1-delta)))^(1/(alpha-1));


%%% Grid %%%

incr=.01; %distance between gridpoints%

kgrid=[0.5*kss:incr:1.5*kss];
nk=length(kgrid); %number of gridpoints%


%%% Compute instantaneous utility for each state z %%%

ul=zeros(nk,nk)+NaN; %instantaneous utility if z=zl%
uh=ul; %instantaneous utility if z=zh%

for i=1:nk; %loop for k%
    for j=1:nk; %loop for k'%
    k=kgrid(i);
    kpr=kgrid(j);
    cl=zl*k^alpha+(1-delta)*k-kpr; %consumption if z=zl%
    ch=zh*k^alpha+(1-delta)*k-kpr; %consumption if z=zh%
    
    if cl>0 ul(j,i)=log(cl);
    end;
    
    if ch>0 uh(j,i)=log(ch);
    end;
    
    end;
end;


%%% Value Function Iteration %%%

Vl=zeros(1,nk); % initial guess for V(k,z=zl) % 
Vh=Vl; % initial guess for V(k,z=zh) %
dist=10;
iter=1;

while dist>exp(-5) & iter<500;
    Vprl=P(1,:)*[Vl;Vh]; % continuation utility if z=zl %
    Vprh=P(2,:)*[Vl;Vh]; % continuation utility if z=zh %
    Wl=ul+beta*repmat(Vprl',1,nk); % total utility if z=zl %
    Wh=uh+beta*repmat(Vprh',1,nk); % total utility if z=zh %
    [Vprl, ikprl]=max(Wl); %take max across k' for each k%
    [Vprh, ikprh]=max(Wh); %take max across k' for each k%
    dist=max(abs([Vprl Vprh]-[Vl Vh])); %check convergence%
    Vl=Vprl; %update V(k,z=zl)%
    Vh=Vprh; %update V(k,z=zh)%
    iter=iter+1 %iteration%
end;


%%% law of motion %%%

ikpr=[ikprl;ikprh];
kpr=[kgrid(ikprl);kgrid(ikprh)];

figure(1);
plot(kgrid,kgrid(ikprl),'b',kgrid,kgrid(ikprh),'r',kgrid,kgrid,'k:');
legend('z=zl','z=zh','45 degrees',2);
title('law of motion');


%%% Simulation %%%

nsim=1000; %number of simulations%
rand('state',0);
shock=rand(nsim,1);
iz=1; %initial state%
ik=190; %initial capital stock%

for i=1:nsim;
    ksim(i)=kgrid(ik);
    zsim(i)=Z(iz);
    ik=ikpr(iz,ik); %update k%
    izother=(iz==1)*2+(iz==2)*1;
    iz=(shock(i)>=p)*izother+(shock(i)<p)*iz; %update z%
end;

csim=zsim(1:nsim-1).*(ksim(1:nsim-1).^alpha)-ksim(2:nsim)+(1-delta)*ksim(1:nsim-1);

figure(2);
plot(ksim(100:nsim));
title('capital');

figure(3);
plot(csim(100:nsim-1));
title('consumption');