
clear all
close all
warning off
clc


% PARAMETERS ==============================================================
%preferences & technology
beta     = 0.96;
psi      = 0.2;
rho      = 5;
bb       = 2.5; % controls bequest motive, donoted by b in paper (like bonds)
lambda   = 0.68212;
K        = 46;
Tmax     = 81;
F        = 0.025;
Rf       = 1.02;
mu       = 0.04;
% shocks
sig_U    = 0.15;
sig_N    = 0.1;
sig_EPS  = 0.18;
nU       = 5;
nN       = 5;
nEPS     = 5;
phi_NEPS = 0.15;
% earnings profile
earn_coefs = [-2.1700;0.1682;-0.0323;0.0020];
% housing profile
h_coefs    = [0.703998;-0.0352276;0.0007205;-0.0000049]; % truncate at 0
% mortality profile
p_t = 1 - [0.000823 0.000894 0.000962 0.001005 0.001016 0.001003 0.000983 0.000967 0.000960 0.000970 0.000994 0.001027 0.001065 0.001115 0.001154 0.001209 0.001271 0.001351 0.001460 0.001603 0.001769 0.001943 0.002120 0.002311 0.002520 0.002747 0.002989 0.003242 0.003512 0.003803 0.004118 0.004464 0.004837 0.005217 0.005591 0.005963 0.006346 0.006768 0.007261 0.007866 0.008596 0.009473 0.010450 0.011456 0.012407 0.013320 0.014299 0.015323 0.016558 0.018029 0.019723 0.021607 0.023723 0.026143 0.028892 0.031988 0.035476 0.039238 0.043382 0.047941 0.052953 0.058457 0.064494 0.071107 0.078342 0.086244 0.094861 0.104242 0.114432 0.125479 0.137427 0.150317 0.164187 0.179066 0.194979 0.211941 0.229957 0.249020 0.269112 0.290198 0.312231 1.000000]';



% STATE & CHOICE GRIDS ====================================================
nX    = 500;
nCvec = [1000 20];
nEvec = [100 20];
nGRID = 2; % positive integer
Xmin  = 0;
Cmin  = 0;
Emin  = 0;
Xmax  = 15; % see figure 2A
Cmax  = 7.5; %7.5; % see figure 2A
Emax  = 1;
%Xvec = linspace(Xmin,Xmax,nX)';
Xvec  = UNEQspace(Xmin,Xmax,nX);
Cvec  = linspace(Cmin,Cmax,nCvec(1))';
Evec  = linspace(Emin,Emax,nEvec(1))';



% SHOCK PROCESSES =========================================================
% [NEPS_shock probNEPS] = qnwnorm([nN;nEPS],[-0.5*sig_N^2;0],[sig_N.^2,sig_N*sig_EPS*phi_NEPS;sig_N*sig_EPS*phi_NEPS,sig_EPS.^2]);
% [U_shock probU]       = qnwnorm(nU,-0.5*sig_U^2,sig_U.^2);
load('shocks','NEPS_shock','probNEPS','U_shock','probU')
Uvec                  = exp( U_shock );
Ngrid                 = exp(NEPS_shock(:,1));
EPSgrid               = NEPS_shock(:,2);
prob                  = kron(probU,probNEPS);
Ugrid                 = reshape(repmat(Uvec',[nN*nEPS,1]),[nN*nU*nEPS,1]);
Ngrid                 = repmat(Ngrid,[nU,1]);
EPSgrid               = repmat(EPSgrid,[nU,1]);
nUNEPS                = nU*nN*nEPS;
Rs                    = Rf + mu + EPSgrid;



% SOLVE AGENT'S PROBLEM ===================================================

% initializations
Cpol = zeros(nX,2,Tmax+1);
Epol = zeros(nX,2,Tmax+1);
Ppol = ones(nX,Tmax+1);
V    = zeros(nX,2,Tmax+1);


for tt = 0:Tmax % loop over ages, beginning at Tmax
    t      = Tmax - tt;
    time_t = clock;
    
    % determine age t+1 income
    earn_t1 = exp([1 (t+1) (t+1)^2/10 (t+1)^3/100]*earn_coefs); % [1,1] t+1 value
    if (t+1)>K
       y_t1 = lambda*ones(nUNEPS,1);
       q_t1 = ones(nUNEPS,1);
    else
       y_t1 = Ugrid;
       q_t1 = earn_t1*Ngrid;
    end
    yq_t1 = y_t1.*q_t1;
    
    % determine age t housing share
    h_t = max( [1 t t^2 t^3]*h_coefs , 0); % [1,1]

    % fit splines to t+1 value function
    if t<Tmax
       pp1 = csapi(Xvec,V(:,1,t+2));
       pp2 = csapi(Xvec,V(:,2,t+2));
    end
    
    % set up some temporary objects (required for parallelization)
    Ppol_t = zeros(nX,1);
    Cpol_t = zeros(nX,2);
    Epol_t = zeros(nX,2);
    V_t    = zeros(nX,2);
    
    % invest in stocks ====================================================
    parfor ix=1:nX
        for ip=1:2 % ip=2 denotes cost have been paid in the past
            ix_c = 0;
            ix_e = 0;
            for ig=1:nGRID % iterate over finer & finer choice grids
                % set up grids for policy choices
                nE = nEvec(ig);
                nC = nCvec(ig);
                if ig==1
                   Cvec_2 = Cvec; 
                   Evec_2 = Evec;
                else
                   Cvec_2 = linspace(Cvec_2(max(ix_c-1,1)),Cvec_2(min(ix_c+1,nCvec(ig-1))),nC)';
                   Evec_2 = linspace(Evec_2(max(ix_e-1,1)),Evec_2(min(ix_e+1,nEvec(ig-1))),nE)';
                end
                % compute implied bond & stock position
                %Cvec_2(Cvec_2>Xvec(ix)) = []; nC = length(Cvec_2);
                s_2            = Evec_2*ones(1,nC) .* (Xvec(ix)-ones(nE,1)*Cvec_2'); % [nE,nC]
                bRf_2          = Rf*((1-Evec_2)*ones(1,nC)) .* (Xvec(ix)-ones(nE,1)*Cvec_2'); % [nE,nC]
                s_2(s_2<0)     = NaN;
                bRf_2(bRf_2<0) = NaN;
                % expected continuation value
                EV = zeros(nE,nC);
                for is=1:nUNEPS  
                    qx_p = s_2*Rs(is) + bRf_2 + (1-h_t)*yq_t1(is) - (ip==1)*F*q_t1(is); % [nE,nC] check how much extrapolation is required for these xp-values !!!        
                    if t==Tmax
                       VP = 1;
                    else
                       x_p = qx_p(:)/q_t1(is);
                       VP  = reshape( fnval(pp2,x_p) , [nE,nC]);
                    end
                    EV = EV + prob(is) * ( p_t(t+1)*(VP*q_t1(is)).^(1-rho) + (1-p_t(t+1))*bb*(qx_p/bb).^(1-rho) ); 
                end
                % perform maximization step
                max_this = ( (1-beta*p_t(t+1))*(ones(nE,1)*Cvec_2').^(1-1/psi) + beta*EV.^((1-1/psi)/(1-rho)) ).^(1/(1-1/psi)); % [nE,nC] 
                max_this(imag(max_this)~=0) = NaN;
                [max_val max_arg] = max(max_this); % max over E dimension for each C
                [V_t(ix,ip) ix_c] = max(max_val);
                ix_e              = max_arg(ix_c);
                Cpol_t(ix,ip)     = Cvec_2(ix_c);
                Epol_t(ix,ip)     = Evec_2(ix_e); 
            end
        end
    end

    
    % don't invest in stocks ==============================================
    parfor ix=1:nX
        ix_c = 0;
        ix_e = 0;
        for ig=1:nGRID % iterate over finer & finer choice grids
            % set up finer grids for policy choices
            nC = nCvec(ig);
            if ig==1
               Cvec_2 = Cvec;
            else
               Cvec_2 = linspace(Cvec_2(max(ix_c-1,1)),Cvec_2(min(ix_c+1,nCvec(ig-1))),nC)';
            end
            %Cvec_2(Cvec_2>Xvec(ix)) = []; nC = length(Cvec_2);
            % compute implied bond & stock position
            bRf_2          = Rf*(Xvec(ix)-Cvec_2); % [nC,1]
            bRf_2(bRf_2<0) = NaN;
            % expected continuation value
            EV = zeros(nC,1);
            for is=1:nUNEPS  
                qx_p = bRf_2 + (1-h_t)*yq_t1(is);      
                if t==Tmax
                   VP = 1;
                else
                   x_p = qx_p/q_t1(is);
                   VP  = fnval(pp1,x_p);
                end
                EV = EV + prob(is) * ( p_t(t+1)*(VP*q_t1(is)).^(1-rho) + (1-p_t(t+1))*bb*(qx_p/bb).^(1-rho) ); 
            end
            % perform maximization step
            max_this        = ( (1-beta*p_t(t+1))*Cvec_2.^(1-1/psi) + beta*EV.^((1-1/psi)/(1-rho)) ).^(1/(1-1/psi)); % [nC,1]
            max_this(imag(max_this)~=0) = NaN;
            [max_val ix_c]  = max(max_this); 
            if max_val > V_t(ix,1) % it's better not to pay the fixed cost now
               V_t(ix,1)    = max_val;
               Epol_t(ix,1) = 0;
               Cpol_t(ix,1) = Cvec_2(ix_c);
               Ppol_t(ix)   = 0;
            end  
%             plot(Cvec,max_this); pause
        end
    end
    
    % collect results of parfor loops
    Ppol(:,t+1)   = Ppol_t;
    Cpol(:,:,t+1) = Cpol_t;
    Epol(:,:,t+1) = Epol_t;
    V(:,:,t+1)    = V_t;
    
    % print iteration summary
    fprintf('age = %2d , time = %4.2fs\n',t,etime(clock,time_t));
    subplot(1,3,1); plot(Xvec,V(:,:,t+1)); xlabel('x'); title(['V',num2str(t)])
    subplot(1,3,2); plot(Xvec,Cpol(:,:,t+1)); xlabel('x'); title(['Cpol',num2str(t)])
    subplot(1,3,3); plot(Xvec,Epol(:,:,t+1)); xlabel('x'); title(['Epol',num2str(t)])
    pause(0.5)
    
end


for i = 1:82;
    gen = 83-i;
    disp(gen)
    subplot(1,3,1); plot(Xvec,V(:,:,gen)); xlabel('x'); title(['V',num2str(gen-1)])
    subplot(1,3,2); plot(Xvec,Cpol(:,:,gen)); xlabel('x'); title(['Cpol',num2str(gen-1)])
    subplot(1,3,3); plot(Xvec,Epol(:,:,gen)); xlabel('x'); title(['Epol',num2str(gen-1)])
    pause
end

%save('output2','Cpol','Epol','Ppol','V')






