function [simul_output,parameters]=main_estimation_policy(THETA,PSI,chi,taumu,tau_b,savedata,irflag,tau_gamma1)
%% This function recieves a set of parameters, and returns the value of the 
%  objective for these parameters. The objective is the standard SMM
%  objective function. 
%  Note that the parameters which are set outside the model are set
%  directly here as well.

psi_l1= PSI(1); 
psi_l2= PSI(2); 
rhou_l = PSI(3);
psi_t1= PSI(4); 
psi_t2= PSI(5); 
rhou_t = PSI(6);
eta_cp= PSI(7); 

mu_l1 = [THETA(1),THETA(2)]  ; % Note: (phi)s from the paper were replaced by (mu)s to avoid types of psi and phi
mu_l2 = [THETA(3),THETA(4)] ; 
mu_t1 = [THETA(5),0]  ;
mu_t2 = [THETA(6),0]   ;
mu_0  = [THETA(7), 1];
sd_eps = [THETA(9) THETA(10)];   

tau_gamma = 10^9*[tau_gamma1/(10^9) THETA(8)]./(4160.^(1-taumu)); % Fixed cost of emplyoment of the wife.

%% Read the file with external parameters
theta = csvread('external_estimates.csv');   % [sig2_u1, sig2_u2, sig2_v1, sig2_v2, 
                                            % sig_u1u2, sig_v1v2, 
                                            % sig2_w0_1, sig2_w0_2, 
                                            % beta1_1, beta1_2, beta2_1,
                                            % beta2_2, beta0_1, beta0_1]
                                            % (betas for age profile)

%% User Parameters
 
% Economic parameters
TTT = 41;      % number of years. In the current version of the paper these are also number of periods

minage = 24;    % This is the age at t0. Absolute age required for deterministic age profile

div  = 1;      % Number of subdivisions for the year (1 for annual, 4 for quarterly, 12 for monthly and so on). 
TT   = TTT*div;  % The actual number of periods. 

KT1   = [ones(1,3)*2,ones(1,10),ones(1,TT-13)*2];  % Two types of families in terms of when kids arrive (according to the median in the data we use)
KT2   = ones(1,TT)*2;  % 
KTlarge = [KT1; KT2];
P_KTlarge = [1,0];    % Can be used if want to have a mix of families with and without kids. Not used in the published version of the paper. 


mu_0_E2 = [1 1];    

epspoints = 2;  % Number of grid points for epsilon (heterogeneous leisure preferences)
[epsl2_xgh, epsl2_wgh] = hermquad(epspoints);
eps_grid1 = (sqrt(2)*sd_eps(1)*epsl2_xgh);
eps_grid2 = (sqrt(2)*sd_eps(2)*epsl2_xgh);
P_mu_l2_shift  = flip((1/sqrt(pi))*epsl2_wgh)';
mu_l2_shift = exp([eps_grid1 eps_grid2]);

% time constant params
rr  = 0.05;   % Annual interest rate

% Derived parameters
r    = (1+rr)^(1/div)-1; % Interest rate adjustment for the subdivision of periods
beta = 1/(1+r); % discount factor equal r. 

% Optimization parameters 
rng(100);  % Seed

% Asset grid
amax = log(4000000);   % maximum log assets 
apoints  =27;   % 50 is what used before. number of grid points on assets grid 
amin = 0;    

% Time constraints grid
L_bar = 1;   % Bounded at maximum hours possible our of 16 hours a day * 5 * 52, assuming that the rest is sleep (if including weekends, change 5 to 7)

% Wage and wage growth grids
wpoints= 11;    % points on the wage grid (F_t-1) for each period. same size for husband and wife. 
maxw   = 3;     % number of s.d. around the mean for F_t-1 grid.
vpoints= 3;     % points on the v grid for each period. same size for husband and wife. Needs to be odd number
upoints= 3;     % points on the u grid for each period. same size for husband and wife. Needs to be odd number

% simulation parameters
n      = 100000;  % Number of workers in each simulation
A0_pzero= 0.23; 
A0_mu   = 10.44; 
A0_sig  = 1.61; 

%% Construct the parameters struct (revised below in the loop)
parameters.TT      = TT      ;
parameters.eta_cp  = eta_cp ;
parameters.mu_0   = mu_0  ; 
parameters.psi_l1  = psi_l1 ;
parameters.psi_l2  = psi_l2 ;
parameters.rhou_l  = rhou_l ;
parameters.mu_l1  = mu_l1 ;
parameters.mu_l2  = mu_l2 ;
parameters.psi_t1  = psi_t1 ;
parameters.psi_t2  = psi_t2 ;
parameters.rhou_t  = rhou_t ;
parameters.mu_t1  = mu_t1 ;
parameters.mu_t2  = mu_t2 ;
parameters.chi     = chi     ;  
parameters.taumu   = taumu   ;  
parameters.amax    = amax    ;
parameters.amin    = amin    ;
parameters.apoints = apoints ;
parameters.L_bar   = L_bar   ;
parameters.beta    = beta    ;
parameters.r       = r       ;        
parameters.div     = div     ;        
parameters.minage  = minage  ;
parameters.wpoints = wpoints ;
parameters.maxw    = maxw    ;
parameters.vpoints = vpoints ;
parameters.upoints = upoints ;
parameters.theta   = theta   ;
parameters.n       = n       ;
parameters.A0_pzero= A0_pzero;
parameters.A0_mu   = A0_mu   ;
parameters.A0_sig  = A0_sig  ;
parameters.tau_b   = tau_b     ;
parameters.tau_gamma= tau_gamma;
parameters.mu_l2_shift  = mu_l2_shift      ;
parameters.P_mu_l2_shift= P_mu_l2_shift      ;
parameters.KTlarge = KTlarge ;
parameters.P_KTlarge = P_KTlarge ;
parameters.mu_0_E2  = mu_0_E2;
parameters.sd_eps = sd_eps;

%% Generate the wage grids
wage_grids_outputs = wage_grids(parameters);
save inputvf

%% Solve the value function 
parameters1 = parameters;
parameters1.KT = KTlarge(1,:);
parameters1.mu_l2  = mu_l2.*mu_l2_shift(1,:);
output_vf_large(1,1) = solve_value_function_kid_lbar(parameters1,wage_grids_outputs);

parameters1 = parameters;
parameters1.KT = KTlarge(2,:);
parameters1.mu_l2  = mu_l2.*mu_l2_shift(1,:);
output_vf_large(2,1) = solve_value_function_kid_lbar_short(parameters1,wage_grids_outputs,output_vf_large(1,1));

parameters1 = parameters;
parameters1.KT = KTlarge(1,:);
parameters1.mu_l2  = mu_l2.*mu_l2_shift(2,:);
output_vf_large(1,2) = solve_value_function_kid_lbar(parameters1,wage_grids_outputs);

parameters1 = parameters;
parameters1.KT = KTlarge(2,:);
parameters1.mu_l2  = mu_l2.*mu_l2_shift(2,:);
output_vf_large(2,2) = solve_value_function_kid_lbar_short(parameters1,wage_grids_outputs,output_vf_large(1,2));

%% Simulate a single cohort using the initial distribution of wages
for k=1:size(KTlarge,1)
    for u=1:size(mu_l2_shift,1)
        parameters1 = parameters;
        parameters1.KT = KTlarge(k,:);
        parameters1.mu_l2  = mu_l2.*mu_l2_shift(u,:);
        simul_output = simul_all(parameters1,wage_grids_outputs,output_vf_large(k,u));
        simul_output_large(k,u)=simul_output;
    end
end

%% Create a mix of simulated households
simul_output = simul_all_het(parameters,simul_output_large);

%% Store temp results for each function call 
cc=clock;
timestamp=strcat(num2str(cc(1)),'_',num2str(cc(2)),'_',num2str(cc(3)),'_',num2str(cc(4)*100+cc(5)));
dirname = strcat(timestamp);
mkdir(dirname);

tempparm = struct2cell(parameters);
tempnames = fieldnames(parameters);
cell2csv(strcat(dirname,'/parameters_used.txt'),tempparm);
cell2csv(strcat(dirname,'/parameters_names.txt'),tempnames);

F1  = simul_output.F1_simul       ;
F2  = simul_output.F2_simul       ;
u1  = simul_output.u1_simul       ;
u2  = simul_output.u2_simul       ;
v1  = simul_output.v1_simul       ;
v2  = simul_output.v2_simul       ;
lw_1  = simul_output.lw_1_simul   ;
lw_2  = simul_output.lw_2_simul   ;
Y1_simul  = simul_output.Y1_simul ;
Y2_simul  = simul_output.Y2_simul ;
Y_simul  = simul_output.Y_simul   ;
atY_simul = simul_output.atY_simul;
C_simul  = simul_output.C_simul   ;
A_simul  = simul_output.A_simul   ;
L1_simul  = simul_output.L1_simul ;
L2_simul  = simul_output.L2_simul ;
T1_simul  = simul_output.T1_simul ;
T2_simul  = simul_output.T2_simul ; 
H1_simul  = simul_output.H1_simul ;
H2_simul  = simul_output.H2_simul ;
H1_simul=(H1_simul>=80).*H1_simul;  % Replace low values with zeros for comparability with procedure in the data
H2_simul=(H2_simul>=80).*H2_simul;
kt_simul  = simul_output.kt_simul;

F1_nk  = simul_output.F1_simul_nk       ;
F2_nk  = simul_output.F2_simul_nk       ;
u1_nk  = simul_output.u1_simul_nk       ;
u2_nk  = simul_output.u2_simul_nk       ;
v1_nk  = simul_output.v1_simul_nk       ;
v2_nk  = simul_output.v2_simul_nk       ;
Y1_simul_nk  = simul_output.Y1_simul_nk;
Y2_simul_nk  = simul_output.Y2_simul_nk ;
Y_simul_nk  = simul_output.Y_simul_nk   ;
atY_simul_nk = simul_output.atY_simul_nk;
C_simul_nk  = simul_output.C_simul_nk   ;
A_simul_nk  = simul_output.A_simul_nk   ;
L1_simul_nk  = simul_output.L1_simul_nk ;
L2_simul_nk  = simul_output.L2_simul_nk ;
T1_simul_nk  = simul_output.T1_simul_nk ;
T2_simul_nk  = simul_output.T2_simul_nk ; 
H1_simul_nk  = simul_output.H1_simul_nk ;
H2_simul_nk  = simul_output.H2_simul_nk ;
H1_simul_nk =(H1_simul_nk>=80).*H1_simul_nk;
H2_simul_nk =(H2_simul_nk>=80).*H2_simul_nk;
kt_simul_nk = simul_output.kt_simul_nk;

Z1 = H1_simul<80;    % That's the definition we use in the data
Z2 = H2_simul<80;    % That's the definition we use in the data
Z1_nk = H1_simul_nk<80;    
Z2_nk = H2_simul_nk<80;    

age=linspace(25,65,41);

% The list of moments:
% m1: mean of H2 (including zeros) - families with young kids 
% m2: mean of H1 (including zeros) - families with young kids
% m3: mean of H2 (including zeros) - families without young kids 
% m4: mean of H1 (including zeros) - families without young kids
% m5: mean of T2 
% m6: mean of T1 
% m7: 75-50 of H2 families with young kids 
% m8: 75-50 of H2 families without young kids 
% m9: employment probabilities for wife with young kids 
% m10: employment probabilities for wife without young kids 
% m11: Change in consumption when kid is born

% read data moments and compare to them
moments= moment_match(T1_simul,T2_simul,H1_simul,H2_simul,C_simul,kt_simul,Z2,size(T1_simul,1),5,10);
warning('off','stats:regress:RankDefDesignMat');
moments_nk= moment_match(T1_simul_nk,T2_simul_nk,H1_simul_nk,H2_simul_nk,C_simul_nk,kt_simul_nk,Z2_nk,size(T1_simul_nk,1),5,10);
moments_nontarget= moment_match(T1_simul,T2_simul,H1_simul,H2_simul,C_simul,kt_simul,Z2,size(T1_simul,1),25,30);
warning('on','stats:regress:RankDefDesignMat');

simmom = [moments(1:2);...
          moments_nk(1:2);...
          moments(3:4);...
          moments(6);...
          moments_nk(6);...
          moments(5);...
          moments_nk(5);...
          moments(7);...
          ];

simnom_nt = [moments_nontarget(1:2);...
             moments_nontarget(6);...
             moments_nontarget(5);...
             ];
      
moments_data1 = csvread('moments.csv');
moments_data2 = csvread('moments_nk.csv');
moments_data3 = csvread('moments_nontarget.csv');

datamom =[moments_data1(1:2);...
          moments_data2(1:2);...
          moments_data1(3:4);...
          moments_data1(6);...
          moments_data2(6);...
          moments_data1(5);...
          moments_data2(5);...
          moments_data1(7);...
          ];
datamom_nt =[moments_data3(1:2);...
             moments_data3(6);...
             moments_data3(5);...
             ];
      
SIGMA = diag(diag(csvread('SIGMA.csv')));
diff = datamom-simmom;
F = diff'*(SIGMA^-1)*diff;

%% Store temp results for every round
moments_out=[datamom,simmom,F*ones(size(datamom))];
nontarget_moments_out=[datamom_nt,simnom_nt];
save(strcat(dirname,'/estimation_output'),'moments_out','THETA','nontarget_moments_out');
dlmwrite(strcat(dirname,'/moments_match.csv'),moments_out);
dlmwrite(strcat(dirname,'/nontarget_moments_out.csv'),nontarget_moments_out);
dlmwrite(strcat(dirname,'/THETA.csv'),THETA');

%% Save data if requested
if savedata==1  
    save(strcat(dirname,'/data'));
end

%% Generate IR upon request
if irflag==1
    hcutoff = 10^-9; % 10^-9 definition of zero hours in elasticities calculations 
    apctile = 20; % Low assets below this simulated percentile at time zero
    shock_t =6; % Time the shock hits. 6 is age 30
    % ver: 1: v1, 2: v2, 3: u1, 4: u2
    for i=1:4
        impulse_responses(parameters,wage_grids_outputs,output_vf_large,dirname,hcutoff,apctile,shock_t,i);
    end
    
    %% Generate the IR figure
    IR_figure(dirname,shock_t,TT)
    %% Generate IR final tables and insurance calculations
    IR_tables(dirname,shock_t,TT)
    
end