%Econ202b HW3 P3. Aiyagari problem. Anton Cheremukhin
clear all
close all
clc

% parameters from the task
beta=1/(1+0.0416);   alpha=0.36;    delta=0.08;
mu=2;   sigma=0.4;  ro=0.9; 
for sigma=0.2:0.2:0.4
    for mu=1.001:0.9995:3
        for ro=0.9:-0.3:0
% parameter for computations    used for convergence detection
epsilon = 0.000001;   eps=0.00001;
%define transition matrix
L0=(-3*sigma):(sigma):(3*sigma);
B0=[-10*sigma, -2.5*sigma;  -2.5*sigma, -1.5*sigma; -1.5*sigma, -0.5*sigma; -0.5*sigma, 0.5*sigma
0.5*sigma, 1.5*sigma;   1.5*sigma, 2.5*sigma;   2.5*sigma, 10*sigma];   
L=exp(L0);
P=zeros(7,7);
sl=sigma*sqrt(1-ro^2);
for i=1:7
    for j=1:7
        %probability of going from i to j
        P(i,j)=normcdf(B0(j,2),L0(i)*ro,sl)-normcdf(B0(j,1),L0(i)*ro,sl);
    end
end

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

% define grid as an array of values from 0 to 30
A_min = 0;      A_max = 20;     A_step = (A_max-A_min)/200;
A_grid = [A_min:A_step:A_max];
V_new = ones(size(A_grid,2),7);
%cycle over K, compute corresponding r and w, then Ea compare to K
r_min=0.005;   r_max=0.05;
while abs(r_max-r_min) > 0.0001
    r=(r_max+r_min)/2;
    K=((r + delta)/alpha)^(1/(alpha-1));
    w=(1-alpha)*K^alpha;
    fi=min(0,w*L(1)/r);

    % value function iteration algorithm, computing arrays of V_new for all A from A_grid
    V_old = ones(size(A_grid,2),7)*0.5;
    A_pf = ones(size(A_grid,2),7);
    while (max(max(abs(V_new - V_old))') > eps)
        V_old = V_new;
        for n = 1 : size(A_grid,2)
            a = A_min + A_step*(n-1);
            % Calculate maximum over all a_primes from our grid
            for i=1:7
                sum=zeros(size(A_grid));
                for j=1:7
                    sum=sum+P(i,j)*V_old(:,j)';
%                    sum=sum+P(i,j)*interp1(A_grid, V_old(:,j), A_grid, 'linear');
                end
                [V_new(n,i),A_pf(n,i)] = max(((max(zeros(size(A_grid)),(w*L(i)-r*fi)*ones(size(A_grid)) + (1+r)*a - A_grid)).^(1-mu)-1)/(1-mu) + beta*sum);
            end
        end
    end

    %find invariant distribution
    pi_old=zeros(1,size(A_grid,2));
    pi_new=ones(1,size(A_grid,2))/size(A_grid,2);
    while (abs(max(pi_new - pi_old)) > epsilon)
        pi_old = pi_new;
        pi_new=zeros(1,size(A_grid,2));
        for n = 1 : size(A_grid,2)
            for i = 1 : 7
                pi_new(1,A_pf(n,i))=pi_new(1,A_pf(n,i))+pi_old(1,n)*pr_new(1,i);
            end
        end
    end
%calculate Ea
    Ea=A_grid*pi_new'-fi;
    Va=((A_grid-Ea).^2)*pi_new';
    if Ea>K
        r_max=r;
    else
        r_min=r;
    end
end

sr=delta*alpha/(r+delta);

disp('sigma         mu        ro');
disp(num2str([sigma mu ro]));
disp('K=Ea         r            w           s             Va');
disp(num2str([K r   w   sr  Va]));
        end
    end 
end


%plot the results
%A_pf2 = A_grid(A_pf);
%V=V_new;
%z_grid=zeros(7,size(A_grid,2));
%for i=1:7
%    z_grid(i,:)=(w*L(i)-r*fi)*ones(1,size(A_grid,2))+(1+r)*A_grid;
%end

%z_pf=zeros(size(A_pf2));
%for i=1:7
%    z_pf(:,i)=(w*L(i)-r*fi)*ones(size(A_pf2,1),1)+(1+r)*A_pf2(:,i);
%end


%figure(1)
%plot(A_grid(:,:)',V(:,:),...
%                'LineWidth',2)
%title('value function of a')
%figure(2)
%plot(z_grid(1,1:50)',A_pf2(1:50,1),...
%               'LineWidth',2)
%title('policy function a of z')
%figure(5)
%plot(z_grid(1,1:50)',z_pf(1:50,1),...
%               'LineWidth',2)
%title('policy function z of z')
%figure(4)
%plot(A_grid,pi_new,...
%               'LineWidth',2)
%title('distribution of a')