% This function takes as input the parameters, the outputs from the
% wage grids and value function solution and simulates:
% (1) Wages for each earners in the household
% (2) Asset and time use choices using the policy functions
% (3) Consumption and income from the budget constraint

function simul_output = simul_all(parameters,wage_grids_outputs,output_vf)

%% Read parameters
TT     = parameters.TT      ;
eta_cp = parameters.eta_cp ;
mu_0  = parameters.mu_0  ;
psi_l1 = parameters.psi_l1 ;
psi_l2 = parameters.psi_l2 ;
rhou_l = parameters.rhou_l ;
psi_t1 = parameters.psi_t1 ;
psi_t2 = parameters.psi_t2 ;
rhou_t = parameters.rhou_t ;
mu_t1 = parameters.mu_t1 ;
mu_t2 = parameters.mu_t2 ;
amax   = parameters.amax   ;
amin   = parameters.amin   ;
apoints= parameters.apoints;
L_bar  = parameters.L_bar;
beta   = parameters.beta   ;
r      = parameters.r      ;        
div    = parameters.div;
wpoints= parameters.wpoints;
theta  = parameters.theta  ;
n      = parameters.n ;
KT     = parameters.KT ;
CHI    = parameters.chi    ;
TAUMU  = parameters.taumu  ;
TAU_B  = parameters.tau_b  ;
TAU_GAMMA= parameters.tau_gamma;
A0_pzero = parameters.A0_pzero;
A0_mu    = parameters.A0_mu;
A0_sig   = parameters.A0_sig;

v1_grid    = wage_grids_outputs.v1_grid   ;
v2_grid    = wage_grids_outputs.v2_grid   ;
u1_grid    = wage_grids_outputs.u1_grid   ;
u2_grid    = wage_grids_outputs.u2_grid   ;
v1_prob    = wage_grids_outputs.v1_prob   ;
v2_prob    = wage_grids_outputs.v2_prob   ;
u1_prob    = wage_grids_outputs.u1_prob   ;
u2_prob    = wage_grids_outputs.u2_prob   ;
lw_grid_1  = wage_grids_outputs.lw_grid_1 ;
lw_grid_2  = wage_grids_outputs.lw_grid_2 ;
gamma_u    = wage_grids_outputs.gamma_u   ; 
gamma_v    = wage_grids_outputs.gamma_v   ;   
sig_u2_hat = wage_grids_outputs.sig_u2_hat;
sig_v2_hat = wage_grids_outputs.sig_v2_hat;
log_w_det_1= wage_grids_outputs.log_w_det_1;
log_w_det_2= wage_grids_outputs.log_w_det_2;

vpoints= size(v1_grid,1); 
upoints= size(u1_grid,1); 

sd_u1   = sqrt(theta(1));
sd_u2   = sqrt(theta(2));
sd_v1   = sqrt(theta(3));
sd_v2   = sqrt(theta(4));
cov_u1u2= theta(5);
cov_v1v2= theta(6);
sd_w0_1 = sqrt(theta(7));
sd_w0_2 = sqrt(theta(8));

A_grid   = output_vf.A_grid;
A_star_E2_0 = output_vf.A_star_E2_0;
A_star_E2_1 = output_vf.A_star_E2_1;
A_star   = output_vf.A_star;
L1_star_E2_0  = output_vf.L1_star_E2_0;
L1_star_E2_1  = output_vf.L1_star_E2_1;
L1_star  = output_vf.L1_star;
L2_star_E2_0  = output_vf.L2_star_E2_0;
L2_star_E2_1  = output_vf.L2_star_E2_1;
L2_star  = output_vf.L2_star;
T1_star_E2_0  = output_vf.T1_star_E2_0;
T1_star_E2_1  = output_vf.T1_star_E2_1;
T1_star  = output_vf.T1_star;
T2_star_E2_0  = output_vf.T2_star_E2_0;
T2_star_E2_1  = output_vf.T2_star_E2_1;
T2_star  = output_vf.T2_star;
V_star_E2_1 = output_vf.V_star_E2_1;
V_star_E2_0 = output_vf.V_star_E2_0;

% Simulated shocks for wage process (will be scaled by the parameters, but
% drawn only once. 
d1 = rand(n,TT);  % v1 
d2 = rand(n,TT);  % v2 
d3 = randn(n,TT);  % u1 
d4 = randn(n,TT);  % u2 
d5 = randn(n,1);    % w0_1 
d6 = randn(n,1);    % w0_2 
d7 = rand(n,1);     % proportion of zero assets at time zero
d8 = randn(n,1);    % log assets at time zero for non-zero asset households

%% Simulate v1,v2,u1,u2 for all periods (both indices and levels)
% v1
unif  = d1;
v1_cdf = cumsum(v1_prob);
v1_ind = (vpoints)*ones(size(unif));
v1 = zeros(size(unif));

for i = 1:vpoints-1
    v1_indtemp= (unif<v1_cdf(i));
    v1_ind = v1_ind - v1_indtemp; 
end

for tau=1:TT
    v1(:,tau) = v1_grid(v1_ind(:,tau),tau);
end

% v2
unif  = d2;
v2_cdf = cumsum(v2_prob);
v2_ind = (vpoints)*ones(size(unif));
v2_hat = zeros(size(unif));

for i = 1:vpoints-1
    v2_indtemp= (unif<v2_cdf(i));
    v2_ind = v2_ind - v2_indtemp; 
end

for tau=1:TT
    v2_hat(:,tau) = v2_grid(v2_ind(:,tau),tau);
end
v2 = gamma_v*v1 + v2_hat;   % Account for the correlation between the shocks of the two earners

% u1
u1 = d3*sd_u1;

% u2
u2_hat = d4*sd_u2;
u2 = gamma_u*u1 + u2_hat;   % Account for the correlation between the shocks of the two earners

%% Simulate the wage processes for husband and wife
log_w_det_1=log_w_det_1(2:TT+1)';  
log_w_det_2=log_w_det_2(2:TT+1)';  

F1 = zeros(n,TT);
F2 = zeros(n,TT);
w0_1 = sd_w0_1*d5;
w0_2 = sd_w0_1*d6;
F1(:,1) = w0_1 + v1(:,1);
F2(:,1) = w0_2 + v2(:,1);

for tau = 1:TT-1
   F1(:,tau+1) = F1(:,tau) + v1(:,tau+1); 
   F2(:,tau+1) = F2(:,tau) + v2(:,tau+1); 
end
F1 = ones(n,1)*log_w_det_1 + F1; 
F2 = ones(n,1)*log_w_det_2 + F2; 

% In the very few cases that the simulated wages go off the grid, assign
% them the maximal or minimal point on the grid
maxtemp1 = ones(n,1)*max(lw_grid_1);
maxtemp2 = ones(n,1)*max(lw_grid_2);
F1(F1>maxtemp1)= maxtemp1(F1>maxtemp1);
F2(F2>maxtemp2)= maxtemp2(F2>maxtemp2);

mintemp1 = ones(n,1)*min(lw_grid_1);
mintemp2 = ones(n,1)*min(lw_grid_2);
F1(F1<mintemp1)= mintemp1(F1<mintemp1);
F2(F2<mintemp2)= mintemp2(F2<mintemp2);

%same for transitory shocks
maxtemp1 = ones(n,1)*max(u1_grid);
maxtemp2 = ones(n,1)*max(u2_grid);
u1(u1>maxtemp1)= maxtemp1(u1>maxtemp1);
u2(u2>maxtemp2)= maxtemp2(u2>maxtemp2);

mintemp1 = ones(n,1)*min(u1_grid);
mintemp2 = ones(n,1)*min(u2_grid);
u1(u1<mintemp1)= mintemp1(u1<mintemp1);
u2(u2<mintemp2)= mintemp2(u2<mintemp2);

lw_1 = F1 + u1; 
lw_2 = F2 + u2; 

%% Prepare matrices
V_E2_1_simul = zeros(size(F1));
V_E2_0_simul = zeros(size(F1));
A_simul  = zeros(size(F1));
L1_simul = zeros(size(F1));
L2_simul = zeros(size(F1));
T1_simul = zeros(size(F1));
T2_simul = zeros(size(F1));
H1_simul = zeros(size(F1));
H2_simul = zeros(size(F1));
Y1_simul = zeros(size(F1));
Y2_simul = zeros(size(F1));
Y_simul  = zeros(size(F1));
atY_simul= zeros(size(F1));
C_simul  = zeros(size(F1));
kt_simul = zeros(size(F1));
E2_simul = zeros(size(F1));

%% Initialize assets
indnonzero = d7>A0_pzero;
A_simul(:,1) = exp(A0_sig*d8 + A0_mu).*indnonzero;
A_simul(:,1) = A_simul(:,1)./(4160^(1-TAUMU(KT(1,1)))); %scale correctly 
A_simul(:,1) = min(A_simul(:,1),max(A_grid));

test= zeros(size(F1));
V_star= output_vf.V_star;

%% Use the interpolated VF for employment choice and choose the right assets accordingly
for tau=1:TT
    V_E2_1_simul(:,tau) = interpn(A_grid,lw_grid_1(:,tau),lw_grid_2(:,tau),u1_grid(:,tau),u2_grid(:,tau),...
                              squeeze(V_star_E2_1(:,:,:,:,:,tau)),...
                              A_simul(:,tau),F1(:,tau),F2(:,tau),u1(:,tau),u2(:,tau));
    V_E2_0_simul(:,tau) = interpn(A_grid,lw_grid_1(:,tau),lw_grid_2(:,tau),u1_grid(:,tau),u2_grid(:,tau),...
                              squeeze(V_star_E2_0(:,:,:,:,:,tau)),...
                              A_simul(:,tau),F1(:,tau),F2(:,tau),u1(:,tau),u2(:,tau));
    test(:,tau) = interpn(A_grid,lw_grid_1(:,tau),lw_grid_2(:,tau),u1_grid(:,tau),u2_grid(:,tau),...
                              squeeze(V_star(:,:,:,:,:,tau)),...
                              A_simul(:,tau),F1(:,tau),F2(:,tau),u1(:,tau),u2(:,tau));
    E2_simul(:,tau) = (V_E2_1_simul(:,tau)>=V_E2_0_simul(:,tau));
    temp1 = max(interpn(A_grid,lw_grid_1(:,tau),lw_grid_2(:,tau),u1_grid(:,tau),u2_grid(:,tau),...
                                 squeeze(A_star_E2_1(:,:,:,:,:,tau)),...
                                 A_simul(:,tau),F1(:,tau),F2(:,tau),u1(:,tau),u2(:,tau)),amin);
    temp1 = min(temp1,max(A_grid));                         
    temp0 = max(interpn(A_grid,lw_grid_1(:,tau),lw_grid_2(:,tau),u1_grid(:,tau),u2_grid(:,tau),...
                                 squeeze(A_star_E2_0(:,:,:,:,:,tau)),...
                                 A_simul(:,tau),F1(:,tau),F2(:,tau),u1(:,tau),u2(:,tau)),amin);
    temp0 = min(temp0,max(A_grid));
    if tau<TT
        A_simul(:,tau+1) = temp1.*E2_simul(:,tau) + temp0.*(1-E2_simul(:,tau));
    end
                          
end

A_simul(:,TT+1) = 0;    % terminal condition for assets. Zero at the end of the last period (or beginning of the TT+1 period)

%% Simulate time use goods using the policy functions, and income, consumption and hours using the budget and time constraints
for tau=1:TT
    temp1 = interpn(A_grid,lw_grid_1(:,tau),lw_grid_2(:,tau),u1_grid(:,tau),u2_grid(:,tau),...
                              squeeze(L1_star_E2_1(:,:,:,:,:,tau)),...
                              A_simul(:,tau),F1(:,tau),F2(:,tau),u1(:,tau),u2(:,tau));
    temp0 = interpn(A_grid,lw_grid_1(:,tau),lw_grid_2(:,tau),u1_grid(:,tau),u2_grid(:,tau),...
                              squeeze(L1_star_E2_0(:,:,:,:,:,tau)),...
                              A_simul(:,tau),F1(:,tau),F2(:,tau),u1(:,tau),u2(:,tau));
	L1_simul(:,tau) = temp1.*E2_simul(:,tau) + temp0.*(1-E2_simul(:,tau));
                          
    temp1 = interpn(A_grid,lw_grid_1(:,tau),lw_grid_2(:,tau),u1_grid(:,tau),u2_grid(:,tau),...
                              squeeze(L2_star_E2_1(:,:,:,:,:,tau)),...
                              A_simul(:,tau),F1(:,tau),F2(:,tau),u1(:,tau),u2(:,tau));
    temp0 = interpn(A_grid,lw_grid_1(:,tau),lw_grid_2(:,tau),u1_grid(:,tau),u2_grid(:,tau),...
                              squeeze(L2_star_E2_0(:,:,:,:,:,tau)),...
                              A_simul(:,tau),F1(:,tau),F2(:,tau),u1(:,tau),u2(:,tau));
    L2_simul(:,tau) = temp1.*E2_simul(:,tau) + temp0.*(1-E2_simul(:,tau));

    temp1 = interpn(A_grid,lw_grid_1(:,tau),lw_grid_2(:,tau),u1_grid(:,tau),u2_grid(:,tau),...
                              squeeze(T1_star_E2_1(:,:,:,:,:,tau)),...
                              A_simul(:,tau),F1(:,tau),F2(:,tau),u1(:,tau),u2(:,tau));
    temp0 = interpn(A_grid,lw_grid_1(:,tau),lw_grid_2(:,tau),u1_grid(:,tau),u2_grid(:,tau),...
                              squeeze(T1_star_E2_0(:,:,:,:,:,tau)),...
                              A_simul(:,tau),F1(:,tau),F2(:,tau),u1(:,tau),u2(:,tau));
    T1_simul(:,tau) = temp1.*E2_simul(:,tau) + temp0.*(1-E2_simul(:,tau));

    temp1 = interpn(A_grid,lw_grid_1(:,tau),lw_grid_2(:,tau),u1_grid(:,tau),u2_grid(:,tau),...
                              squeeze(T2_star_E2_1(:,:,:,:,:,tau)),...
                              A_simul(:,tau),F1(:,tau),F2(:,tau),u1(:,tau),u2(:,tau));
    temp0 = interpn(A_grid,lw_grid_1(:,tau),lw_grid_2(:,tau),u1_grid(:,tau),u2_grid(:,tau),...
                              squeeze(T2_star_E2_0(:,:,:,:,:,tau)),...
                              A_simul(:,tau),F1(:,tau),F2(:,tau),u1(:,tau),u2(:,tau));
    T2_simul(:,tau) = temp1.*E2_simul(:,tau) + temp0.*(1-E2_simul(:,tau));
end

for tau=1:TT
    kt   = KT(tau);

    chi  = CHI(kt);
    taumu= TAUMU(kt);
    tau_b= TAU_B(kt);
    tau_gamma= TAU_GAMMA(kt);

    H1_simul(:,tau)= L_bar - L1_simul(:,tau) - T1_simul(:,tau);
    H2_simul(:,tau)= L_bar - L2_simul(:,tau) - T2_simul(:,tau);
    Y1_simul(:,tau)= exp(F1(:,tau)+u1(:,tau)).*H1_simul(:,tau);
    Y2_simul(:,tau)= exp(F2(:,tau)+u2(:,tau)).*H2_simul(:,tau);
    Y_simul(:,tau) = tau_b + Y1_simul(:,tau)+Y2_simul(:,tau);
    atY_simul(:,tau) = chi.*Y_simul(:,tau).^(1-taumu);
    C_simul(:,tau)= A_simul(:,tau) + atY_simul(:,tau) - (1/(1+r))*A_simul(:,tau+1);
end

% We rescale here to non-normalized numbers
for tau=1:TT
    kt   = KT(tau);
    taumu= TAUMU(kt);

    sc_L = 4160;
    sc_C = 4160^(1-taumu);

    Y1_simul(:,tau)  = sc_L*Y1_simul(:,tau) ;
    Y2_simul(:,tau)  = sc_L*Y2_simul(:,tau) ;
    Y_simul(:,tau)   = sc_L*Y_simul(:,tau)   ;
    atY_simul(:,tau) = sc_C*atY_simul(:,tau);
    C_simul(:,tau)   = sc_C*C_simul(:,tau)   ;
    A_simul(:,tau)   = sc_C*A_simul(:,tau)   ;
    H1_simul(:,tau)  = sc_L*H1_simul(:,tau) ;
    H2_simul(:,tau)  = sc_L*H2_simul(:,tau) ;
    L1_simul(:,tau)  = sc_L*L1_simul(:,tau) ;
    L2_simul(:,tau)  = sc_L*L2_simul(:,tau) ;
    T1_simul(:,tau)  = sc_L*T1_simul(:,tau) ;
    T2_simul(:,tau)  = sc_L*T2_simul(:,tau) ; 
    kt_simul(:,tau)  = kt ; 

end

%% Assign to struct
simul_output.F1_simul = F1;
simul_output.F2_simul = F2;
simul_output.u1_simul = u1;
simul_output.u2_simul = u2;
simul_output.v1_simul = v1;
simul_output.v2_simul = v2;
simul_output.w0_1_simul = w0_1;
simul_output.w0_2_simul = w0_2;
simul_output.lw_1_simul = lw_1;
simul_output.lw_2_simul = lw_2;
simul_output.Y1_simul  = Y1_simul;
simul_output.Y2_simul  = Y2_simul;
simul_output.Y_simul  = Y_simul;
simul_output.atY_simul= atY_simul;
simul_output.C_simul  = C_simul;
simul_output.A_simul  = A_simul;
simul_output.H1_simul = H1_simul;
simul_output.H2_simul = H2_simul;
simul_output.L1_simul = L1_simul;
simul_output.L2_simul = L2_simul;
simul_output.T1_simul = T1_simul;
simul_output.T2_simul = T2_simul;
simul_output.kt_simul = kt_simul;
simul_output.E2_simul = E2_simul;

