*! version 1.4  updated 9-02-11
// This procedure estimates the robust standard errors for meta regressions by way of various weighting schemes
// Program by Eric Hedberg (eric.hedberg@me.com) based on
// Hedges, Larry V., Elizabeth Tipton, and Matthew C. Johnson. 2010. Robust variance estimation
//     in meta-regression with dependent effect size estimates. Research Synthesis Methods.
//     (www.interscience.wiley.com) DOI: 10.1002/jrsm.5

program robumeta, eclass byable(recall)
    version 11.1
    syntax varlist(min=1 numeric) [if] [in],  ///
    [study(varlist max=1 numeric)]  ///
    [weighttype(string)] ///
    [variance(varlist max=1 numeric)]  ///
    [uweights(varlist max=1 numeric)]   ///
    [Level(cilevel)]  ///
    [rho(numlist max = 1 < 1)] ///
    [rhoforweights] [eform]

    *mark sample
    marksample touse

    *create macros of variables
    tokenize `varlist'
    local t `1'
    macro shift
    local x `*'

    *specify the temporary variables used in the program
    tempvar keepers cons w wh wfinal prelim_hat prelim_resid  ///
    prime_hat prime_resid v_mean v_n studyweight studynumber 
   

    *specifiy the temporary scalars and matrixes
    tempname A1  A2  b   B1  B2  C1  C2  D   e   E   F   I   J    ///
	k   kw  kXJWX   kXWJX   omega_squared   omega_squared_o      ///
	Q1  QE  QR  sigmahat  min_n max_n sumk    sumk2   T   T_XB    T_XBJT_XB    ///
	tau_squared tau_squared_o trW trW_1  TWT     TWX     V   VkXJWX  VkXWJX   ///
	Vw2XJX  VwkXJX_XX   VwkXX   VXJWX   VXJX    VXJXVXW2X    ///
	VXJXVXWJWX  VXW2X   VXW2X   VXWJWX  VXWJWX  VXWJX   W    ///
	w_k     w2  w2XJX   wkXJX_XX    wkXX    wXX     X   XB   ///
	XJWX    XJX     XJX     XW2X    XW2X    XWeeWX  XWJWX   XWJX     ///
	XWT     XWX     XX

    *generate constant term
    quietly : gen `cons' = 1 if `touse'

    *listwise mark and count obs
    quietly : egen `keepers' = rowmiss(`t' `x') if `touse'
    quietly : replace `touse' = 0 if `keepers' > 0
    quietly : count if `touse'
    local nobs = r(N)

    *count predictors, create macro p for number of predictors
    local p = 0
    foreach v in `x' cons {
        local ++p
    }

    capture confirm existence `study'
    if _rc == 6 {
        quietly : gen `studynumber' = _n if `touse'
    }
    else {
        quietly : gen `studynumber' = `study' if `touse'
    }

    *capture ids
    quietly : levelsof `studynumber', local(idlist)

    *count ids, create macro m
    local m = 0
    foreach j in `idlist' {
        local ++m
    }

    *set up iccs for later
    capture confirm existence `rho'
    if _rc != 6 {
        local rhot = `rho'
        capture confirm existence `rhoforweights'
        if _rc != 6 {
            local rhow = `rho'
        }
    }

    *check logic of icc options
    capture confirm existence `rhoforweights'
    if _rc != 6 {
        capture confirm existence `rho'
        if _rc == 6 {
            display as error "must specify a value of rho if you want to use rho for weights"
            exit
        }
    }

    *defualt is random weights, set if necessary
    capture confirm existence `weighttype'
    if _rc == 6 {
        local weighttype "random"
    }
    *if uweights are specified, make sure fixed is used
    capture confirm existence `uweights'
    if _rc != 6 {
    	if "`weighttype'" == "fixed" {
    		*ok
    	}
    	else {
    		di as error "need to use fixed weighing scheme for user weights"
    		exit 9	
    	}
    }	
    *figure out how user specified the weights, check arguments,
    if "`weighttype'" == "random" | "`weighttype'" == "hierarchical" | "`weighttype'" == "inverse" {
        capture confirm existence `variance'
        if _rc == 6 {
            display as error "you must specify the variance estimation with a `weighttype' weighting scheme"
            exit
        }
        else {
            *study average variance
            quietly : by `studynumber', sort rc0: egen `v_mean' = mean(`variance') if `touse'
            *number of cases per study
            quietly : by `studynumber', sort rc0: egen `v_n' = count(`variance') if `touse'
			quietly sum `v_n'
			scalar `min_n' = r(min) 
			scalar `max_n' = r(max)
            *study constant weight
            quietly : gen `studyweight' = 1 /(`v_n' * `v_mean') if `touse'
        }
    }
    else if "`weighttype'" == "fixed" {
        capture confirm existence `variance'
        if _rc == 6 {
            capture confirm existence `uweights'
            if _rc != 6 {
                capture confirm existence `rhow'
                if _rc != 6 {
                    display as error "cannot specify rho for weights with user weights"
                    exit
                }
                quietly : gen `studyweight' = `uweights' if `touse'
                quietly : by `studynumber', sort rc0: egen `v_n' = count(`uweights') if `touse'
				*number of cases per study
				quietly sum `v_n'
				scalar `min_n' = r(min) 
				scalar `max_n' = r(max)
			}
            else {
                display as error "you must specify the variance estimation or user weights with a `weighttype' weighting scheme"
                exit
            }
        }
        else {
            *study average variance
            quietly : by `studynumber', sort rc0: egen `v_mean' = mean(`variance') if `touse'
            *number of cases per study
            quietly : by `studynumber', sort rc0: egen `v_n' = count(`variance') if `touse'
            *study constant weight
            quietly : gen `studyweight' = 1 /(`v_n' * `v_mean') if `touse'
            quietly sum `v_n'
			scalar `min_n' = r(min) 
			scalar `max_n' = r(max)
        }
    }
    else {
        display as error "unknown weight specification `weighttype'"
        exit
    }

    *Calculate inverse weights, fixed weights, random effects weights, hierarchical weights

    /*****************/
    /*   Inverse     */
    /*****************/
    if "`weighttype'" == "inverse" {
        quietly : gen `wfinal' = 1/`variance' if `touse'
        capture confirm existence `rhot'`rhow'
        if _rc != 6 {
            display as error "rho for weights or tau-square has no impact on this procedure"
            exit
        }
    }
    /*****************/
    /*    Fixed      */
    /*****************/
    else if "`weighttype'" == "fixed" {
        capture confirm existence `rho'
        if _rc == 6 {
            quietly : gen `wfinal' = `studyweight' if `touse'
        }
        else {
            quietly : gen `wfinal' = 1/((1 + ((`v_n' - 1) * `rho')) * `v_mean')
        }
    }
    /*****************/
    /*  Hierarchical */
    /*****************/
    else if "`weighttype'" == "hierarchical" {
		capture confirm existence `rhot'`rhow'
        if _rc != 6 {
            display as error "rho has no impact on this procedure"
            exit
        }
        *calculate slopes for Q1

        quietly : gen `wh' = 1/`variance' if `touse'

        mkmat `wh' if `touse', matrix(`W')
        matrix `W' = diag(`W')
        mkmat `t' if `touse', matrix(`T')
        matrix colnames `T' = `t'
        mkmat `x' `cons' if `touse', matrix(`X')
        matrix `b' = inv(`X''*`W'*`X') * `X''*`W'*`T'

        matrix `XWX' = `X'' * `W' * `X'
        matrix `V' = inv(`XWX')

        *calculate matrix quantities

        matrix `TWT' = J(1, 1, 0)
        matrix `TWX' = J(1, `p', 0)
        matrix `XWX' = J(`p', `p', 0)
        matrix `XWT' = J(`p', 1, 0)

        matrix `XWJWX' = J(`p', `p', 0)
        matrix `kXJWX' = J(`p', `p', 0)
        matrix `kXWJX' = J(`p', `p', 0)
        matrix `XJX' = J(`p', `p', 0)
        matrix `XJWX' = J(`p', `p', 0)
        matrix `XWJX' = J(`p', `p', 0)
        matrix `XW2X' = J(`p',`p',0)

        matrix `T_XBJT_XB' = J(1, 1, 0)

        scalar `sumk' = 0
        scalar `sumk2' = 0

        foreach j in `idlist' {
            *different weights for each type

            mkmat `wh' if `touse' & `studynumber' == `j', matrix(`W')
            matrix `W' = diag(`W')


            quietly : levelsof `studyweight' if `touse' & `studynumber' == `j', local(sw)
            quietly : levelsof `v_n' if `touse' & `studynumber' == `j', local(sn)

            scalar `k' = `sn'
            scalar `sumk' = `k' + `sumk'
            scalar `sumk2' = (`k' * `k') + `sumk2'


            mkmat `t' if `touse' & `studynumber' == `j', matrix(`T')
            matrix colnames `T' = `t'
            mkmat `x' `cons' if `touse' & `studynumber' == `j', matrix(`X')
            matrix colnames `X' = `x' _cons
            matrix `J' = J(`sn', `sn', 1)

            matrix `TWT' = (`T'' * `W' * `T') + `TWT'
            matrix `TWX' = (`T'' * `W' * `X') + `TWX'
            matrix `XWX' = (`X'' * `W' * `X') + `XWX'
            matrix `XWT' = (`X'' * `W' * `T') + `XWT'
            matrix `XWJWX' = (`X'' * `W' * `J' * `W' * `X') + `XWJWX'
            matrix `XW2X' = (`X'' * `W' * `W' * `X') + `XW2X'

            matrix `XJWX' = (`X''*`J'*`W'*`X') + `XJWX'
            matrix `XWJX' = (`X''*`W'*`J'*`X') + `XWJX'

            matrix `kXJWX' = (`k' * (`X''*`J'*`W'*`X')) + `kXJWX'
            matrix `kXWJX' = (`k' * (`X''*`W'*`J'*`X')) + `kXWJX'

            matrix `XJX' = (`X'' * `J' * `X') + `XJX'

            matrix `XB' = `X'*`b'
            matrix `T_XB' = `T' - `XB'
            matrix `T_XBJT_XB' = (`T_XB'' * `J' * `T_XB') + `T_XBJT_XB'
        }

        *Calcuate parameters for tau and omega

        matrix `QE' = `TWT' - (`TWX' * inv(`XWX') * `XWT')
        scalar `QE' = `QE'[1,1]

        scalar `Q1' = `T_XBJT_XB'[1,1]


        matrix `VkXJWX' = `V' * `kXJWX'
        matrix `VkXWJX' = `V' * `kXWJX'
        matrix `VXJX' = `V' * `XJX'
        matrix `VXWJWX' = `V' * `XWJWX'
        matrix `VXJWX' = `V' * `XJWX'
        matrix `VXWJX' = `V' * `XWJX'
        matrix `VXW2X' = `V' * `XW2X'

        matrix `VXJXVXWJWX' = `VXJX' * `VXWJWX'

        matrix `VXJXVXW2X' = `VXJX' * `VXW2X'

        scalar `A1' = `sumk2' - trace(`VkXJWX') - trace(`VkXWJX') + trace(`VXJXVXWJWX')
        scalar `B1' = `sumk' - trace(`VXJWX') - trace(`VXWJX') + trace(`VXJXVXW2X')
        quietly : sum `variance'
        scalar `trW_1' = r(sum)
        scalar `C1' = `trW_1' - trace(`VXJX')

        quietly : sum `wh'
        scalar `trW' = r(sum)
        scalar `A2' = `trW' - trace(`VXWJWX')
        scalar `B2' = `trW' - trace(`VXW2X')
        scalar `C2' = `sumk' - `p'

        *tau and omega

        scalar `omega_squared' = ((`A2'*(`Q1'-`C1')) - (`A1'*(`QE'-`C2')))/((`B1'*`A2') - (`B2'*`A1'))

        if `omega_squared' < 0 {
            scalar `omega_squared_o' = `omega_squared'
            scalar `omega_squared' = 0
        }
		else {
			scalar `omega_squared_o' = `omega_squared'
		}
		scalar `tau_squared' =   ((`QE'-`C2')/`A2')-(`omega_squared'*(`B2'/`A2'))

        if `tau_squared' < 0 {
            scalar `tau_squared_o' = `tau_squared'
            scalar `tau_squared' = 0
        }
		else {
			scalar `tau_squared_o' = `tau_squared'
		}

        *Calcuate weights

        quietly : gen `wfinal' = 1/(`variance' + `omega_squared' + `tau_squared')
    }
    /*****************/
    /*    Random     */
    /*****************/
    else if "`weighttype'" == "random" {
        *if no rho is specified, value of 1 is assumed
        capture confirm existence `rho'
        if _rc == 6 {
            local rhot = 1
        }
        capture confirm existence `rhow'
        if _rc == 6 {
            quietly : gen `w' = `studyweight' if `touse'
        }
        else {
            quietly : gen `w' = 1/((1 + ((`v_n' - 1) * `rhow')) * `v_mean') if `touse'
        }

        *calculate matrix quantities

        matrix `TWT' = J(1, 1, 0)
        matrix `wXX' = J(`p', `p', 0)
        matrix `TWX' = J(1, `p', 0)
        matrix `XWX' = J(`p', `p', 0)
        matrix `XWJWX' = J(`p', `p', 0)
        matrix `XW2X' = J(`p',`p',0)
        matrix `XWT' = J(`p', 1, 0)
        matrix `wkXX' = J(`p',`p',0)
        matrix `wkXJX_XX' = J(`p', `p', 0)
        matrix `w2XJX' = J(`p', `p', 0)
        scalar `kw' = 0

        foreach j in `idlist' {
            mkmat `w' if `touse' & `studynumber' == `j', matrix(`W')
            quietly : levelsof `studyweight' if `touse' & `studynumber' == `j', local(sw)
            quietly : levelsof `v_n' if `touse' & `studynumber' == `j', local(sn)

            scalar `w_k' = `sw' / `sn'
            scalar `w2' = `sw'^2

            mkmat `t' if `touse' & `studynumber' == `j', matrix(`T')
            matrix colnames `T' = `t'
            mkmat `x' `cons' if `touse' & `studynumber' == `j', matrix(`X')
            matrix colnames `X' = `x' _cons
            matrix `J' = J(`sn', `sn', 1)
            matrix `W' = diag(`W')
            matrix `TWX' = (`T'' * `W' * `X') + `TWX'
            matrix `TWT' = (`T'' * `W' * `T') + `TWT'
            matrix `wXX' = (`sw' * `X'' * `X') + `wXX'
            matrix `XWX' = (`X'' * `W' * `X') + `XWX'
            matrix `XWT' = (`X'' * `W' * `T') + `XWT'
            matrix `XWJWX' = (`X'' * `W' * `J' * `W' * `X') + `XWJWX'
            matrix `XW2X' = (`X'' * `W' * `W' * `X') + `XW2X'
            matrix `wkXX' = (`w_k' * (`X'' * `X')) + `wkXX'
            matrix `wkXJX_XX' = (`w_k' * ((`X'' * `J' * `X') - (`X'' * `X'))) + `wkXJX_XX'
            matrix `w2XJX' = (`w2' * `X'' * `J' * `X' ) + `w2XJX'
            scalar `kw' = (`sw' * `sn') + `kw'
        }

        *Calcuate parameters for tau

        matrix `V' = inv(`XWX')

        matrix `QE' = `TWT' - (`TWX' * inv(`XWX') * `XWT')
        scalar `QE' = `QE'[1,1]

        matrix `Vw2XJX' = `V' * `w2XJX'
        scalar `D' = trace(`Vw2XJX')

        matrix `VwkXX' = `V' * `wkXX'
        scalar `E' = trace(`VwkXX')

        matrix `VwkXJX_XX' = `V' * `wkXJX_XX'
        scalar `F' = trace(`VwkXJX_XX')

        *Calcuate weights

        scalar `tau_squared' = (`QE' - `m' + `E' + (`rhot' * `F')) / (`kw' - `D')
        if `tau_squared' < 0 {
            scalar `tau_squared_o' = `tau_squared'
            scalar `tau_squared' = 0
        }
		else {
			scalar `tau_squared_o' = `tau_squared'
		}
        quietly : gen `wfinal' = 1/(`v_n' * (`v_mean' + `tau_squared'))
    }

    /*****************/
    /*    Slopes     */
    /*****************/
    matrix `XWX' = J(`p', `p', 0)
    matrix `XWT' = J(`p', 1, 0)

    foreach j in `idlist' {
        mkmat `t' if `touse' & `studynumber' == `j', matrix(`T')
        matrix colnames `T' = `t'
        mkmat `x' `cons' if `touse' & `studynumber' == `j', matrix(`X')
        matrix colnames `X' = `x' _cons
        mkmat `wfinal' if `touse' & `studynumber' == `j', matrix(`W')
        matrix `W' = diag(`W')
        matrix `XWX' = (`X'' * `W' * `X') + `XWX'
        matrix `XWT' = (`X'' * `W' * `T') + `XWT'
    }
    matrix `b' = (inv(`XWX') * `XWT')'
    matrix colnames `b' = `x' _cons
    matrix score `prime_hat' = `b'
    gen `prime_resid' = `t' - `prime_hat'

    *Variance covariance matrix estimation for standard errors
    matrix `XWeeWX' = J(`p', `p', 0)
    foreach j in `idlist' {
        mkmat `wfinal' if `touse' & `studynumber' == `j', matrix(`W')
        matrix `W' = diag(`W')
        mkmat `x' `cons' if `touse' & `studynumber' == `j', matrix(`X')
        matrix colnames `X' = `x' _cons
        mkmat `prime_resid' if `touse' & `studynumber' == `j', matrix(`e')

        matrix `sigmahat' = `e' * `e''

        matrix `XWeeWX' = (`X'' * `W' * `sigmahat' * `W' * `X') + `XWeeWX'
    }
    matrix `V' = inv(`XWX') * `XWeeWX' * inv(`XWX')
    local df = `nobs' - `p'

    display _newline
    display as text "Robust standard error estimation using " as result "`weighttype'" as text " model weights"

    /*****************************/
    /*  small sample adjustment  */
    /*****************************/
    local df = `m' - `p'
    matrix `V' = (`V') * (`m' / (`m'-`p'))

    *name the rows and columns of the matrixes

    matrix colnames `V' = `x' _cons
    matrix rownames `V' = `x' _cons

    *post results
    
    ereturn local depvar "`t'"
    ereturn local cmd "gmeta"
    ereturn scalar N_g = `m'

    /*********************/
    /*  Display results  */
    /*********************/

    display _col(55) as text "N Level 1" _col(69) "=" _col(69) as result %9.0f `nobs'
    display _col(55) as text "N Level 2" _col(69) "=" _col(69) as result %9.0f `m'
	display _col(55) as text "Min Level 1 n" _col(69) "=" _col(69) as result %9.0f `min_n'
	display _col(55) as text "Max Level 1 n" _col(69) "=" _col(69) as result %9.0f `max_n'
    display _col(55) as text _col (5) "Average" _col(69) "=" _col(69) as result  %9.2f `nobs' / `m'

    display _col(55) as text "T-Test DF" _col(69) "=" _col(69) as result  %9.0f `df'

    capture confirm existence `rho'
    if _rc != 6 {
        display _col(55) as text  "Assumed rho" _col(69) "=" _col(69) as result  %9.2f `rhot'
    }

    if "`weighttype'" == "fixed" {
    }
    else if "`weighttype'" == "random" {
        display _col(55) as text  "tau-squared" _col(69) "=" _col(69) as result  %9.4f `tau_squared'
    }
    else if "`weighttype'" == "hierarchical" {
        display _col(55) as text  "tau-squared" _col(69) "=" _col(69) as result  %9.4f `tau_squared'
        display _col(55) as text  "omega-squared" _col(69) "=" _col(69) as result  %9.4f `omega_squared'
    }

    *post values of rho
    capture confirm existence `rho'
    if _rc != 6 {
        ereturn scalar rho = `rho'
    }
	ereturn post `b' `V', dof(`df') obs(`nobs') depname(`t') esample(`touse')
	*display results
	if "`eform'" == "eform" {
		ereturn display , level(`level') eform("exp(b)") 
	}
	else {
		 ereturn display , level(`level')
	}
	if "`weighttype'" == "fixed" {
    }
    else if "`weighttype'" == "random" {
        ereturn scalar tau2 = `tau_squared'
        ereturn scalar tau2o = `tau_squared_o'
        ereturn scalar QE = `QE'
    }
    else if "`weighttype'" == "hierarchical" {
        ereturn scalar tau2 = `tau_squared'
        ereturn scalar tau2o = `tau_squared_o'
        ereturn scalar omega2 = `omega_squared'
        ereturn scalar omega2o = `omega_squared_o'
        ereturn scalar QE = `QE'
        ereturn scalar Q1 = `Q1'
    }
end