*! Version 1.0.0 (STB-58: sxd2)
program define optbud, eclass
	version 6.0
	local varlist "required existing"     
	local if "optional"
	local in "optional"
	local options "FIRST(string) PREV(string) VAR(integer 0) B(real 100) C1(real 1) C2(real 10) CODING(integer 0)" 
	parse "`*'"

   	if "$S_E_cmd"=="optbud" {
     		error 301
   }

if ("`first'"=="") {
	di in red "WARNING: Please support the first stage covariates"
	exit
	}

if (`var'==0 | "`prev'"=="" | "`b'"=="" | "`c1'"=="" | "`c2'"=="" ) {
	di in red "WARNING: Some of the following items are missing:"
	di in red "         var (position of variable to be optimised) or"
	di in red "	    prev (prevalences) or"
	di in red "	    b (available budget) or"	
	di in red "	    c1 (first stage cost) or"	
	di in red "	    c2 (second stage cost) or"	

	exit
	}

preserve
local c11=`c1'
local c21=`c2'
msopt `varlist',first(`first') sample(`prev') coding(`coding')

tempname wzykk
matrix `wzykk'=J(e(lev),1,0)
local i=1
while `i'<=e(lev) {
     matrix `wzykk'[`i',1]=el(e(wzy`i'),`var',`var')
     local i=`i'+1}

*******create sqrt(wzykk) matrix***************
tempname sqwzy wzyprev
matrix `sqwzy' =J(e(lev),1,0)
matrix `wzyprev' =J(e(lev),1,0)
local i=1
tempvar sqr mult
while `i'<=e(lev) {
     scalar `sqr'=sqrt(`wzykk'[`i',1])
     matrix `sqwzy'[`i',1]=`sqr'
     scalar `mult'=`wzykk'[`i',1]*`prev'[`i',1]
     matrix `wzyprev'[`i',1]=`mult'
     local i=`i'+1}

*****computation start here********************
tempname prev1
matrix `prev1'=`prev'

*****check if the denominator <= 0 ************
tempname denom invI
tempvar denom1
matrix `denom'=`wzykk''*`prev'
matrix `invI'=syminv(e(info))
scalar `denom1'=`invI'[`var',`var']-`denom'[1,1]

local c=1
while `denom1'<=0{
	tempvar maxi propmax
	matrix wzyprev=`wzyprev'
	svmat wzyprev
	egen `propmax'=max(wzyprev1)
	scalar `maxi'=`propmax'[1]
	drop `propmax' wzyprev1 

	tempname test
	matrix `test'=J(e(lev),1,0)
	local i=1
	while `i'<=e(lev) {
	     local check1=sign(`wzyprev'[`i',1]-`maxi'+0.0001)
	     matrix `test'[`i',1]=`check1'
	     local i=`i'+1}

	local j=1
	while `j'<=e(lev){
		if `test'[`j',1]==1{
		tempvar post`c'
		scalar `post`c''=`j'
		local c1=`c1'+`prev1'[`j',1]*`c2'
		matrix `prev1'[`j',1]=0
		local j=`j'+e(lev)}

	local j=`j'+1
	}

	matrix `denom'=`wzykk''*`prev1'

	scalar `denom1'=`invI'[`var',`var']-`denom'[1,1]

	local i=1
	while `i'<=e(lev) {
	matrix `wzyprev'[`i',1]=`wzykk'[`i',1]*`prev1'[`i',1]
	local i=`i'+1}

local c=`c'+1
}

********update the first and second stage cost***************
tempvar divisor fstc scdc
tempname numer
matrix `numer' = `sqwzy''*`prev1'
scalar `divisor' = `numer'[1,1]
scalar `fstc'  =`c1'
scalar `scdc'  =`c2'
matrix `numer' = `numer'*sqrt(`fstc'*`scdc')

matrix `numer'=`numer'/sqrt(`denom1')

tempname fstcost
matrix `fstcost'=J(1,1,`c1')
***************************************************************

******get the sampling fraction*******************************
matrix `numer' = `numer'+`fstcost'

matrix `numer'[1,1]=1/`numer'[1,1]

tempname nobs
matrix `nobs' = `b'*`numer'

tempvar front
scalar `front' = (`b'-`nobs'[1,1]*`c1')/(`nobs'[1,1]*`c2')

tempname prop
matrix `prop' = (`front'/`divisor')*`sqwzy'


****check if any prop>1***************************************
local check=0
local i=1
while `i'<=e(lev) {
      local check=`check'+(`prop'[`i',1]>1)
      local i=`i'+1
	}
**************************************************************


****correcting for prop >1 *****************
while `check'>0 {
	tempvar maxi propmax
	matrix prp_BC=`prop'
	svmat prp_BC
	egen `propmax'=max(prp_BC1)
	scalar `maxi'=`propmax'[1]
	drop `propmax' prp_BC1

	tempname test
	matrix `test'=J(e(lev),1,0)
	local i=1
	while `i'<=e(lev) {
	     	local check1=sign(`prop'[`i',1]-`maxi'+0.0001)
	     	matrix `test'[`i',1]=`check1'
	     	local i=`i'+1
		}
	
	local j=1
	while `j'<=e(lev){
		if `test'[`j',1]==1{
		tempvar post`c'
		scalar `post`c''=`j'
		local c1=`c1'+`prev1'[`j',1]*`c2'
		matrix `prev1'[`j',1]=0
		local j=`j'+e(lev)
		}

local j=`j'+1
}

*****computation to update the proportion and nobs start here********************
matrix `denom'=`wzykk''*`prev1'
scalar `denom1'=`invI'[`var',`var']-`denom'[1,1]
matrix `numer' = `sqwzy''*`prev1'
scalar `divisor' = `numer'[1,1]
scalar `fstc'  =`c1'
scalar `scdc'  =`c2'
matrix `numer' = `numer'*sqrt(`fstc'*`scdc')

matrix `numer'=`numer'/sqrt(`denom1')

matrix `fstcost'=J(1,1,`c1')

matrix `numer' = `numer'+`fstcost'

matrix `numer'[1,1]=1/`numer'[1,1]

matrix `nobs' = `b'*`numer'

scalar `front' = (`b'-`nobs'[1,1]*`c1')/(`nobs'[1,1]*`c2')

matrix `prop' = (`front'/`divisor')*`sqwzy'

***********************************************

******** all the prop>1 are given sampling fraction = 1************
local i=1
while `i'<=`c' {
      matrix `prop'[`post`i'',1]=1
      local i=`i'+1
	}


****re-check if any prop>1****************
local check=0
local i=1
while `i'<=e(lev) {
      local check=`check'+(`prop'[`i',1]>1)
      local i=`i'+1
	}
***************************************

local c=`c'+1
}

local i=1
tempname scsamp
local scsize=0
matrix `scsamp'=J(e(lev),1,0)
while `i'<=e(lev) {
    	tempvar prop`i'
    	scalar `prop`i'' = round(`prop'[`i',1],.001)
	matrix `scsamp'[`i',1]=round(`prev'[`i',1]*`prop'[`i',1]*`nobs'[1,1],1)
	local scsize=`scsize'+`scsamp'[`i',1]
    	di in blue "the optimal sampling fraction (sample size) for grp_yz `i' = " `prop`i'' " ("`scsamp'[`i',1] ")"
    	local i=`i'+1
	}

tempvar num
scalar `num'=`nobs'[1,1]
di in white "the optimal number of obs = " int(`num')


***computing the minimum variance obtained by the optimal design*******
tempvar weight 
tempname covar varsi result minvar
local c: word count `varlist'
matrix `result'=J(`c',`c',0)
local i=1
while `i'<=e(lev) {
	scalar `weight'=`prev'[`i',1]/`prop'[`i',1]*(1-`prop'[`i',1])
	matrix `varsi'=`weight'*e(wzy`i')
	matrix `result'=`result'+`varsi'
	local i=`i'+1}

matrix `covar'=(`invI'+`result')*1/int(`num')
matrix `minvar'=vecdiag(`covar')

***display the minimum variance attained
local varpost: word `var' of `varlist' 
di in yellow "the minimum variance for `varpost' : " `covar'[`var',`var']

*********** check the total budget spent **********************
local budspnt=int(`num')*`c11'+`scsize'*`c21'
di in yellow "total budget spent:  " `budspnt'

est matrix nobs `nobs'
est matrix frac `prop'
est local cmd "optbud"
restore, preserve
end