*! Version 1.0.0  (STB-58: sxd2)
program define optfixn, eclass
	version 6.0
	local varlist "required existing"     
	local if "optional"
	local in "optional"
	local options "FIRST(string) N1(string) N2(integer 1) VAR(integer 0) coding(integer 0)" 
	parse "`*'"

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


preserve

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

if (`var'==0 | "`n1'"=="" | "`n2'"=="") {
	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 "	    n1 (number of obs) or"
	di in red "	    n2 (the second stage sample sizes)"	
	exit
	}

msopt `varlist',first(`first') sample(`n1') 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}

tempvar length
*********get the nobs ******
local i=1
scalar `length'=0
while `i'<=e(lev) {
	scalar `length'=`length'+`n1'[`i',1]
	local i=`i'+1
	}
	
tempname prev
matrix `prev'=`n1'/`length'

*******create sqrt(wzykk) matrix***************

tempvar sqr
tempname sqwzy
matrix `sqwzy' =J(e(lev),1,0)

local i=1
while `i'<=e(lev) {
	scalar `sqr'=sqrt(`wzykk'[`i',1])
	matrix `sqwzy'[`i',1]=`sqr'
	local i=`i'+1
	}

*****computation start here********************
tempvar denom1
tempname numer denom
matrix `numer'=(`n2'/`length')*`sqwzy'

matrix `denom'=`sqwzy''*`prev'

scalar `denom1'=`denom'[1,1]

tempname prop
matrix `prop'=`numer'/`denom1'
***********************************************

****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 *****************
tempname sample1
matrix `sample1'=`n1'
tempname prev1
matrix `prev1'=`prev'
local c=1
while `check'>0 {
	tempvar propmax maxi
	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'
		matrix `sample1'[`j',1]=0
		matrix `prev1'[`j',1]=0
		local j=`j'+e(lev)
		}

	local j=`j'+1
	}
	local i=1
	local n11=0
	while `i'<=e(lev) {
		local n11=`n11'+`sample1'[`i',1]
		local i=`i'+1
		}
	local n21=`n2'-(`length'-`n11')

*****update the estimates********************
	matrix `numer'=(`n21'/`length')*`sqwzy'

	matrix `denom'=`sqwzy''*`prev1'

	scalar `denom1'=`denom'[1,1]

	matrix `prop'=`numer'/`denom1'

*******all 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
}

****display sampling fraction for each (y,z) stratum
local i=1
tempname scsamp
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(`n1'[`i',1]*`prop'[`i',1],1)
	di in blue "the optimal sampling fraction(sample size) for grp_yz `i' = " `prop`i'' " ("`scsamp'[`i',1] ")"
	local i=`i'+1
}


***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}

tempname one nobs
matrix `one'=J(e(lev),1,1)
matrix `nobs'=`n1''*`one'
matrix `covar'=(syminv(e(info))+`result')*1/int(`nobs'[1,1])
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 second stage sample size ********
tempname sum
matrix `sum'=`prop''*`n1'

di in yellow "Total second stage sample size =" `sum'[1,1]



est matrix prop `prop'
est local cmd "optfixn"
restore, preserve
end