%Econ202b HW3 P2. Lucas problem. Anton Cheremukhin
clear;
% parameters from the task
beta=0.9;  mu=3;
% parameter for computations    used for convergence detection
epsilon = 0.0000001;
%define transition matrix
Z=0.5:0.5:1.5;
P=ones(3,3)/3;

%invariant distribution of states of nature
pr_new=ones(1,3)/3; pr_old=zeros(1,3);
while (abs(max(pr_new - pr_old)) > epsilon)
    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
M_min = 1;      M_max = 5;     M_step = (M_max-M_min)/200;   y=1;
M_grid = [M_min:M_step:M_max];
V_new = ones(size(M_grid,2),3);
    % value function iteration algorithm, computing arrays of V_new for all A from A_grid
    V_old = ones(size(M_grid,2),3)*0.5;
    M_pf = ones(size(M_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(M_grid,2)
            m = M_min + M_step*(n-1);
            % Calculate maximum over all a_primes from our grid
            for i=1:3
                sum=zeros(size(M_grid));
                for j=1:3
                    sum=sum+P(i,j)*V_old(:,j)';
                end
                [V_new(n,i),M_pf(n,i)] = max((Z(i)*(max(zeros(size(M_grid)),(m+y)*ones(size(M_grid)) - M_grid)).^(1-mu)-1)/(1-mu) + beta*sum);
            end
        end
    end

    %find invariant distribution
    pi_old=zeros(1,size(M_grid,2));
    pi_new=ones(1,size(M_grid,2))/size(M_grid,2);
    while (abs(max(pi_new - pi_old)) > epsilon)
        pi_old = pi_new;
        pi_new=zeros(1,size(M_grid,2));
        for n = 1 : size(M_grid,2)
            for i = 1 : 3
                pi_new(1,M_pf(n,i))=pi_new(1,M_pf(n,i))+pi_old(1,n)*pr_new(1,i);
            end
        end
    end
    Em=M_grid*pi_new'
    Vm=((M_grid-Em*ones(size(M_grid))).^2)*pi_new'
    

%plot the results
%M_pf = M_grid(M_pf);

figure(1)
plot(M_grid',V_new(:,1),M_grid',V_new(:,2),M_grid',V_new(:,3),...
                'LineWidth',2)
title('value function of m')
figure(2)
plot(M_grid,M_pf(:,1),M_grid,M_pf(:,2),M_grid,M_pf(:,3),...
               'LineWidth',2)
title('policy function m_ of m')
figure(3)
plot(M_grid,pi_new,...
               'LineWidth',2)
title('distribution of m')