*! version 0.0 updated 02-March-2017
// Update by Marcelo Tyszler (tyszler.jobs@gmail.com):
//
// Post-estimation function for reg_sandwich
//

capture program drop test_sandwich
program define test_sandwich, eclass byable(recall) sortpreserve
	version 14.2


    set type double
    syntax [varlist(default=none)], [cons]

	tempname cluster_list

	tempvar Fconstant ///
			 clusternumber ///
			 temp_Ftest ///
			 selectvar ///
			 Q_Ftest

	*verify that this is run after reg_sandwich:

	if e(cmd) !="reg_sandwich" {
		display as error "{it:test_sandwich} can only be used after {it:reg_sandwich}"
		error 301
		exit
	}

	* create placeholder for mata selection
	qui: gen `selectvar'=.

	* F-test:
	*
	* The Q statistic needs the definition of C matrix and a c vector
	*
	* p is the number of coefficients (including the constant)
	* q is the number of coefficients to be tested (1 <= q <= p)
	*
	* C is the contrast matrix (q x p) such that:
	*  H0: Cb = c
	*  For example for the test beta_s = 0, would have q=1,  c = 0 and C = [ 0 .. 1 .. 0], where entry s (between 1 and p) is 1 and 0 otherwise
	*
	*  For the F-test, C will be a (q x p) matrix where entry rs (1 <= r <= q and 1 <= s <=p) will be 1 and 0 otherwise
	*  and c will be a (q x 1) vector of 0s
	*
	*  For example if p = 3, an F test of beta_1 = beta_3 = 0 would have c = [0 0]' and C = [1 0 0; 0 0 1]


	*count coefficients to be tested,
	* Simultaneously check if they belong to the list of original coefficients
	*load original variables:
	local x = e(indepvars)
	local type_VCR = e(type_VCR)
	local constant_used = e(constant_used)

	local q_Ftest = 0
	capture confirm existence `varlist'
	if _rc != 6 {
		foreach current_x in `varlist' {
			* verify coefficient is in the original list
			if strpos("`x'","`current_x'")==0{
				display as error "F-test error: {it:`current_x'} does not belong to the list of coefficients from the {it:robumeta} estimation"
				error 101
				exit
			}

			local ++q_Ftest
		}
	}
	* verify constant term:
	* remove leading and trailing whitespace
	if "`cons'" != "" {
		if "`constant_used'" == "0" {
			display as error "Constant was not included in the estimation. It cannot be included in the tests."
			error 101
			exit
		}
		* increment q:
		local ++q_Ftest
	}

	* Define C
	local p = 0
    foreach v in `x' {
        local ++p
    }

	if "`constant_used'" == "1" {
		* increment q:
		local ++p
	}

	matrix `temp_Ftest' = J(1, `p', 0) // C is initialized as a 1 x p matrix of zeros
	mata: C_Ftest = J(`q_Ftest', `p', 0) // C is initialized as a q x p matrix of zeros

	if "`constant_used'" == "1" {
		matrix colnames `temp_Ftest' = `x' _cons
	}
	else {
		matrix colnames `temp_Ftest' = `x'
	}

	local current_row = 1
	capture confirm existence `varlist'
	if _rc != 6 {
		* for each var listed in ftest, check which column it corresponds to
		foreach current_q in `varlist'{
			local coln = colnumb(`temp_Ftest',"`current_q'")
			mata: C_Ftest[`current_row', `coln'] = 1
			local ++current_row
		}
	}

	* If option constant is active, last column needs to be active:
	if "`cons'" != "" {
		mata: C_Ftest[`current_row', `p'] = 1
		local ++current_row
	}
	mata : C_Ftest = sort(C_Ftest, -1..-`p')

    * F-test:
	*
	* To compute the degress of freedom we need P:
	* Psi = (I-Hx)i'*Ai*Wi*Xi*M*C*gs
	*
	* These matrices are needed to compute the terms Psi'*Theta*Ptj:
	*  gs'*C'*M*Xi'*Wi*Ai*(I-Hx)i*Theta*(I-Hx)j'*Aj*Wj*Xj*M*C*gt
	*
	* We saved the "middle" portion, which is independent of C and gs:
	*
	* Using the fact that Hx = X*M*X'W and
	* (I-X*M*X'*W)i*T*(I-X*M*X'*W)j' =
	*
	* if i==j
	* Tj - Tj*(Wj*Xj*M*Xj') - (Xj*M*Xj'*W)*Tj + Xj*(M*X'*W*V*W*X*M)*Xj'
	*
	* For OLS this simplifies to:
	* Tj - Xj*M*Xj'
	*
	* For WLSp, this simplifies to (Dj = I):
	* Tj - Wj*Xj*M*Xj' - Xj*M*Xj'Wj + Xj'MXWWXM*Xj'
	*
	* For WLSa, this simplified to:
	* Tj - Xj*M*Xj
	*
	*
	* and we call M*Xi'*Wi*Ai*(I-Hx)i*Theta*(I-Hx)j'*Aj*Wj*Xj*M:
	* Pi_Theta_Pi_relevant
	*
	*
	* if i!=j
	* - Ti*Wi*Xi*M*Xj'   - Xi*M*Xj'*Wj*Tj     + Xi*(M*X'*W*T*W*X*M)*Xj'
	*
	* For OLS this simplifies to:
	* - Xi*M*Xj'
	*
	* For WLSp, this simplifies to:
	* - Wi*Xi*M*Xj'   - Xi*M*Xj'*Wj     + Xi*(M*X'*W*W*X*M)*Xj'
	*
	* For WLSa, this simplified to:
	* - Xi*M*Xj'
	*
	* For OLS and WLSa we call M*Xi'*Wi*Ai*Xi:
	* Pi_relevant (and ignore the (min) sign, since it will be cancelled out after multiplication)
	*
	* For WLSp we call  M*Xi'*Wi*Ai
	* Pi_Pj_relevant, (this is more efficient to save)
	*
	* and additionally save M*Xi'*Wi*Ai as PPi

	*

	local m = e(N_clusters)

	if "`type_VCR'" == "WLSp" {

		if "`constant_used'" == "1" {
			quietly : gen double `Fconstant' = 1 if e(sample)
		}

		local cluster = e(clustvar)
		capture confirm numeric variable `cluster'
		if _rc==0 {
			* numeric
		   quietly : gen double `clusternumber' = `cluster' if e(sample)
		}
		else {
			* string
			quietly: encode `cluster' if e(sample), gen(`clusternumber')
		}

		quietly sort `clusternumber' `_sortindex'
		qui: tab `clusternumber' if e(sample), matrow(`cluster_list')


		local endi = 0

		forvalues i = 1/`m'{

			local starti = `endi'+1

			if "`e(absorb)'"~=""{
				qui: sum  `x' if e(sample) & `clusternumber' == `cluster_list'[`i',1]
				local endi  =  `starti' + r(N) -1
				mata: X`i' = Ur[`starti'..`endi',1..`p']
			}
			else {

				qui: replace `selectvar' = e(sample) & `clusternumber' == `cluster_list'[`i',1]
				mata: X`i' = .

				if "`constant_used'" == "1" {
					mata: st_view(X`i', ., "`x' `Fconstant'","`selectvar'")

				}
				else {
					mata: st_view(X`i', ., "`x'","`selectvar'")
				}

				mata: st_local("rows_number", strofreal(rows(X`i')))
				local endi  =  `starti' + `rows_number' - 1
			}


			mata: PP`i' = Big_PP[`starti'..`endi',1..`p']
			mata: P`i'_relevant = Big_P_relevant[`starti'..`endi',1..`p']'
		}

	}

	mata: Omega_Ftest = C_Ftest*MXWTWXM*C_Ftest'

	* Symmetric square root of the Moore-Penrose inverse of Omega_Ftest
	mata: evecs = .
	mata: evals = .
	mata: symeigensystem(Omega_Ftest, evecs, evals)
	mata: sq_Omega_Ftest =  evecs*diag(editmissing(evals:^(1/2),0))*evecs'

	mata: matrix_Ftest = invsym(sq_Omega_Ftest)

	mata: st_local("Sum_temp_calc2", test_sandwich_ftests("`type_VCR'", `q_Ftest', `m', `p', Big_PThetaP_relevant,  Big_P_relevant,  MXWTWXM,  matrix_Ftest,  C_Ftest))


	if "`type_VCR'" == "WLSp" {
		forvalues i = 1/`m'{

			mata: mata drop X`i'
			mata: mata drop PP`i'
			mata: mata drop P`i'_relevant
		}
	}

	* eta needs to be computed according to equation (14):
	* eta = q*(q+1) / [sum(s=1 to q) sum(t=1 to q) Var(d_st)]

	local eta_Ftest = (`q_Ftest'*(`q_Ftest'+1))/`Sum_temp_calc2'

	* z = Omega^(-1/2)(Cb-c)
	* D = Omega^(-1/2)*C*VR*C'*Omega^(-1/2)
	* Q = z'inv(D)z (equation 6)

	mata: b = st_matrix("e(b)")
	mata: V = st_matrix("e(V)")

	mata: z_Ftest = invsym(sq_Omega_Ftest)*(C_Ftest*b')
	mata: D_Ftest = invsym(sq_Omega_Ftest)*C_Ftest*V*C_Ftest'*invsym(sq_Omega_Ftest)
	mata: st_matrix("`Q_Ftest'", z_Ftest'*invsym(D_Ftest)*z_Ftest)

	* Now we can compute the F-statistic:
	* (eta - q + 1)/(eta*q) * Q  follows F(q, eta - q + 1) distribution
	local F_stat = ((`eta_Ftest' - `q_Ftest' + 1)/(`eta_Ftest'*`q_Ftest'))* `Q_Ftest'[1,1]
	local F_df1 = `q_Ftest'
	local F_df2 = `eta_Ftest' - `q_Ftest' + 1

	local F_pvalue = Ftail(`F_df1',`F_df2',`F_stat')


	* F-test:
	* display some results
	display in b _newline
	display in b  "Small Sample Corrected F-test:"
	display _col(10) in b  "F(" as result %5.4f `F_df1' "," as result %5.4f `F_df2' ")" _col(30) "=" _col(35) as result  %5.4f `F_stat'
	display _col(10)  "Prob > F" _col(30) "=" _col(35) as result  %5.4f `F_pvalue'


	* F-test:
	* save some results
	ereturn scalar F_stat = `F_stat'
	ereturn scalar F_df1 = `F_df1'
	ereturn scalar F_df2 = `F_df2'
	ereturn scalar F_pvalue = `F_pvalue'
	ereturn scalar F_eta = `eta_Ftest'


	* Clean:
	mata: mata drop sq_Omega_Ftest
	mata: mata drop C_Ftest
	mata: mata drop D_Ftest
	mata: mata drop Omega_Ftest
	mata: mata drop V
	mata: mata drop b
	mata: mata drop evals
	mata: mata drop evecs
	mata: mata drop matrix_Ftest
	mata: mata drop z_Ftest


end