%Econ202c HW1 P1. Anton Cheremukhin
clear all;
close all;
clc;
% parameters from the task
beta=25/27;  sigma=1.01;  delta=247/2500; A=1341/2500;   teta=1/3;   gamma=1.02; ksi=1.015;
% parameter for computations    used for convergence detection
epsilon = 0.00000001;
eps=0.0000000001;
%define transition matrix
alpha=0.75:0.25:1.25;
P=[0.15 0.7 0.15        %case (4)
   0.1  0.8 0.1
   0.15 0.7 0.15];

%P=[0.8  0.1 0.1        %case(5)
%   0.1  0.8 0.1
%   0.1  0.1 0.8];

%P=[0.95  0.025 0.025        %case(6)
%   0.025  0.95 0.025
%   0.025  0.025 0.95];

P=ones(3,3)/3;              %case (7)


%invariant distribution of states of nature
pr_new=ones(1,3)/3; pr_old=zeros(1,3);
while (abs(max(pr_new - pr_old)) > eps)
    pr_old = pr_new;
    pr_new=zeros(1,3);
    for i = 1 : 3
        for j = 1 : 3
            pr_new(1,j)=pr_new(1,j)+pr_old(1,i)*P(i,j);
        end
    end
end
%norm L to 1 on average

% define grid as an array of values from 0 to 30
k_min = 1-100*0.0025;      k_max = 1+100*0.0025;     k_step = (k_max-k_min)/200;
k_grid = [k_min:k_step:k_max];      N=201;
V_new = ones(size(k_grid,2),3);
% value function iteration algorithm, computing arrays of V_new for all A from A_grid
    V_old = ones(size(k_grid,2),3)*0.5;
    k_pf = ones(size(k_grid,2),3);
    while (max(max(abs(V_new - V_old))') > epsilon )
        abs(max(max(abs(V_new - V_old))'));
        V_old = V_new;
        for n = 1 : size(k_grid,2)
            k = k_min + k_step*(n-1);
            % Calculate maximum over all a_primes from our grid
            for i=1:3
                sum=zeros(size(k_grid));
                for j=1:3
                    sum=sum+P(i,j)*V_old(:,j)';
                end
                %[V_new(n,i),k_pf(n,i)] = max(alpha(i)*(max(zeros(size(k_grid)),( A*k^teta + (1-delta)*k - ksi*gamma*k_grid)).^(1-sigma)-1)/(1-sigma) + beta*ksi*gamma*sum);
                [V_new(n,i),k_pf(n,i)] = max(alpha(i)*log(max(ones(size(k_grid))/10^7,( A*k^teta + (1-delta)*k - ksi*gamma*k_grid))) + beta*ksi*gamma*sum);
            end
        end
    end

%find invariant distribution
   pi_old=zeros(1,size(k_grid,2));
   pi_new=ones(1,size(k_grid,2))/size(k_grid,2);
   while (abs(max(pi_new - pi_old)) > eps)
       pi_old = pi_new;
       pi_new=zeros(1,size(k_grid,2));
       for n = 1 : size(k_grid,2)
           for i = 1 : 3
               pi_new(1,k_pf(n,i))=pi_new(1,k_pf(n,i))+pi_old(1,n)*pr_new(1,i);
           end
       end
   end
% compute first and second moments
   Ek=k_grid*pi_new';
   Vk=((k_grid-Ek*ones(size(k_grid))).^2)*pi_new';

   i_grid=ksi*gamma*k_grid(k_pf(:,:))-(1-delta)*k_grid'*ones(1,3);
   Ei=pr_new*(i_grid'*pi_new');
   Vi=pr_new*(((i_grid-Ei*ones(size(i_grid))).^2)'*pi_new');
   
   c_grid=A*(k_grid'*ones(1,3)).^teta-i_grid;
   Ec=pr_new*(c_grid'*pi_new');
   Vc=pr_new*(((c_grid-Ec*ones(size(c_grid))).^2)'*pi_new');
  
   y_grid=c_grid+i_grid;
   Ey=pr_new*(y_grid'*pi_new');
   Vy=pr_new*(((y_grid-Ey*ones(size(y_grid))).^2)'*pi_new');

   k_grid3=k_grid'*ones(1,3);
   Cov_ic=pr_new*(((i_grid-Ei*ones(size(i_grid))).*(c_grid-Ec*ones(size(c_grid))))'*pi_new');
   Cov_iy=pr_new*(((i_grid-Ei*ones(size(i_grid))).*(y_grid-Ey*ones(size(y_grid))))'*pi_new');
   Cov_ik=pr_new*(((i_grid-Ei*ones(size(i_grid))).*(k_grid3-Ek*ones(size(k_grid3))))'*pi_new');
   Cov_cy=pr_new*(((c_grid-Ec*ones(size(c_grid))).*(y_grid-Ey*ones(size(y_grid))))'*pi_new');
   Cov_ck=pr_new*(((c_grid-Ec*ones(size(c_grid))).*(k_grid3-Ek*ones(size(k_grid3))))'*pi_new');
   Cov_ky=pr_new*(((k_grid3-Ek*ones(size(k_grid3))).*(y_grid-Ey*ones(size(y_grid))))'*pi_new');
   
   %     i             c                y                   k
   Corr=[1  Cov_ic/sqrt(Vc*Vi)  Cov_iy/sqrt(Vy*Vi)  Cov_ik/sqrt(Vk*Vi)        %i
        Cov_ic/sqrt(Vc*Vi)  1   Cov_cy/sqrt(Vy*Vc)  Cov_ck/sqrt(Vk*Vc)        %c
        Cov_iy/sqrt(Vy*Vi)  Cov_cy/sqrt(Vy*Vc)  1   Cov_ky/sqrt(Vk*Vy)        %y
        Cov_ik/sqrt(Vk*Vi)  Cov_ck/sqrt(Vc*Vk)   Cov_ky/sqrt(Vk*Vy)  1        %k
   ]
   FM=[Ek   Vk  Ec  Vc  Ei  Vi  Ey  Vy]

%same using monte-carlo simulations
Nb=10000;
k=ones(1,Nb+1);   c=ones(1,Nb);
in=ones(1,Nb);  y=ones(1,Nb);
s=2*ones(1,Nb);
for i=1:Nb
    rn=rand(1,1);
    if rn<P(s(i),1)
        s(i+1)=1;
    elseif rn>1-P(s(i),3)
        s(i+1)=3;
    else
        s(i+1)=2;
    end
    j=interp1(k_grid,1:size(k_grid,2),k(i));
    k(i+1)=k_grid(k_pf(int8(j),s(i)));
    y(i)=A*(k(i))^teta;
    in(i)=ksi*gamma*k(i+1)-(1-delta)*k(i);
    c(i)=y(i)-in(i);
    if c(i)<0
        c(i)=0;
        in(i)=y(i);
        k(i+1)=(in(i)+(1-delta)*k(i))/(ksi*gamma);
    end
end

c=c(101:Nb);    y=y(101:Nb);    in=in(101:Nb);    k=k(101:Nb);

Eco=(c*ones(size(c))')/size(c,2);
Vco=(((c-Eco).^2)*ones(size(c))')/size(c,2);
Ein=(in*ones(size(in))')/size(in,2);
Vin=(((in-Ein).^2)*ones(size(in))')/size(in,2);
Eyp=(y*ones(size(y))')/size(y,2);
Vyp=(((y-Eyp).^2)*ones(size(y))')/size(y,2);
Eca=(k*ones(size(k))')/size(k,2);
Vca=(((k-Eca).^2)*ones(size(k))')/size(k,2);

Covinco=(((in-Ein).*(c-Eco))*ones(size(in))')/size(in,2);
Corrinco=Covinco/sqrt(Vco*Vin);
Covypco=(((y-Eyp).*(c-Eco))*ones(size(in))')/size(in,2);
Corrypco=Covypco/sqrt(Vco*Vyp);
Covypin=(((y-Eyp).*(in-Ein))*ones(size(in))')/size(in,2);
Corrypin=Covypin/sqrt(Vin*Vyp);

Corr2=[1  Corrinco  Corrypin  Corrypin        %i
        Corrinco  1   Corrypco  Corrypco        %c
        Corrypin  Corrypco  1   1        %y
        Corrypin  Corrypco  1  1        %k
   ]
FM2=[Eca   Vca  Eco  Vco  Ein  Vin  Eyp  Vyp]

%plot the results
k_pf2 = k_grid(k_pf);

figure(1)
plot(k_grid',V_new(:,1),k_grid',V_new(:,2),k_grid',V_new(:,3),...
                'LineWidth',2)
title('value function of k')
figure(2)
plot(k_grid,k_pf2(:,1),k_grid,k_pf2(:,2),k_grid,k_pf2(:,3),...
               'LineWidth',2)
title('policy function k of k')

figure(3)
plot(k_grid,pi_new,...
               'LineWidth',2)
title('distribution of k')
figure(4)
plot(1:300,y(1:300),1:300,c(1:300),1:300,in(1:300),...
               'LineWidth',2)
legend('output','consumption','investment')