program define ivprob, eclass
version 6.0
syntax varlist  [if] [in], ENDOG(varlist) IV(varlist) EXOG(varlist)
* parse "`*'"
local depvar "`varlist'"
tempvar tmpy touse
marksample touse

local niv : word count `iv'
local nendog : word count `endog'

/* need to eliminate variables that predict success/failure
perfectly and restrict sample accordingly */
qui probit `depvar' `endog' `exog' if `touse'
replace `touse'=e(sample)
matrix tmp = get(_b)
local tmp : colnames tmp
tokenize "`tmp'", parse(" ")
local i 1
while `i'<=`nendog' {
	mac shift
	local i=`i'+1
}
local exog "`1'"
mac shift
while "`1'" != "_cons" {
	local exog "`exog' `1'"
	mac shift
}
local nexog : word count `exog'


local K=`nexog'+`niv'		            /* K=total number of exogenous variables */
matrix Dhat=J(`K'+1,`nendog'+`nexog'+1,0)  /* Inititialize Dhat matrix */

/* fill in remaining columns of Dhat= I_1: lower right identity matrix */
matrix Dhat[`niv'+1,`nendog'+1]=I(`nexog'+1)

parse "`endog'", parse(" ")
local e = 1			/* do 1st stage regression for each endogenous variable */
while `e'<=`nendog' {
	qui regress `1' `iv' `exog' if `touse'
	matrix tmp=get(_b)
	matrix Dhat[1,`e']=tmp'
	tempvar fitted`e' resid`e'
	qui predict `fitted`e'' `if'
	qui replace  `fitted`e'' =. if `1' ==.
	qui predict `resid`e'' `if', residuals
	qui replace `resid`e'' =. if `1' ==.
	local fitted "`fitted' `fitted`e''"
	local resid "`resid' `resid`e''"
	mac shift
  	local e = 1 + `e'
}

/* get lee model estimates */
cap estimates drop lastpbt
qui probit `depvar' `fitted' `exog' if `touse'
matrix tmp = get(_b)
matrix beta=tmp[1,1..`nendog']
matrix means=get(mns)
matrix colnames means = `endog' `exog' _cons
estimates hold lastpbt


* Run probit of depvar on resid, and exog variables
qui probit `depvar' `iv' `exog' `resid' if `touse'
local obs = e(N)
local df = e(df_m)
matrix tmp = get(_b)
matrix alpha=J(1,`K'+1,0)  /* initialize alpha */

matrix alpha[1,1]=tmp[1,1..`K']  /* take coefficients on all exog variables */
matrix alpha[1,`K'+1]=tmp[1,`K'+`nendog'+1]  /*get coefficient on constant (skipping coefs on residuals */
matrix lamda=tmp[1,`K'+1..`K'+`nendog'] 	/* coefficients on residuals go into lamda */
matrix tmp = get(VCE)			/* now get VCE matrix, skipping residuals */
matrix Jinv=J(`K'+1,`K'+1,0)		/* initialize Jinv matrix and move VCE matrix into it */
matrix Jinv[1,1]=tmp[1..`K',1..`K']
local i 1
while `i'<=`K'  {
	matrix Jinv[`K'+1,`i']=tmp[`K'+`nendog'+1,`i']
	matrix Jinv[`i',`K'+1]=Jinv[`K'+1,`i']
	local i=`i'+1
}
matrix Jinv[`K'+1,`K'+1]=tmp[`K'+`nendog'+1,`K'+`nendog'+1]

local e 1
gen `tmpy'=0
parse "`endog'", parse(" ")
while `e'<=`nendog' {
	local tmp=lamda[1,`e']-beta[1,`e']
	qui replace `tmpy'=`tmpy'+`1'*`tmp'
	mac shift
	local e=`e'+1
}
qui regress `tmpy' `iv' `exog' if `touse'
matrix V=get(VCE)

matrix Omega=V+Jinv
matrix OmInv=inv(Omega)
matrix delta1=Dhat' * OmInv  * Dhat
matrix rownames delta1= `endog' `exog' _cons
matrix colnames delta1= `endog' `exog' _cons
matrix delta1=inv(delta1)

matrix delta2=Dhat' * OmInv * alpha'
matrix delta=delta1 * delta2
matrix rownames delta = `endog' `exog' _cons
matrix delta=delta'
matrix bnew=delta
matrix Vnew=delta1
estimates unhold lastpbt

tempname b V
matrix `b'=e(b)
matrix `V'=e(V)
local nobs=e(N)
local tr2_p=e(r2_p)
local tchi2= e(chi2)
local tdf_m= e(df_m)
local tll=   e(ll)
local tll_0= e(ll_0)

est post `b' `V', obs(`nobs') depname(`depvar') esample(`touse')
est local depvar "`depvar'"
est scalar r2_p=`tr2_p'
est scalar chi2= `tchi2'
est scalar df_m= `tdf_m'
est scalar ll=   `tll'
est scalar ll_0= `tll_0'
est local cmd="probit"

est matrix mns means
estimates repost b=bnew V=Vnew , rename
estimates disp

end