/*
build_cond_expectation_1.ado performs the following operation:
1) It computes the conditional expectations E(D_gt|X) and E(Y_gt|X) that are
needed to compute the Wald-DID and Wald-TC estimators with covariates, when these
estimators are requested by the user*/

quietly capture program drop build_cond_expectation_1
quietly program build_cond_expectation_1

	version 12

	syntax varlist(max=3 numeric) ///
	[, inf_method(name) reg_method(name) ///
	continuous(varlist num fv) qualitative(varlist num fv) ///
	sieve_expansion(varlist num) i(numlist) j(numlist) target(name)]

	tokenize `varlist', parse(" ,")

	args outcome G T

	quietly {

		tempname mean_outcome coefs sd_coefs

		local n_quali=wordcount("`qualitative'")
		local quali_list ""

		/*include qualitative variables in the list of covariates*/
		forvalues l=1(1)`n_quali' {

			local a: word `l' of `qualitative'
			local quali_list "`quali_list' i.`a'"
		}

		local covariates "`quali_list' `continuous'"

		/*computation of unconditional expectations of the form
		E(D_gt) or E(Y_gt)*/
		sum `outcome' if `G'==`i' & `T'==`j', meanonly
		scalar `mean_outcome'=r(mean)

		/*computation of conditional expectations of the form
		E(D_ij|X) or E(Y_ij|X), according to computation method chosen by user*/
		if "`inf_method'"=="param" {

			`reg_method' `outcome' `covariates' if `G'==`i' & ///
			`T'==`j'
			matrix `coefs'=e(b)
			matrix `sd_coefs'=e(V)
		}
		else {

			if "`inf_method'"=="sieve" {

				reg `outcome' `quali_list' `sieve_expansion' ///
				if `G'==`i' & `T'==`j'
				matrix `coefs'=e(b)
				matrix `sd_coefs'=e(V)
			}
		}

		/*evaluation of regression function E(D_ij|X) or E(Y_ij|X)
		(ij different from 11) at values of X in subgroup {G=1,T=1}*/
		predict `target' if `G'==1 & `T'==1

		/*if prediction is missing for some observations in {G=1,T=1},
		regression function=unconditional mean for those observations*/
		if "`inf_method'"=="sieve" | "`reg_method'"=="reg" {

			replace `target'=`mean_outcome' if `target'==. & `G'==1 & `T'==1
		}
		else {

			if "`reg_method'"=="logit" {

				replace `target'=`mean_outcome' if `target'==. & `G'==1 & `T'==1
			}

			if "`reg_method'"=="probit" {

				replace `target'=`mean_outcome' if `target'==. & `G'==1 & `T'==1
			}
		}

		/*check if support of some qualitative covariates larger on
		{G=1,T=1} than on {G=i,T=j}*/
		tempvar patch_var

		gen `patch_var'=0

		/*loop over the list of qualitative regressors*/
		forvalues l=1(1)`n_quali' {

			/*identify qualitative regressor at step l of the loop*/
			local quali_var`l': word `l' of `qualitative'

			/*count how many values the variable takes in pop
			{G=1,T=1}*/
			quietly levelsof `quali_var`l'' if `G'==1 ///
			& `T'==1, local(coucou1)

			/*count how many values the variable takes in pop
			{G=i,T=j}*/
			levelsof `quali_var`l'' if `G'==`i' ///
			& `T'==`j', local(coucou2)

			/*count how many values of the variable appear in {G=1,T=1} without
			appearing in {G=i,T=j}*/
			local both : list coucou1 - coucou2
			local count=wordcount("`both'")

			/*if there are actually values that appear in {G=1,T=1} and
			not in {G=i,T=j}, identify all the obs in {G=1,T=1} for which
			the variable takes one of those problematic vales*/
			if `count'!=0 {

				forvalues m=1(1)`count' {
					local word`m': word `m' of `both'
					replace `patch_var'=`patch_var'+1 if ///
					`quali_var`l''==`word`m'' & `G'==1 ///
					& `T'==1
				}
			}
		}

		/*for observations in {G=1,T=1} that were identified as problematic,
		regression function=unconditional mean*/
		replace `target'=`mean_outcome' if `patch_var'!=0

		drop `patch_var'
	}
end