%Solving Bewley model using policy function iteration with Chebyshev
%polinomials and solving for equilibrium probability distribution over a
%discretized state space.

%Calling:   TRP_AR.m to calculate Markov transition matrix from AR(1)
                %process

clc;
clear all;

%Switch
switch_figure   =   1;

%Setting parameters
beta       =   0.951;
alpha      =   0.33;
delta      =   0.06;
b          =   0;
r_init     =   0.0475; %0.02
a_max      =   100;
rho     =  0; %0.95    %persistence of lny
sigma   =  0.2;    %conditional standard deviation of lny
gamma   =  2;

%Setting iteration parameters
eta         =  0.01;            %Reducing linesearch stepsize for r iteration
maxiter_PFI =  1000;            %Maximum iteration for policy function iteration
conv        =  1e-6;            %Convergence criterion for the policy function iteration
conv_r      =  1e-4;            %Convergence criterion for calculating r;
n_y         =  15;              %number of y grids
n_a         =  15;              %number of asset grids
p_a         =  5;               %order of Chebyshev polinomial (p_a<n_a)
n_a_GE      =  200;             %Number of asset grids used for GE

%Creating transition matrix for an AR(1) process
ss      =  3; %2       %multiplier for the variance to determine the width of the grid

%Creating y grid
lny_min     =   -ss*sigma/(sqrt(1-rho^2));
lny_max     =   ss*sigma/(sqrt(1-rho^2));
lny_inc     =   (lny_max-lny_min)/(n_y-1);
grid_lny    =   lny_min:lny_inc:lny_max;
y           =   exp(grid_lny);

%Transition matrix
P_y         =   TRP_AR(grid_lny,0,rho,sigma^2);

%Creating grid over a using Chebyshev grids
a_max_GE   =   a_max;
a_min      =   b;
z           =   -cos((2*(1:n_a)-1)/(2*n_a)*pi);
grid_a      =   (z+1)*(a_max-a_min)/2+a_min;

%iterating to get A matrix of Chebyshev weights for each y
%initializing A matrix
A_init      =   zeros(p_a,n_y);
A_init(1,:) =   1;

%Matrix of Chebyshev polinomials
Phi_C       =   cos(acos(z')*(0:p_a-1));

%Initializing the search for r
r1          =   r_init;
r           =   0;   

%Initial Chebyshev coefficients
A1          =   A_init;
A0          =   zeros(p_a,n_y);

%Starting the iteration for the interest rate r
iter=0;
while abs(r1-r)>conv_r
    iter       =   iter+1;
    r          =   r1;
    w          =   (1-alpha)*((r+delta)/alpha)^(alpha/(alpha-1));
    A0         =   zeros(p_a,n_y);
    C0         =   Phi_C*A0;
    C1         =   Phi_C*A1;

    i=0;
    fprintf('Iteration #%2.0f\n',iter);
    display('Calculating policy function...');
    while and(max(max(abs(C1-C0)))>conv,i<maxiter_PFI)
        i=i+1;
        A0      =   A1;                                             %p_a x n_y matrix of Chebyshev coefficients
        for j=1:n_y                                                 
            C0(:,j) =   Phi_C*A0(:,j);                                  %n_a x 1 vector of unconstrained current consumption
            ap      =   max((1+r)*grid_a'+w*ones(n_a,1)*y(j)-C0(:,j),b);                      %n_a x 1 vector of next period assets
            Phi_Cp  =   cos(acos(2*(ap-a_min)/(a_max-a_min)-1)*(0:p_a-1));        %n_a x p_a matrix of Chebyshev points over ap
            cp      =   min(Phi_Cp*A0,(1+r)*ap*ones(1,n_y)+w*ones(n_a,1)*y-b);   %n_a x n_y matrix of cprime
            Y       =   (beta*(1+r)*(cp.^(-gamma))*P_y(j,:)').^(-(1/gamma));             %n_a matrix of expected marginal utilities
            A1(:,j) =   inv(Phi_C'*Phi_C)*Phi_C'*Y;
        end;    
    C1          =   Phi_C*A1;
    %max(max(abs(C1-C0)))
    end;
    if (i>=maxiter_PFI); display('Warning: No PFI convergence!'); end;

    %Calculating consumption function

    if switch_figure==1
        Cons        =   min(Phi_C*A1,(1+r)*grid_a'*ones(1,n_y)+w*ones(n_a,1)*y-b);
        figure(1);
        plot(grid_a,Cons(:,1),grid_a,Cons(:,round(n_y/2)),grid_a,Cons(:,n_y));
    end;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%Calculate general equilibrium
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%Create grid_a discrete logarithmic grid for a
    lna_inc_GE  =   (log(a_max_GE+1)-log(a_min+1))/(n_a_GE-1);
    grid_lna_GE =   log(a_min+1):lna_inc_GE:log(a_max_GE+1);
    grid_a_GE   =   a_min+exp(grid_lna_GE)-1;                                               %a grid
    Phi_C_GE    =   cos(acos(2*(grid_a_GE'-a_min)/(a_max_GE-a_min)-1)*(0:(p_a-1)));         %Chebyshev polinomial matrix
    c_GE        =   min(Phi_C_GE*A1,(1+r)*grid_a_GE'*ones(1,n_y)+w*ones(n_a_GE,1)*y-b);     %n_a_GE x n_y matrix of consumption
    ap_GE       =   (1+r)*grid_a_GE'*ones(1,n_y)+w*ones(n_a_GE,1)*y-c_GE;                   %n_a_GE x n_y matrix of next period asset holdings

%Calculating the n_a_GE*n_y x n_a_GE*n_y transition matrix
    a_ind       =   repmat((1:n_a_GE)',[1,n_y,n_a_GE,n_y]);     %Index array for a;
    ap_ind      =   permute(a_ind,[3,2,1,4]);                   %Index array for aprime
    y_ind       =   repmat(1:n_y,[n_a_GE,1,n_a_GE,n_y]);        %Index array for y
    ap_lind     =   sub2ind(size(ap_GE),a_ind,y_ind);           %linear indices for aprime;
    ap_array    =   ap_GE(ap_lind);                             %array for aprime

%Calculating the transition matrix
    P_trans     =   reshape(permute(repmat(P_y,[1,1,n_a_GE,n_a_GE]),[3,1,4,2]).*...
        (ap_ind==(floor(n_a_GE*(log(ap_array-a_min+1)-log(a_min+1))/(log(a_max_GE+1)-log(a_min+1))+1))),[n_a_GE*n_y,n_a_GE*n_y]);

%Solving for the discrete steady state distribution
    lambda      =   reshape([(eye(n_a_GE*n_y)-P_trans');ones(1,n_a_GE*n_y)]\[zeros(n_a_GE*n_y,1);1],[n_a_GE,n_y]); 

%Calculating capital demand
    A           =   sum(sum(lambda.*ap_GE));                %Aggregate asset demand
    L           =   sum(sum(lambda.*(ones(n_a_GE,1)*y)));   %Aggregate labor supply

%Capital supply
    K           =   ((r+delta)/alpha)^(1/(alpha-1))*L;  %Aggregate asset supply

    F_KA        =   alpha*A^(alpha-1)*L^(1-alpha);  %Marginal product of capital at aggregate asset demand

%Calculating and writing out new r;
    r1          =   r+eta*((F_KA-delta)-r);
    fprintf('New value for r: %1.4f%%\n\n',100*r1);
end;


Y           =   K^alpha*L^(1-alpha); %output
f           =   lambda';             %density

%Calculating cdf
F           =   zeros(n_a_GE,n_y);
for i=1:n_a_GE
    for j=1:n_y
        F(i,j) =    sum(sum(lambda(1:i,1:j)));
    end;
end;

if switch_figure == 1
    figure(2);
    subplot(1,2,1); plot(grid_a_GE,c_GE); title('Consumption');
    subplot(1,2,2); plot(grid_a_GE,ap_GE); title('Asset Choice');
    figure(3);
    maxgrid     =   110;
    [A_mesh,Y_mesh] =   meshgrid(grid_a_GE(1:maxgrid),y);
    mesh(A_mesh,Y_mesh,lambda(1:maxgrid,:)'); title('Equilibrium density'); xlabel('Asset'); ylabel('Income');
    figure(4);
    [A_mesh,Y_mesh] =   meshgrid(grid_a_GE(1:110),y);
    mesh(A_mesh,Y_mesh,F(1:110,:)'); title('Equilibrium CDF'); xlabel('Asset'); ylabel('Income');

end;

save PFI_rho0.mat grid_a y grid_a_GE c_GE ap_GE f F r w L K Y;

