% script para resolucao do ex 6 da lista 1

% parametros do modelo

alfa = 0.3;
beta = 0.95;
gama = 2;
delta=0.1;

thetas = [0.75,1,1.25]; % estados dos choques de consumo
num_estados = length(thetas);
P = [0.15,0.7,0.15;0.1,0.8,0.1;0.15,0.7,0.15]; % matriz de transição


%P = [0.8,0.1,0.1;0.1,0.8,0.1;0.1,0.1,0.8];


% calculo do capital em estado estacionário

kss= ((1/beta-(1-delta))/alfa)^(1/(alfa-1));
css = kss^alfa-delta*kss;
iss = kss^alfa-css;
%%
% cria grid de capital 
espessura = 5e-3;
kgrid = (kss-400*espessura):espessura:(kss+400*espessura);
kprox = kgrid;
nk = length(kgrid);

% matriz 3d que guardara as utilidades para todas as combinações possiveis
% de k,k' e theta.
% primeira dimensão: valores possiveis dos estados theta
% segunda dimensão: k'
% terceira dimensão: k
u=zeros(num_estados,nk,nk)+NaN;

for i=1:nk
    for j=1:nk
        c=kgrid(i)^alfa+(1-delta)*kgrid(i)-kprox(j);
        if c>=0
            u(:,j,i) = thetas.*(c^(1-gama)-1)/(1-gama);
        end
    end
end
                       
%% solucao numerica para funcao valor

% cria matriz que guardara as funções valores. 
% cada linha da matriz corresponde a um estado diferente
% Chute inicial = 0;
V = zeros(num_estados,nk); 
dist = ones(1,num_estados)*10;
iter=1;
ikprox=zeros(num_estados,nk)+NaN;
while max(dist)>exp(-5) && iter<500
    for i=1:num_estados
        %calculo da esperanca da função valor dado estado atual:
        esperanca_v = P(i,:)*V;
        W=squeeze(u(i,:,:))+beta*repmat(esperanca_v',1,nk); %utilidade total para todas as combinções de k e k'
        [Vpr, ikprox(i,:)]=max(W); % toma o maximo pra cada coluna em relação as linhas
        dist(i)=max(abs(Vpr-V(i,:))); % diferença elemento a elemento do V anterior com V atual
        V(i,:)=Vpr; %update V(k)%
    end
    iter=iter+1;
end

%% plot das funcoes valores pros 3 estados
figure(1)
plot(kgrid,V(1,:),'b','DisplayName','Função Valor para theta = 0.75')
hold on
plot(kgrid,V(2,:),'r','DisplayName','Função Valor para theta = 1')
plot(kgrid,V(3,:),'g','DisplayName','Função Valor para theta = 1.25')
legend
title('Funções Valores')
hold off
%% graficos da policy function

figure(2)
plot(kgrid,kgrid(ikprox(1,:)),'b','DisplayName','Policy p/ Theta = 0.75')
hold on
plot(kgrid,kgrid(ikprox(2,:)),'r','DisplayName','Policy p/ Theta= 1')
plot(kgrid,kgrid(ikprox(3,:)),'g','DisplayName','Policy p/ Theta = 1.25')
plot(kgrid,kgrid,'k--','DisplayName','Reta 45º')
ylabel('k''=\phi(k)');
xlabel('k');
legend
title('Policy Function')
hold off

%% Simulando a economia por 1000 períodos 
T = 1e3; % tempo de simulação 
estado_inicial = 2; %indice do estado incial da economia 
seed = 1; 
%simula uma cadeia de markov com T realizações, com estados thetas e matriz de transição P
[chain,state] = markovsimul(P,thetas,T,estado_inicial,seed); 

k_simulacao = zeros(1,T);

k_simulacao(1)=kgrid(1); % capital incial k0

for i =2:T
    estado_ontem = state(i-1);
    indice_capital_passado = find(kgrid==k_simulacao(i-1));
    k_simulacao(i)=kgrid(ikprox(estado_ontem,indice_capital_passado));
    
end

% consumption path:
csim=k_simulacao(1:T-1).^alfa+(1-delta)*k_simulacao(1:T-1)-k_simulacao(2:T); %simulated consumption%
isim = k_simulacao(1:T-1).^alfa-csim;
% calculo da correlação entre consumo e investimento
% interesse no elemento 1,2 da matriz (correlação de c com i)
matriz_correlacao = corrcoef(csim(100:end),isim(100:end));
disp(matriz_correlacao(1,2))
%% plot dos trajetórias das variáveis
figure(3)
subplot(4,1,1)
plot(k_simulacao)
ylabel('Capital')

subplot(4,1,2)
plot(csim)
ylabel('Consumo')

subplot(4,1,3)
plot(isim)
ylabel('Investimento')

subplot(4,1,4)
plot(chain)
ylim([0.7 1.3])
ylabel('Choques')