/*
build_cond_expectation_2.ado performs the following operation:
1) It computes the conditional expectations E(Y_d0t|X) that are
needed to compute the Wald-TC estimator with covariates, when this estimator is
requested by the user*/

quietly capture program drop build_cond_expectation_2
quietly program build_cond_expectation_2

	version 12

	syntax varlist(max=4 numeric) ///
	[, inf_method(name) reg_method(name) ///
	continuous(varlist num fv) qualitative(varlist num fv) ///
	sieve_expansion(varlist num) d(numlist) target0(name) ///
	target1(name)]

	tokenize `varlist', parse(" ,")

	args Y G T D
	*preserve

	quietly {

		tempname mean_outcome coefs sd_coefs

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

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

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

		local covariates "`quali_list' `continuous'"


		forvalues i=0(1)1 {

			/*computation of unconditional expectations of the form
			E(Y_d0i)*/
			sum `Y' if `G'==0 & `T'==`i' ///
			& `D'==`d'
			scalar `mean_outcome'=r(mean)

			/*conditional expectation E(Y_d0i|X)*/

			if r(min)==r(max) {

				gen `target`i''=`mean_outcome' if `G'==1 & `T'==1
			}
			else {

				if "`inf_method'"=="param" {

					`reg_method' `Y' `covariates' if `D'==`d' & `G'==0 ///
					& `T'==`i'
					matrix `coefs'=e(b)
					matrix `sd_coefs'=e(V)
				}
				else {

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

						reg `Y' `quali_list' `sieve_expansion' ///
						if `D'==`d' & `G'==0 & `T'==`i'
						matrix `coefs'=e(b)
						matrix `sd_coefs'=e(V)
					}
				}

				/*evaluation of regression function E(Y_d0i|X)
				at values of X in subgroup {G=1,T=1}*/
				predict `target`i'' 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`i''=`mean_outcome' if `target`i''==. & `G'==1 & `T'==1
				}
				else {

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

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

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

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


			/*check if support of some qualitative covariates larger on
			{G=1,T=1} than on {G=0,T=i,D=d}*/
			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 distinct values the variable takes in pop
				{G=1,T=1}*/
				levelsof `quali_var`l'' if `G'==1 ///
				& `T'==1, local(coucou1)

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

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

				/*if there are actually values that appear in {G=1,T=1} and
				not in {D=d,G=0,T=i}, 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`i''=`mean_outcome' if `patch_var'!=0
		}
	}
end