*! version 1.1.2  31jan2000   (STB-55: sg141)
program define treatreg, eclass
	version 6.0

	if replay() {
		if "`e(cmd)'" != "treatreg" {
			noi di in red "results of treatreg not found"
			exit 301
		}
		if "`e(method)'"== "two-step" {
			Display2 `0'
		}
		else    Display `0'
		exit
	}

	if _caller() < 6.0 {
		di in red "treatreg is for version 6 only"
		exit 198
	}
	Estimate `0'
end

program define Estimate, eclass
	version 6
	gettoken dep 0 : 0 , parse(" =,[")
	unab dep : `dep'
	confirm variable `dep'
	gettoken equals rest : 0, parse(" =")

	if "`equals'" == "=" { local 0 `"`rest'"' }
	 
				/* Primary */
	syntax [varlist(default=none)] [if] [in] [pw fw iw aw]    	/*
		*/ , TReat(string)   [ HAzard(string)			/* 
		*/ noConstant FIRst Level(integer $S_level)           /*
		*/ MLOpts(string) noLOG Robust  Cluster(varname)	/*
		*/ TWOstep noSKIP * ]		

				/* set up command options */

	local ind `varlist'
	local if0 `if'
	local in0 `in'
	local option0 `options'
	local nc `constant'
	local wtype `weight'
	local wtexp `"`exp'"'
	if "`weight'" != "" { local wgt `"[`weight'`exp']"'  }
	
	if "`weight'" == "pweight" | "`cluster'" != "" {
		local robust "robust"
	}

	if "`hazard'" != "" { confirm new var `hazard' }
	if "`cluster'" != "" { local clusopt "cluster(`cluster')" }

	if "`log'" == "" {
		local showlog noisily
	}
	else local showlog quietly
        if "`skip'" != "" {
                if "`robust'" != "" {
di in blue "model LR test inappropriate with robust covariance estimates"
			local skip
                }
                if "`constan'" != "" {
di in blue "model LR test inappropriate with noconstant option"
                        local skip
                }
                if "`ind'" == "" {
di in blue "model LR test inappropriate with constant-only model"
                      local skip
                }
                if "`twostep'" != "" {
di in blue "model LR test inappropriate for two-step model"
                       local skip
                }
                if "`skip'" == "" {
                       di in blue "    performing Wald test instead"
		}
	}
				/* Parse ML options */

	mlopts stdopts, `option0'

				/* Process selection equation */
	tokenize `"`treat'"', parse(" =")
	if "`2'" ~= "=" {
		di in red "invalid syntax"
		exit 198
	}
	unab 1 : `1'
	local trtdep `1'
	local 1 " "
	local 2 " "
	local 0 `*'
	syntax [varlist(default=none)] [, noConstant ] 			

				/* set up selection eq options */
	local trtind `varlist'
	local trtnc `constant'
	 
				/* ensure valid selection eq */
	if "`trtnc'" != "" & "`trtind'" ==""{
		noi di in red "no variables specified for treatment" 
		exit 198
	}

				/* ensure valid main eq */
	if "`nc'" != "" & "`ind'" == "" {
		noi di in red "no variables pecified for primary equation"
		exit 198
	}

				/* mark sample */

	tempname touse 
	mark `touse' `wgt' `if0' `in0'
	markout `touse' `dep' `ind' `trtdep' `trtind'

				/* ensure trtvar are bivariate 1 0 */
	tempname val1 val2
	qui sum `trtdep' if `touse' 
	if r(N) == 0 {         
		di in red "no observations for treatment equation"
		exit 198
	}
	scalar `val1' = r(min)
	scalar `val2' = r(max)
	if `val1' == `val2' {
		di in red "treatments does not vary"
		exit 2000
	}

	qui count if  `trtdep' != 1 & `trtdep' != 0 & `touse'
	if r(N) != 0 {
		di in red "invalid treatment values. only 0 and 1 are allowed"
		exit 420 
	}
	 
				/* test collinearity */
	_rmcoll `trtind' `wgt' if `touse', `trtnc'
	local trtind "`r(varlist)'"
	_rmdcoll `dep' `ind' `wgt' if `touse', `nc'			
	local ind "`r(varlist)'"

				/* model estimation */

	TwoStep "`dep'" "`ind'" "`nc'" "`trtdep'" "`trtind'" /*
		*/ "`trtnc'" "`wgt'" "`wtype'" "`wtexp'"     /*
		*/ `touse' "`first'" "`hazard'"  


	if "`twostep'" == "" {

		tempname b1 athrho lnsigma llreg llprob
		scalar `llprob' =e(ll_p)

				/* initial values */
		mat `b1' = e(b)
		mat `b1'= `b1'[1, 1..colsof(`b1')-1]
		scalar `athrho' = max(min(e(rho), .85), -.85)
		scalar `athrho' = 1/2*ln((1+`athrho')/(1-`athrho'))
		scalar `lnsigma' = ln(e(sigma))
		mat `b1' = `b1', `athrho', `lnsigma'

		local hazard `e(hazard)'
		qui reg `dep' `ind' `trtdep' `wgt' if `touse', `nc'
		scalar `llreg' = e(ll)
		if "`skip'" == "noskip" {
				/* constant only model */

			TwoStep "`dep'" "" "`nc'" "`trtdep'" "" "`trtnc'"   /*
				*/ "`wgt'" "`wtype'" "`wtexp'" `touse'      /*
				*/ "`first'" "`hazard'"

			tempname b0 athrho lnsigma 
			mat `b0' = e(b)
			mat `b0'= `b0'[1, 1..colsof(`b0')-1]
			scalar `athrho' = max(min(e(rho), .85), -.85)
			scalar `athrho' = 1/2*ln((1+`athrho')/(1-`athrho'))
			scalar `lnsigma' = ln(e(sigma))
			mat `b0' = `b0', `athrho', `lnsigma'
			local continu "continue"
			est clear
			`showlog' di
			`showlog' di in green "Fitting constant-only model:"
			Mmethod "`dep'" "" "`nc'" "`trtdep'" "" /*
				*/ "`trtnc'" "`wgt'" "`wtype'" "`wtexp'"    /*
				*/ `touse' "`b0'"  "`robust'"   /*	 
				*/ "`clusopt'" "`log'" "" "`stdopts'"      
			`showlog' di
			`showlog' di in green "Fitting full model:"
		}

		Mmethod "`dep'" "`ind'" "`nc'" "`trtdep'" "`trtind'" /*
			*/ "`trtnc'" "`wgt'" "`wtype'" "`wtexp'"    /*
			*/ `touse' "`b1'"  "`robust'"   /*		 
			*/ "`clusopt'" "`log'" "`continu'" "`stdopts'"      

				/* comparison test */

		if "`robust'" != "" {
			est local chi2_ct "Wald"
			qui test [athrho]_b[_cons] = 0
			est scalar chi2_c = r(chi2)
		}
		else {
			est local chi2_ct "LR"
			est scalar chi2_c = 2*(e(ll) - `llreg'- `llprob' )
		}
		est scalar p_c = chiprob(1, e(chi2_c))

		Reparm
		est local title "Treatment effects model -- MLE"
		if "`hazard'" != "" { est local hazard "`hazard'" }
		Display, level(`level')
	}
	else {
		if "`robust'`cluster'`wgt'" != ""{
			di in red "robust, cluster and weights not allowed" /*
				*/ " with the twostep option"
			est clear 
			exit 198
		}

		est local title "Treatment effects model -- two-step estimates"
		Display2, level(`level')
		est local depvar `dep' `trtdep'
		est local method "twostep" 
		est local ll_p
	}
	est local cmd "treatreg"	
	est local predict "treatr_p"
end

program define TwoStep, eclass
	args 		dep	/*
		*/	ind	/*
		*/	nc	/*
		*/	trtdep	/*
		*/	trtind	/*
		*/	trtnc	/*
		*/	wgt	/*
		*/	wtype	/*
		*/	wtexp	/*
		*/	touse	/*
		*/	first   /*
		*/	hazard
	if "`first'" != "" {local first noi }

	tempname bprb vprb bregx vreg breg llprob
	tempvar xbprb lambda tao

				/* Step1 -- probit */
	qui{
		`first' probit `trtdep' `trtind' `wgt' if `touse',  /*
			*/ nolog `trtnc'
		mat `vprb' = get(VCE)
		mat `bprb' = get(_b)
		scalar `llprob' = e(ll)
		predict double `xbprb', xb
		gen double `lambda' = normd(`xbprb')/normprob(`xbprb') /*
			*/ if `trtdep' == 1 & `touse'
		replace  `lambda' = -normd(`xbprb')/(1-normprob(`xbprb')) /*
			*/ if `trtdep' == 0 & `touse'
		if "`hazard'" != "" { gen double `hazard' = `lambda' }
		gen double `tao' = `lambda'*(`lambda' + `xbprb')
	}
				/* Step 2 -- regress */
	
	qui{	
		if "`nc'" =="" {
			tempvar one
			gen byte `one' = 1
		}
		reg `dep' `ind' `trtdep' `one' `lambda' `wgt' if `touse', noc 
		mat `breg' = get(_b) 
		mat `vreg' = get(VCE)
		local blambda _coef[`lambda']
		sum `tao' `wgt' if `touse' 
		local mtao r(mean)
		local sigma =sqrt(e(rss)/e(N) + `mtao'*`blambda'^2)
		local rho = `blambda'/`sigma'
		if `sigma' == 0 { local sigma = .001}
		if `rho' == . { local rho = 0 }
		if abs(`rho') > 1 {
			local rho = `rho'/abs(`rho') 
		}
		local rho2= `rho'*`rho'
	}

				/* Cov adjustment */

	tempname  XsI Vbs F Q 
	if e(rmse) != 0 {
		matrix `XsI' = `vreg' / (e(rmse)^2)
	}
	else {
		di in red "insufficient information"
		exit 2000
	}

	local fulnam `ind' 
	local fulnam "`fulnam' `trtdep'"
	if "`nc'" == ""{
		tempvar one
		gen byte `one' = 1
		local fulnam "`fulnam' _cons"
	}
	local fulnam "`fulnam' lambda"
	local kreg : word count `fulnam'
	local kreg1 = `kreg' + 1
	qui mat accum `F' = `ind' `trtdep' `one' `lambda' `trtind' /*
		*/ [iw = `tao'], `trtnc', if `touse'
	mat `F' = `F'[1.. `kreg', `kreg1' ...]
	mat rowname `F' = `fulnam'
	mat `Q' = `rho2'*`F'*`vprb'*`F''

 	tempvar rho2tao 
	qui gen double `rho2tao' = 1 - `rho2'*`tao' if `touse' 
	qui mat accum `Vbs' = `ind' `trtdep' `one' `lambda' [iw=`rho2tao'], /*
		*/ nocons, if `touse'
	mat `Vbs' = `sigma'^2 * `XsI' *(`Vbs' + `Q')*`XsI' 

				/* Build the eqn names */
	local eqnam 
	tokenize `fulnam'
	local i 1
	while `i' < `kreg' {
		local eqnam `eqnam' `dep':``i''
		local i= `i' + 1
	}

	local trtname `trtdep'
	local kprb : word count `trtind'
	local i 1
	tokenize `trtind'
	while `i' <= `kprb' {
		local eqnam `eqnam' `trtname':``i''
		local i = `i' + 1
	}
	if "`trtnc'" =="" {
		local kprb = `kprb' + 1
		local eqnam `eqnam' `trtname':_cons
	}
	local eqnam `eqnam' hazard:lambda
	 
	tempname bfull Vfull zeros Vmills CovMill zeroPrb b0 bmills
	local kreg0 = `kreg' -1
	mat `b0' = `breg'[1, 1..`kreg0']
	mat `bmills' = `breg'[1, `kreg']
	mat `bfull' = `b0', `bprb', `bmills'
	mat `Vmills' = `Vbs'[`kreg', `kreg']
	mat `CovMill' = `Vbs'[1..`kreg0', `kreg']
	mat `Vbs' = `Vbs'[1..`kreg0', 1..`kreg0']
	mat `zeros' = J(`kreg0', `kprb', 0)
	mat `zeroPrb' = J(`kprb', 1, 0)
	mat `Vbs' = /*
		*/ ( `Vbs'	, `zeros'	, `CovMill' \	/*
		*/   `zeros''	, `vprb'	, `zeroPrb'  \  /*
		*/   `CovMill'', `zeroPrb''	, `Vmills' )
	mat colnames `bfull' = `eqnam'
	mat rownames `Vbs' = `eqnam'
	mat colnames `Vbs' = `eqnam'

	qui count if `touse'
	est post `bfull' `Vbs', obs(`r(N)')
	est scalar N = `r(N)'
	capture qui test `ind' `trtdep', min
	if _rc == 0 {
		est scalar chi2 = `r(chi2)'
		est scalar df_m = `r(df)'
		est scalar p    = `r(p)'
	}
	est local chi2type "Wald"
	if "`hazard'" !="" { est local hazard `hazard'}
	est scalar rho= `rho'
	est scalar sigma = `sigma'
	est scalar lambda = [hazard]_b[lambda]
	est scalar selambda = [hazard]_se[lambda]
	est scalar ll_p = `llprob'
end


program define Mmethod, eclass
	args 		dep	/*
		*/	ind	/*
		*/	nc	/*
		*/	trtdep	/*
		*/	trtind	/*
		*/	trtnc	/*
		*/	wgt	/*
		*/	wtype	/*
		*/	wtexp	/*
		*/	touse	/*
		*/	b0 	/*
		*/	robust	/*
		*/	clusopt	/*
		*/      log     /*
		*/ 	continu /*
		*/	stdopts	
	ml model lf treat_ll					/* 
		*/ (`dep': `dep' = `ind' `trtdep', `nc')	/*
		*/ (`trtdep': `trtdep' = `trtind', `trtnc')	/*
		*/ /athrho /lnsigma 				/*
		*/ if `touse' `wgt',				/* 
		*/ `robust' `clusopt'				/*
		*/ collinear missing max nooutput nopreserve	/*
		*/ init(`b0', copy) search(off) `log' `continu' /* 
		*/ `stdopts' `mlopts' 
	est local method "ml"
end	

program define Display2
	syntax [, level(integer $S_level) ]
	#delimit ;
	di in gr _n "`e(title)'" _col(49) "Number of obs"
		in gr _col(68) "="
		in ye _col(70) %9.0g e(N) ;

	local chifmt = cond(e(chi2)<1e+6,"%9.2f","%9.2e") ;
	di in gr _n _col(49) "`e(chi2type)' chi2(" in ye e(df_m) in gr ")"
		in gr _col(68) "="
		in ye _col(70) `chifmt' e(chi2) ;
	di in gr _col(49) "Prob > chi2"
		in gr _col(68) "="
		in ye _col(73) %6.4f chiprob(e(df_m),e(chi2)) _n ;
	#delimit cr
	est display, level(`level')
	di in gr "     rho | " in ye %10.5f e(rho)
	di in gr "   sigma | " in ye %10.0g e(sigma)
	di in gr "  lambda | " in ye %10.0g e(lambda) "  " %9.0g e(selambda)
	di in gr _dup(79) "-"

end

program define Reparm, eclass

			/* somewhat superceded by _diparm, but kept for
			 * lambda and so sigma and rho can be saved in e() */

	tempname b v d tau lns rho sig tmp lambda lamse
	mat `b' = get(_b)
	mat `v' = get(VCE)

	mat `tmp' = `b'[1,"athrho:_cons"]
	scalar `tau' = `tmp'[1,1]
	mat `tmp' = `b'[1,"lnsigma:_cons"]
	scalar `lns' = `tmp'[1,1]

	scalar `rho' = (exp(2*`tau')-1) / (exp(2*`tau')+1)
	scalar `sig' = exp(`lns')
	scalar `lambda' = `rho'*`sig'

	mat `d' =  ( `sig'*4*exp(2*`tau')/(exp(2*`tau')+1)^2 , `lambda' )
	mat `v' =  (`v'["athrho:_cons".."lnsigma:_cons",      /*
		*/   "athrho:_cons".."lnsigma:_cons"]  )
	mat `v' = `d'*`v'*`d''
	scalar `lamse' = sqrt(`v'[1,1])

	est scalar rho    = `rho'
	est scalar sigma  = `sig'
	est scalar lambda = `lambda'
	est scalar selambda  = `lamse'
	/*  Double saves for backward compatibility */
	global S_1 = `rho'
	global S_2 = `sig'
	global S_3 = `lambda'
	global S_4 = `lamse'
end

program define Display
	syntax [, Level(integer $S_level)]

	#delimit ;
	di in gr _n "`e(title)'" _col(49) "Number of obs"
		in gr _col(68) "="
		in ye _col(70) %9.0g e(N) ;

	local chifmt = cond(e(chi2)<1e+6,"%9.2f","%9.2e") ;
	di in gr _n _col(49) "`e(chi2type)' chi2(" in ye e(df_m) in gr ")"
		in gr _col(68) "="
		in ye _col(70) `chifmt' e(chi2) ;
	di in gr "Log likelihood = "e(ll) _col(49) "Prob > chi2"
		in gr _col(68) "="
		in ye _col(73) %6.4f chiprob(e(df_m),e(chi2)) _n ;
	#delimit cr

	est display, level(`level') neq(2) plus
	_diparm athrho, level(`level')
	_diparm lnsigma, level(`level')
	di in gr _dup(9) "-" in gr "+" _dup(69) "-"
	_diparm athrho, level(`level') tanh label("rho")
	_diparm lnsigma, level(`level') exp label("sigma")
	DiLambda `level'
	di in gr _dup(79) "-"
	if "`e(vcetype)'" != "Robust" {
		local testtyp LR
	}
	else local testyp Wald
	di in gr  "`testtyp' test of indep. eqns. (rho = 0):" /*
		*/ _col(38) "chi2(" in ye "1" in gr ") = "   /*
		 */ in ye %8.2f e(chi2_c)                     /*
		 */ _col(59) in gr "Prob > chi2 = " in ye %6.4f e(p_c)
	di in gr _dup(79) "-"
end 

program define DiLambda
	args level

	tempname z lb ub

	scalar `z' = invnorm((100 + `level') / 200)
	scalar `lb' = e(lambda) - `z'*e(selambda)
	scalar `ub' = e(lambda) + `z'*e(selambda)

	local llb : di %9.0g `lb'
	if length("`llb'") > 9 {
		local lbfmt "%8.0g"
	}
	else	local lbfmt "%9.0g"

	local uub : di %9.0g `ub'
	if length("`uub'") > 9 {
		local ubfmt "%8.0g"
	}
	else	local ubfmt "%9.0g"

	di in gr _col(3) "lambda |  " /*
		*/  in ye %9.0g e(lambda) "  " %9.0g e(selambda)    /*
		*/  _col(58) `lbfmt' `lb' "   "  `ubfmt' `ub'

end