/*
five_fold_cv.ado performs the following operation:
1) It selects the polynomial order in the series approximation to be used to 
estimate the conditional expectations of Y and D, when the user wishes to estimate
the Wald-DID and/or Wald-TC estimator with covariates, and specifies the sieve 
option, but not the sieve_order one. It calls legendrisation.ado to compute the 
Legendre polynomials in X to be used to estimate the conditional expectations 
of Y and D*/

capture program drop five_fold_cv

program five_fold_cv, sortpreserve

	version 12
	
	syntax varlist(min=3 max=4 numeric) [, continuous(varlist num) ///
	qualitative(varlist numeric) cluster(varlist num) sieve_order(name)]
	
	tokenize `varlist', parse(" ,")		 
	marksample touse
	
	local varlist_count=wordcount("`varlist'")
	
	/*rename outcome, group, time and treatment*/
	if `varlist_count'==3 {
	
		args outcome G T
	}
	
	if `varlist_count'==4 {
	
		args outcome Gb Gf T
	}
		
	preserve

	
		
	quietly {
		
		keep if `touse'	
		tempname mat_cv
		tempvar u T_init T_no_init
		local ncont=wordcount("`continuous'")
		local nquali=wordcount("`qualitative'")
		scalar `sieve_order'=0
		quietly count 
		local N_obs=r(N)	
	
		/*Create random partitions of the data into 5 groups:*/
		
		local K=1
		
		forvalues i=1(1)`K' {
		
			set seed `i'
	
			if "`cluster'"=="" {
				
				tempvar no_clust_`i'
				gen `u'=runiform()
				sort `u'	
				gen `no_clust_`i'' = mod(_n-1,5) + 1
				drop `u'
			}
			
			/*if there is clustering in the data, partition at the cluster level*/
			
			else {
				
				tempvar clust_`i'
				gen `u'=runiform()
				sort `u'
				bysort `cluster': gen `clust_`i''=cond(_n==1,runiform(),.)
				replace `clust_`i''=cond(`clust_`i''<0.2,1, ///
				cond(`clust_`i''<0.4,2,cond(`clust_`i''<0.6,3, ///
				cond(`clust_`i''<0.8,4,cond(`clust_`i''<1,5,.)))))
				sort `cluster' `clust_`i''
				replace `clust_`i''=`clust_`i''[_n-1] if `clust_`i''>=.
				drop `u'
			}
		}
		
		local iter=1
		local time_to_stop=0
		matrix `mat_cv'=J(10,1,0)
		
		if `varlist_count'==4 {
		
			gen `T_init'=(`T'==1)
			gen `T_no_init'=(`T'>1)
		}
		
		
		while `iter'!=0 {
		
			local iter2=`iter'+4
			local iter3=`iter'
			
			forvalues i=`iter'(1)`iter2' {
				
				local n_new_vars=comb(`ncont'+`i',`i')
				
				/*if at a certain order, too many polynomial terms, error*/
				if `n_new_vars'>min(4800,`N_obs'/5) {
				
					scalar `sieve_order'=.
					ereturn clear
					local time_to_stop=1
					continue, break
				}
				
				/*create polynomials at current order*/
				local varz ""
				
				forvalues j=1(1)`n_new_vars' {
				
					tempname v`j'
					local varz "`varz' `v`j''"
				}
				
				legendrisation, continuous_var(`continuous') ///
				order1(`i') new_vars(`varz')
				
				/*regress, compute standard errors and average those for each partition*/
	
				forvalues j=1(1)`K' {
					
					local temp_CV=0
					
					forvalues l=1(1)5 {
						
						tempvar outcome_hat
						
						if "`cluster'"=="" {
							
							if `varlist_count'==3 {
							
								reg `outcome' i.`G'#i.`T'#c.(`varz') ///
								i.`G'#i.`T'#i.(`qualitative') ///
								if `no_clust_`j''!=`l', noc
							}
							else {
							
								reg `outcome' i.`T_no_init'#i.`Gb'#i.`T'#c.(`varz') ///
								i.`T_no_init'#i.`Gb'#i.`T'#i.(`qualitative') ///
								i.`T_init'#i.`Gf'#i.`T'#c.(`varz') ///
								i.`T_init'#i.`Gf'#i.`T'#i.(`qualitative') ///
								if `no_clust_`j''!=`l', noc							
							}
							
							predict `outcome_hat' if `no_clust_`j''==`l'
						}
						else {
							
							if `varlist_count'==3 {
							
								reg `outcome' i.`G'#i.`T'#c.(`varz') ///
								i.`G'#i.`T'#i.(`qualitative') ///
								if `clust_`j''!=`l', noc
							}
							else {
							
								reg `outcome' i.`T_no_init'#i.`Gb'#i.`T'#c.(`varz') ///
								i.`T_no_init'#i.`Gb'#i.`T'#i.(`qualitative') ///
								i.`T_init'#i.`Gf'#i.`T'#c.(`varz') ///
								i.`T_init'#i.`Gf'#i.`T'#i.(`qualitative') ///
								if `clust_`j''!=`l', noc								
							}
							
							predict `outcome_hat' if `clust_`j''==`l'						
						}
						
						replace `outcome_hat'=(`outcome_hat'-`outcome')^2
						sum `outcome_hat'
						local temp_CV=`temp_CV'+(1/`N_obs')*r(sum)
	
						drop `outcome_hat'						
					}
					
					matrix `mat_cv'[`i',1]=((`j'-1)/`j')*`mat_cv'[`i',1]+ ///
					(1/`j')*`temp_CV'					
				}
				
				drop `varz'
				
				/*see whether current order better than best so far in terms of MSE*/
				if `mat_cv'[ `i',1]<=`mat_cv'[`iter3',1] {
				
					local iter3=`i'
				}
			}
			
			if `time_to_stop'==1 {
			
				continue, break
			}
			
			if `iter3'==`iter2' {
				
				/*order larger than 10 not allowed*/
				if `iter3'==10 {
				
					/*display as err _newline(1) ///
					"max power reached" _newline(1)*/
					local iter=0
					scalar `sieve_order'=10
				}
				/*if MSE decreased at each order and order=10 not reached, try 5 more
				orders*/
				else {
					
					local iter=`iter'+5
				}
			}
			else {
			
				local iter=0
				scalar `sieve_order'=`iter3'
			}
		}
			
		ereturn clear
	}
end