close all
clear
clc

%% Structural model
% This code solves for the policy functions of the model described in
% Section 5, and generates the associated statistics to be reported in the
% paper, for the different earnings processes.

% To do so, it relies on additional functions provided in mfiles

% All inputs are provided.

%% 1. Choice of parametrization

%---------- Choose earnings profile to input ---------------------------%

% choose==1 : NL process
% choose==2 : canonical process
% choose==3 : Non-normal, age dependent
% choose==4 : Normal, age dependent

% choose==5 : transitory component, robustness
% choose==6 : persistent component, robustness
% choose==7 : canonical process, robustness

tic

path(path,'mfiles')

global da dy dbet dhip TL T TR gpop a ageoldyoung agetime eff sur pfrac dye dagg azero
   
choose=3; % choice of earnings process
commonbeta=0; %do we want to use the average beta?
sig=2; % CRRA parameter
initialwealth=0; % initial wealth of agents = 0. Other numbers are options
% for empirical correspondences between initial wealth and earnings.
blimit=-0.12; % Borrowing limit.
          %This is for the grid, so it will be the lowest possible across
          % all ages. Then we will also restrict so that this does not
          % violate anyone's natural borrowing limit.

r=0.04; % Interest rate
maxa=500;
da=200;
wiratio=3.1;


bet=0.9445;
if commonbeta==1
bet=0.9390;
end

%--------- Calibration of beta -----------------------------------------%

autocalibrate=1;        %=1, calibration by bisection; 0: no calibration

if commonbeta==1
    autocalibrate=0;
end

    criterio=1; %calibration of beta: deviation of K/Y from 3
    topbet=0;
    bottombet=0;
    topcriterio=0;
    bottomcriterio=0;
    
while criterio>0.01

%------------First start some common age components---------------------%

tlength=1;
TL=36;
agetime=(25:60)';
ageoldyoung=[agetime;(61:agetime(1)+T-1)'];

% Age-efficiency profile
psid_logprofile=csvread('Inputs/age_efficiency_log.csv');    
eff2=exp(psid_logprofile)./mean(exp(psid_logprofile));
eff=eff2;
    
    dbet=1;dhip=1; % Not active anymore in this code
    

if choose==1
    dy=18;
    U=csvread('../Processes/eta_space.csv',1,0);
    U2=csvread('../Processes/eta_tranmatrices.csv',1,0);
    
    ylifelog=reshape(U(:,3),dy,TL);
    Qylife=permute(reshape(U2(:,4),dy,dy,TL-1),[2 1 3]);
    invher=[0.02*ones(5,1); 0.1*ones(8,1); 0.02*ones(5,1)];
    
    U3=csvread('../Processes/eps_space.csv',1,0);
    dye=8;
    ye=exp(reshape(U3(:,3),dye,TL));
    yedtb=repmat([0.025;0.025;0.05;0.4;0.4;0.05;0.025;0.025],1,TL);
end


if choose==2
    
    dy=18;
    dye=8;
    
    [ylifelog,Qylife]=rouwen_finiteT(1,sqrt(0.2332),sqrt(0.0060),dy,TL);
    [ye,Qye]=rouwen_finiteT(0,sqrt(0.0620),sqrt(0.0620),dye,TL) ;
    ye=exp(ye);
    
    [vec,val]=eig(Qylife(:,:,1)');  %find inv distr
    [i,~]=find((abs(val-1))<1e-3);
    invher=vec(:,i)/sum(vec(:,i));
    
    [vec,val]=eig(Qye(:,:,1)');  %find inv distr
    [i,~]=find((abs(val-1))<1e-3);
    yedtb2=vec(:,i)/sum(vec(:,i));
    
    yedtb=repmat(yedtb2,1,TL);
    
    %Note that to keep notation constant I have to drop the first matrix
    %from Qylife (only useful to define dtb at birth):
    
    Qylife=Qylife(:,:,2:end);
end


if choose==3
    dy=18;
    U=csvread('../Processes/Decompositions/r_medianstable_abb_d1.csv',1,0);
    U2=csvread('../Processes/Decompositions/r_tranmatrices_abb_d1.csv',1,0);
    
    ylifelog=reshape(U(:,3),dy,TL);
    Qylife=permute(reshape(U2(:,4),dy,dy,TL-1),[2 1 3]);
    
    invher=[0.02*ones(5,1); 0.1*ones(8,1); 0.02*ones(5,1)];
    
    U3=csvread('../Processes/Decompositions/r_transitory_abb_d1.csv',1,0);
    
    dye=8;
    ye=exp(reshape(U3(:,3),dye,TL));
    yedtb=repmat([0.025;0.025;0.05;0.4;0.4;0.05;0.025;0.025],1,TL);
end



if choose==4
    dy=18;
    U=csvread('../Processes/Decompositions/r_medianstable_abb_ko.csv',1,0);
    U2=csvread('../Processes/Decompositions/r_tranmatrices_abb_ko.csv',1,0);
    
    ylifelog=reshape(U(:,3),dy,TL);
    Qylife=permute(reshape(U2(:,4),dy,dy,TL-1),[2 1 3]);
    
    invher=[0.02*ones(5,1); 0.1*ones(8,1); 0.02*ones(5,1)];
    
    U3=csvread('../Processes/Decompositions/r_transitory_abb_ko.csv',1,0);
    
    dye=8;
    ye=exp(reshape(U3(:,3),dye,TL));
    yedtb=repmat([0.025;0.025;0.05;0.4;0.4;0.05;0.025;0.025],1,TL);
end


if choose==5
    dy=18;
    U=csvread('../Processes/eta_space.csv',1,0);
    U2=csvread('../Processes/eta_tranmatrices.csv',1,0);
    
    ylifelog=reshape(U(:,3),dy,TL);
    Qylife=permute(reshape(U2(:,4),dy,dy,TL-1),[2 1 3]);
    
    invher=[0.02*ones(5,1); 0.1*ones(8,1); 0.02*ones(5,1)];
    
    U3=csvread('../Processes/eps_space_16.csv',1,0);
    
    dye=16;
    ye=exp(reshape(U3(:,3),dye,TL));
    yedtb=repmat([0.001;0.004;0.005;0.02;0.02;0.05;0.2;0.2;0.2;0.2;0.05;0.02;0.02;...
        0.005;0.004;0.001],1,TL);
    %da=100;
    %maxa=250;
end


if choose==6
    dy=36;
    U=csvread('../Processes/eta_space_36.csv',1,0);
    U2=csvread('../Processes/eta_tranmatrices_36.csv',1,0);
    
    ylifelog=reshape(U(:,3),dy,TL);
    Qylife=permute(reshape(U2(:,4),dy,dy,TL-1),[2 1 3]);
    invher=[0.0005;0.0005;0.002;0.002;0.005;0.01;0.01;0.01;0.01;...
        0.05*ones(18,1);0.01;0.01;0.01;0.01;0.005;0.002;0.002;0.0005;0.0005];
    
    dye=8;
    U3=csvread('../Processes/eps_space.csv',1,0);
    
    ye=exp(reshape(U3(:,3),dye,TL));
    yedtb=repmat([0.025;0.025;0.05;0.4;0.4;0.05;0.025;0.025],1,TL);

end


if choose==7
    dy=18;
    U=csvread('../Processes/Robustness/canonical_eta.csv',1,0);
    U2=csvread('../Processes/Robustness/canonical_eta_tranmatrix.csv',1,0);
    
    ylifelog=reshape(U(:,3),dy,TL);
    Qylife=permute(reshape(U2(:,4),dy,dy,TL-1),[2 1 3]);
    
    invher=[0.02*ones(5,1); 0.1*ones(8,1); 0.02*ones(5,1)];    
    U3=csvread('../Processes/Robustness/canonical_eps.csv',1,0);
    
    dye=8;
    ye=exp(reshape(U3(:,3),dye,TL));
    yedtb=repmat([0.025;0.025;0.05;0.4;0.4;0.05;0.025;0.025],1,TL);
end



%% 3. Other parameters and relevant components

%Now pieces that would be useful if there were HIPs, but now they are zeros.
%Form the ylife matrix that will be the state space for earnings.

    invherhip=(1./dhip)*kron(ones(dhip,1),invher);
    
    hipage=zeros(TL,dhip);
    
    logylife=reshape(repmat(reshape(ylifelog,dy,TL),dhip,1),dy*dhip*TL,1)+...
    kron(reshape(hipage',TL*dhip,1),ones(dy,1));
    
    ylife=exp(logylife);

    
%----------- Survival probabilities ------------------------------------%

load 'Inputs/surlife'; % conditional survival probabilities from 20 to 86
if (floor(size(surlife,1)/tlength)-size(surlife,1)/tlength) ~=0
   error('modify surlife'),
end
    surlife=surlife(6:end);
T=size(surlife,1)/tlength+1; % maximum number of periods an agent lives
                             % in the model (she is born at 25)
                             
sur=reshape(surlife,tlength,T-1);
sur2=prod(sur,1)'; %adjusting for our time period;
sur=[sur2; 0];
clear eflife surlife;

% define retirement length
TR=T-TL;

% define the population growth rate
gpop=1.012^tlength;

pfrac=(cumprod([1; sur(1:T-1)]/gpop))./...
      sum(cumprod([1; sur(1:T-1)]/gpop));
  
  % Renormalize eff such that sum(eff.*pfrac(1:TL))./sum(pfrac(1:TL))=1
  
eff=eff./(sum(eff.*pfrac(1:TL))./sum(pfrac(1:TL)));
  

%% 4. Statistics of the earnings processes
normtype=2; %average across workers

[ylife,ylook,ageyn2,earnginiage,ageyn_dye]= stats_earnings_process_dye(ylife,Qylife,invherhip,normtype,ye,yedtb);




%% 5. Other parameters which are relevant for the loop

inv_workers=ageyn2'; %TLxdy, distribution
inc_workers=ylook'; %TLxdy, values
avg_inc_workers=sum(inv_workers.*inc_workers,2)./sum(inv_workers,2);
aveinc_p=pfrac(1:TL)'*avg_inc_workers;

avg_inc_workers_dye=sum(ageyn_dye'.*kron(ones(dy,1),ye).*...
    kron(ylook,ones(dye,1)));
aveinc_p_2=pfrac(1:TL)'*avg_inc_workers_dye';
% This is the relevant average that is kept constant across earnings
% processes


numold=sum(pfrac(TL+1:T));

%% 6. Loop
 %% 6.1. Extrapolation
%%%%%%%%%%%%%% form assets grid for which the value function will be defined
% note that how the grid is formed is crucial for later transformations
% to reconduce one grid to another
 
 if blimit==0
a=linspace(0,sqrt(maxa),da)'.^2;
azero=1;
 elseif blimit<0
apre=linspace(0,sqrt(maxa),da)'.^2;
I=find(apre>-blimit,1,'first');
aneg=-(linspace(0,sqrt(-blimit),I+1)'.^2);
apos=linspace(0,sqrt(maxa),da-I)'.^2;
a=[flipud(aneg); apos(2:end)];
azero=find(a==0,1,'first');
 end
 

%%%%%%%%%%%%%%% take care of the possibility that when interpolating
% the value function later on we fall out of the range of values we
% have for the value function itself (on positive side)
    I=find(a==0);
stpa=sqrt(a(I+1))-sqrt(a(I));
dagg=20; % number of points we will add on the grid for the value
                 % function to extend it by polynomial extrapolation
ai=[a; (sqrt(maxa)+stpa*(1:dagg)').^2]; 
dpol=2;   % dpol-1 is the degree of the polynomial we are using
% polynomial extrapolation
% how do we do it?
% take the example of extrapolation with a quadratic polynomial
% the equation is y=a+bx+cx^2, where y=[y_1 y_2 y_3]'; and analogously
% for x. we can rewrite it in matrix form as
% |y_1|   | 1 x_1  x_1^2 |   | a |
% |y_2| = | 1 x_2  x_2^2 |   | b |
% |y_3|   | 1 x_3  x_3^2 |   | c |
% invert it to get the vector of coefficients:
% |a |   | 1 x_1  x_1^2 |-1   | y_1 |
% |b | = | 1 x_2  x_2^2 |     | y_2 |
% |c |   | 1 x_3  x_3^2 |     | y_3 |
% now use:
% |y_4|   | 1 x_4  x_5^2 |   | a |
% |y_5| = | 1 x_4  x_5^2 |   | b |
% |y_6|   | 1 x_4  x_5^2 |   | c |
% substituiting the expression for the coefficients derived above, obtain:
% |y_4|   | 1 x_4  x_5^2 |   | 1 x_1  x_1^2 |-1   | y_1 |
% |y_5| = | 1 x_4  x_5^2 |   | 1 x_2  x_2^2 |     | y_2 |
% |y_6|   | 1 x_4  x_5^2 |   | 1 x_3  x_4^2 |     | y_3 |
%         call this temp2    call this temp1
% temp2*temp1 is what we will use when we will have to extrapolate the
% value function later on, when we will have the relevant y's.
temp1=(a(da-dpol+1:da)*ones(1,dpol)).^(ones(dpol,1)*(0:dpol-1));
temp2=(ai(da+1:da+dagg)*ones(1,dpol)).^(ones(dagg,1)*(0:dpol-1));
intpcf=temp2/temp1;
clear temp1 temp2

%%%%%%%%%%%%%% grid on tomorrow's assets that the agent is allowed to chose.
% It is bigger than the grid for a, for which the value function
% is defined to save on computation time while not constraining too much
% the agent's choice on a narrow grid. We will interpolate the value function
% on ar later on (using the value function defined on a) in order to compute
% the agent's maximization problem for each period of time
%ar=linspace(0,sqrt(maxa),200)'.^2;
dar=2000; %How many points for interpolation
 if blimit==0
    ar=linspace(0,sqrt(maxa),dar)'.^2;
 elseif blimit<0
    arpre=linspace(0,sqrt(maxa),dar)'.^2;
    I=find(arpre>-blimit,1,'first');
    arneg=-(linspace(0,sqrt(-blimit),I+1)'.^2);
    arpos=linspace(0,sqrt(maxa),dar-I)'.^2;
    ar=[flipud(arneg); arpos(2:end)];
 end


% Now compute pensions
% These depend on the last realization of earnings, but pre-tax rather than
% post-tax

% First reconstruct pre-tax earnings

btax = -(1-0.2177183);
atax = -0.00000009745116;

%ypre=(-btax-sqrt(btax.^2-4*atax*ylook(:,TL)))./(2*atax);

ypre=(-btax-sqrt(btax.^2-4*atax*(kron(ylook(:,TL),ones(dye,1)).*...
kron(ones(dy,1),ye(:,TL)))))./(2*atax);

pens=NaN(dy,1);
pens(ypre<0.19)=1.1*ypre(ypre<0.19);
pens(ypre>=0.19 & ypre<1.10)=0.32*(ypre(ypre>=0.19 & ypre<1.10)-0.19)+...
    1.1*0.19;
pens(ypre>=1.10)=0.15*(ypre(ypre>=1.10)-1.10)+0.32*(1.10-0.19)+...
    1.1*0.19; 

[ypresort,I]=sort(ypre);
pensort=pens(I);
[ypresort,I]=unique(ypresort); % only necessary for Markov 2 case
pensort=pensort(I);

%And scale looking for the closest to 1
ypara1 = interp1(ypresort,pensort,1);

pens=0.45*pens/ypara1; % keep avg worker with 45% replacement rate



% Get natural borrowing limits to guarantee they are respected

blimitage=zeros(T,1); %zero in T and T-1

if blimit<0
for i=T-2:-1:TL
    blimitage(i)=(1/(1+r)).*(blimitage(i+1)+pens(1));
end
for i=TL-1:-1:1
    blimitage(i)=(1/(1+r)).*(blimitage(i+1)+ylook(1,i+1));
end
end

% And allow for exogenous borrowing limit:
     for i=1:T
     blimitage=min(blimitage,-blimit);
     end

     % Find policy functions and stationary distribution
 [invm,invmt,vt,vr,copt,colr,sopt,solr,avgk,aveinc]=...
     vfi_retirement_TL_transitory_pens(dar,dpol,pens,ai,ar,sig,...
                bet,intpcf,ylife,Qylife,r,maxa,...
                 invherhip,ye,yedtb,initialwealth,...
                 blimitage);


u.ye=ye;
u.ylife=ylife;

dimTL=[da dye dy dhip TL];
dimTR=[da dy*dye TR];

%% 7. Statistics on wealth inequality

% Compute statistics on wealth + consumption

[wgini,ginage,dist,cons_aworkers,cons_aold,lcons_aworkers,lcons_aold,...
    ltransiti,lincwork,lvarlabincwork,lthirdincwork,lfourthincwork,...
    lvarcons,percentable,...
    cons_perctiles,lvarlabgroup,lvarconsgroup,negwealth,...
    varage,savbyage]=...
    compute_cw_statistics(ylook,u,invher,copt,colr,sopt,invm,invmt,dimTL,dimTR);

%% 8. Consumption equivalents

% Following Guvenen, Karahan, Ozkan and Song in the computation
[xi,xiearn]=consumption_equivalents(bet,r,invher,ylook,...
    ye,yedtb,pens,vt,Qylife,sig);

%% 9. Save simulation for BPP coefficients

rub=panel_simul_cons(Qylife,copt,sopt,ylook,ylifelog,ye,lvarcons,commonbeta,...
    choose,invher,yedtb);

%% Save and use relevant outputs as necessary

consdtb=[reshape(copt,da*dy*dye*TL,1); reshape(colr,da*dy*dye*(T-TL),1)];

if autocalibrate==1
    if blimit==-0.12 && commonbeta==0
        name2=strcat('Outputs/case',num2str(choose));
    end
end
    
 if commonbeta==1
name2 = strcat('Outputs/case',num2str(choose),'_cb');
 end

save(name2,'lvarconsgroup','savbyage','varage','percentable','xi')


disp('Top wealth holdings')
disp(percentable)
diary off;

criterio=abs(avgk/((aveinc+r*avgk)/tlength)-wiratio);
criterion=(avgk/(aveinc+r*avgk)/tlength-wiratio);

if criterio>0.01
if autocalibrate==1
   [bet,topbet,topcriterio,bottombet,bottomcriterio] = ...
    beta_bisection(bet,criterion,criterio,topbet,topcriterio,bottombet,bottomcriterio);
    fprintf('New beta is ')
    fprintf('%.4f',bet)
    fprintf(' and criterio is ')
    fprintf('%.4f',criterio)
    fprintf('\n')
end
end
if autocalibrate==0
    break
end
%

end


toc