% This wrapper file is used for the calculation of the policy experiments reported in Table 10 of the paper

clc
clear all

%% Set up the parallel environment
npool=7;
p = gcp('nocreate'); % If no pool, do not create new one.
if isempty(p)
    poolsize = 0;
else
    poolsize = p.NumWorkers;
end
if poolsize <npool
    delete(gcp('nocreate'))
    parpool('local',npool)
end
warning('off','MATLAB:mir_warning_maybe_uninitialized_temporary');


%% Load point estimates 
THETA_0=csvread('THETA.csv');   

PSI(1) = 0.211;   % psi_l1
PSI(2) = 0.162;   % psi_l2
PSI(3) = 0.535;   % rho_l
PSI(4) = 0.107;   % psi_t1
PSI(5) = 0.411;   % psi_t2
PSI(6) = -0.288;  % rho_t
PSI(7) = 0.903;   % eta_cp

%% The magnitude of the exercise
scale_exercise = 0.02;	% the magnitude of the transfer (in terms of disposable income per cent)


%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Policy experiment 1: increase b to match a certain % increase in overall transfers
%%%%%%%%%%%%%%%%%%%%%%%%%%

% Baseline parameters
chi = [2.2 2.2];    % level parameter in the tax function
taumu=[0.1 0.1];    % slope parameter in the tax function
tau_b=[12*(526+476)/4160 12*367/4160];    % minimum pre-tax transfers that the household always gets regardless of work. Note that this assumes scaling by 4160
tau_gamma1=2900;

savedata = 0;    % Choose one to save the entire output and the vf iterations
irflag= 1; % Choose 1 to run the IR (and save the results), and 0 otherwise

% Caltulate baseline
tic
[simul_output,parameters] = main_estimation_policy(THETA_0,PSI,chi,taumu,tau_b,savedata,irflag,tau_gamma1);
toc

% Come up with a guess that assumes no behavioral responses
r    = parameters.r;
TT   = parameters.TT;
beta = parameters.beta;
KTlarge=parameters.KTlarge;

policy_calcs = policy_calc(parameters,simul_output);
T1 = policy_calcs.T1;
T2 = policy_calcs.T2;
nkperiods = sum(KTlarge(1,:)==1);

discby=0;
for i=T1+1:T2
    discby = discby + beta^i;
end
dY = scale_exercise*policy_calcs.cohort_atY_disc_kids*(nkperiods/discby);
db = 1/(1-taumu(1))*(policy_calcs.cohort_Y_disc_kids/policy_calcs.cohort_atY_disc_kids)*dY;

% the initial guess
tau_b(1) = parameters.tau_b(1)+db/4160;

% Run with initial guess
tic
[simul_output_new,parameters_new] = main_estimation_policy(THETA_0,PSI,chi,taumu,tau_b,savedata,irflag,tau_gamma1);
toc
policy_calcs_new = policy_calc(parameters_new,simul_output_new);

factor=(policy_calcs.cohort_tax_disc_kids-policy_calcs_new.cohort_tax_disc_kids)/dY;
if abs(factor-1)<=0.02 
    simul_output_new2=simul_output_new;
    parameters_new2 = parameters_new;
end

% repeat until convergence
while abs(factor-1)>0.02 
    db=db/factor;
    tau_b(1) = parameters.tau_b(1)+db/4160; 
    tic
    [simul_output_new2,parameters_new2] = main_estimation_policy(THETA_0,PSI,chi,taumu,tau_b,savedata,irflag,tau_gamma1);
    toc
    policy_calcs_new2 = policy_calc(parameters_new2,simul_output_new2);
    factor=(policy_calcs.cohort_tax_disc_kids-policy_calcs_new2.cohort_tax_disc_kids)/dY;
end

save policy1

%% Policy implications
% Preparations
kt_simul = simul_output.kt_simul;
discfact = zeros(size(kt_simul));
for i=1:size(kt_simul,2)
    discfact(:,i) = (1/(1+r))^(i-1);
end

pre =zeros(size(kt_simul));
pre(:,1:T1)=1; 
kids=zeros(size(kt_simul));
kids(:,T1+1:T2)=1; 
post=zeros(size(kt_simul));
post(:,T2+1:TT)=1; 

% Calculations
C_simul = simul_output.C_simul;
C_simul_new = simul_output_new2.C_simul;
H2_simul = simul_output.H2_simul;
H2_simul_new = simul_output_new2.H2_simul;

H1_simul = simul_output.H1_simul;
H1_simul_new = simul_output_new2.H1_simul;
Z1_simul = H1_simul<10^-9;    
Z1_simul_new = H1_simul_new<10^-9;    
Z2_simul = H2_simul<10^-9;    
Z2_simul_new = H2_simul_new<10^-9;    
E2_simul = 1-Z2_simul;    
E2_simul_new = 1-Z2_simul_new;    
L1_simul = simul_output.L1_simul;
L1_simul_new = simul_output_new2.L1_simul;
T1_simul = simul_output.T1_simul;
T1_simul_new = simul_output_new2.T1_simul;
L2_simul = simul_output.L2_simul;
L2_simul_new = simul_output_new2.L2_simul;
T2_simul = simul_output.T2_simul;
T2_simul_new = simul_output_new2.T2_simul;

U_simul = calcU(C_simul,L1_simul,L2_simul,T1_simul,T2_simul,E2_simul,kt_simul,parameters);
U_simul_new = calcU(C_simul_new,L1_simul_new,L2_simul_new,T1_simul_new,T2_simul_new,E2_simul_new,kt_simul,parameters_new2);

[U_simul_disc,U_simul_new_disc,dU_total,dU_kids,dU_pre,dU_post]=calc_dc(U_simul,U_simul_new,discfact,pre,kids,post);
[C_simul_disc,C_simul_new_disc,dC_total,dC_kids,dC_pre,dC_post]=calc_dc(C_simul,C_simul_new,discfact,pre,kids,post);
[H1_simul_disc,H1_simul_new_disc,dH1_total,dH1_kids,dH1_pre,dH1_post]=calc_dc(H1_simul,H1_simul_new,discfact,pre,kids,post);
[H2_simul_disc,H2_simul_new_disc,dH2_total,dH2_kids,dH2_pre,dH2_post]=calc_dc(H2_simul,H2_simul_new,discfact,pre,kids,post);
[E2_simul_disc,E2_simul_new_disc,dE2_total,dE2_kids,dE2_pre,dE2_post]=calc_dc(E2_simul,E2_simul_new,discfact,pre,kids,post);
[L1_simul_disc,L1_simul_new_disc,dL1_total,dL1_kids,dL1_pre,dL1_post]=calc_dc(L1_simul,L1_simul_new,discfact,pre,kids,post);
[L2_simul_disc,L2_simul_new_disc,dL2_total,dL2_kids,dL2_pre,dL2_post]=calc_dc(L2_simul,L2_simul_new,discfact,pre,kids,post);
[T1_simul_disc,T1_simul_new_disc,dT1_total,dT1_kids,dT1_pre,dT1_post]=calc_dc(T1_simul,T1_simul_new,discfact,pre,kids,post);
[T2_simul_disc,T2_simul_new_disc,dT2_total,dT2_kids,dT2_pre,dT2_post]=calc_dc(T2_simul,T2_simul_new,discfact,pre,kids,post);

target=mean(U_simul_new_disc(:));
f = @(pi)calc_pi(pi,target,C_simul,L1_simul,L2_simul,T1_simul,T2_simul,E2_simul,kt_simul,U_simul_new,parameters,discfact,pre,kids,post);
[pi,~]= fzero(f,[1 1.5]); 

results = [dU_total,dC_total,dH1_total,dH2_total,dE2_total,dL1_total,dL2_total,dT1_total,dT2_total;...
           dU_pre,dC_pre,dH1_pre,dH2_pre,dE2_pre,dL1_pre,dL2_pre,dT1_pre,dT2_pre;...
           dU_kids,dC_kids,dH1_kids,dH2_kids,dE2_kids,dL1_kids,dL2_kids,dT1_kids,dT2_kids;...
           dU_post,dC_post,dH1_post,dH2_post,dE2_post,dL1_post,dL2_post,dT1_post,dT2_post;...
           ones(1,9)*pi];

csvwrite('results_policy1.csv',results);

%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Policy experiment 2
%%%%%%%%%%%%%%%%%%%%%%%%%%
chi = [2.2 2.2];    
taumu=[0.1 0.1];    
tau_b=[12*(526+476)/4160 12*367/4160];    
dgamma = dY/mean(E2_simul(kids==1));

% the initial guess
tau_gamma1_new = tau_gamma1-dgamma;

% Run with initial guess
tic
[simul_output_new,parameters_new] = main_estimation_policy(THETA_0,PSI,chi,taumu,tau_b,savedata,irflag,tau_gamma1_new);
toc
H2_simul_new = simul_output_new.H2_simul;
Z2_simul_new = H2_simul_new<10^-9;    
E2_simul_new = 1-Z2_simul_new;    

factor=mean(E2_simul_new(kids==1))*dgamma/dY

if abs(factor-1)<=0.02 
    simul_output_new3=simul_output_new;
    parameters_new3 = parameters_new;
end

% repeat until convergence
while abs(factor-1)>0.02 
    dgamma=dgamma/factor;
    tau_gamma1_new = tau_gamma1-dgamma;
    tic
    [simul_output_new3,parameters_new3] = main_estimation_policy(THETA_0,PSI,chi,taumu,tau_b,savedata,irflag,tau_gamma1_new);
    toc
    H2_simul_new3 = simul_output_new3.H2_simul;
    Z2_simul_new3 = H2_simul_new3<10^-9;    
    E2_simul_new3 = 1-Z2_simul_new3;    

    factor=mean(E2_simul_new3(kids==1))*dgamma/dY
end

save policy2

% Calculations
C_simul_new_p2 = simul_output_new3.C_simul;
H2_simul_new_p2 = simul_output_new3.H2_simul;

H1_simul_new_p2 = simul_output_new3.H1_simul;
Z1_simul_new_p2 = H1_simul_new_p2<10^-9;    
Z2_simul_new_p2 = H2_simul_new_p2<10^-9;    
E2_simul_new_p2 = 1-Z2_simul_new_p2;    
L1_simul_new_p2 = simul_output_new3.L1_simul;
T1_simul_new_p2 = simul_output_new3.T1_simul;
L2_simul_new_p2 = simul_output_new3.L2_simul;
T2_simul_new_p2 = simul_output_new3.T2_simul;

U_simul_new_p2 = calcU(C_simul_new_p2,L1_simul_new_p2,L2_simul_new_p2,T1_simul_new_p2,T2_simul_new_p2,E2_simul_new_p2,kt_simul,parameters_new3);

[U_simul_disc,U_simul_new_p2_disc,dU_total_p2,dU_kids_p2,dU_pre_p2,dU_post_p2]=calc_dc(U_simul,U_simul_new_p2,discfact,pre,kids,post);
[C_simul_disc,C_simul_new_p2_disc,dC_total_p2,dC_kids_p2,dC_pre_p2,dC_post_p2]=calc_dc(C_simul,C_simul_new_p2,discfact,pre,kids,post);
[H1_simul_disc,H1_simul_new_p2_disc,dH1_total_p2,dH1_kids_p2,dH1_pre_p2,dH1_post_p2]=calc_dc(H1_simul,H1_simul_new_p2,discfact,pre,kids,post);
[H2_simul_disc,H2_simul_new_p2_disc,dH2_total_p2,dH2_kids_p2,dH2_pre_p2,dH2_post_p2]=calc_dc(H2_simul,H2_simul_new_p2,discfact,pre,kids,post);
[E2_simul_disc,E2_simul_new_p2_disc,dE2_total_p2,dE2_kids_p2,dE2_pre_p2,dE2_post_p2]=calc_dc(E2_simul,E2_simul_new_p2,discfact,pre,kids,post);
[L1_simul_disc,L1_simul_new_p2_disc,dL1_total_p2,dL1_kids_p2,dL1_pre_p2,dL1_post_p2]=calc_dc(L1_simul,L1_simul_new_p2,discfact,pre,kids,post);
[L2_simul_disc,L2_simul_new_p2_disc,dL2_total_p2,dL2_kids_p2,dL2_pre_p2,dL2_post_p2]=calc_dc(L2_simul,L2_simul_new_p2,discfact,pre,kids,post);
[T1_simul_disc,T1_simul_new_p2_disc,dT1_total_p2,dT1_kids_p2,dT1_pre_p2,dT1_post_p2]=calc_dc(T1_simul,T1_simul_new_p2,discfact,pre,kids,post);
[T2_simul_disc,T2_simul_new_p2_disc,dT2_total_p2,dT2_kids_p2,dT2_pre_p2,dT2_post_p2]=calc_dc(T2_simul,T2_simul_new_p2,discfact,pre,kids,post);

target=mean(U_simul_new_p2_disc(:));
f = @(pi)calc_pi(pi,target,C_simul,L1_simul,L2_simul,T1_simul,T2_simul,E2_simul,kt_simul,U_simul_new,parameters,discfact,pre,kids,post);
[pi,~]= fzero(f,[0.9 1.5]); 

results = [dU_total_p2,dC_total_p2,dH1_total_p2,dH2_total_p2,dE2_total_p2,dL1_total_p2,dL2_total_p2,dT1_total_p2,dT2_total_p2;...
           dU_pre_p2,dC_pre_p2,dH1_pre_p2,dH2_pre_p2,dE2_pre_p2,dL1_pre_p2,dL2_pre_p2,dT1_pre_p2,dT2_pre_p2;...
           dU_kids_p2,dC_kids_p2,dH1_kids_p2,dH2_kids_p2,dE2_kids_p2,dL1_kids_p2,dL2_kids_p2,dT1_kids_p2,dT2_kids_p2;...
           dU_post_p2,dC_post_p2,dH1_post_p2,dH2_post_p2,dE2_post_p2,dL1_post_p2,dL2_post_p2,dT1_post_p2,dT2_post_p2;...
           ones(1,9)*pi];

csvwrite('results_policy2.csv',results);

save policy_all
