*! version 1.0.0  12apr2004 M. Lokshin, Z. Sajaia              (SJ4-3: st0000)
#delim ;

cap program drop movestay;
program define movestay, sortpreserve;
	version 8;
	if replay() {;
		if "`e(cmd)'" ~= "movestay" error 301;
		movestay_replay `0';
		}; // end if
	else movestay_mx `0';
end; // end program movestay

cap program drop movestay_replay;
program define movestay_replay;
	syntax, [Level(int $S_level)];

	local sigma1  diparm(lns1, exp  label("sigma_1"));
	local rho1    diparm(r1  , tanh label("rho_1"));
	local sigma2  diparm(lns2, exp  label("sigma_2"));
	local rho2    diparm(r2  , tanh label("rho_2"));

	ml display, level(`level') `sigma1' `sigma2' `rho1' `rho2';

	if ("`e(vcetype)'" != "Robust") local testtyp LR; else local testtyp Wald;
	display	in green "`testtyp' test of indep. eqns. :"  _col(38) "chi2(" in yellow "1" in green ") = "
		in yellow %8.2f e(chi2_c) _col(59)
		in green  "Prob > chi2 = "
		in yellow %6.4f e(p_c);
	display	in smcl in green "{hline 78}";
end; // end program movestay_reply

cap program drop movestay_mx;
program define movestay_mx, eclass;

	syntax 	anything(id="equations id" equalok) [pweight iweight fweight] [if] [in] ,
	       	SELect(string) [CLuster(varname) Robust noLOg] [*];

	mlopts mlopts, `options';

	marksample touse;
	markout `touse' `cluster', strok;

	quietly {;

	if ("`weight'"!="") {;	tempvar wvar; 	generate double `wvar' `exp' if `touse'; local weight "[`weight'=`wvar']";};
	if (("`robust'"!="") | ("`weight'"=="pweight") | ("`cluster'"~="")) local lrtest nolrtest;
	if ("`cluster'"!="") local clopt cluster(`cluster');

	// read vector of parameters

	 gettoken reg1 reg2_ :  anything,  match(parns) bind;
           if ("`parns'"!="(") {; local reg1 "`anything'"; local reg2_ ""; };
	    tokenize "`reg1'", parse("()");
	      if ("`1'"!="`reg1'") {; display as error "If more than one, equations must be enclosed in brackets";error;};
            tokenize "`reg1'", parse("=");
                if ("`2'"!="=") {;
                    tokenize "`reg1'";
                    local y_reg1 `1';
                    macro shift;
                    local x_reg1 `*';
	        };
                else  {;
		        if ("`4'"=="=") {; display as error "If more than one, equations must be enclosed in brackets";error;};
                    local y_reg1 `1';
                    local x_reg1 `3';
                };
           if ("`reg2_'"!="") {;
                gettoken reg2  : reg2_ , match(parns) bind;
                if ("`parns'"!="(") local reg2 "`reg2_'";
                tokenize "`reg2'", parse("=");
	 	   if ("`2'"!="=") {;
	               tokenize "`reg2'";
                       local y_reg2 `1';
	               macro shift;
	               local x_reg2 `*';
		   };
                   else {;
                          if ("`4'"=="=") {; display as error "If more than one, equations must be enclosed in brackets";error;};
                     local y_reg2 `1';
                     local x_reg2 `3';
                   };
       	   };
	   else {; local y_reg2 `y_reg1'; local x_reg2 `x_reg1';};

        tokenize "`select'", parse("=");         // define vars for regression equation
	   if ("`2'"!="=") {;
	         tokenize "`select'";
                 local y_prob `1';
	         macro shift;
	         local x_prob `*';
	   };
           else {;
	         local y_prob `1';
                 local x_prob `3';
                   };

	local x_reg1: list uniq x_reg1;
	local x_reg2: list uniq x_reg2;
	local x_prob  `x_reg1' `x_reg2' `x_prob';
	local x_prob: list uniq x_prob;
	local px1: word 1 of `x_prob';

        if ("`y_reg1'"=="`y_reg2'") {; local eq1 = "`y_reg1'_1"; local eq2 = "`y_reg2'_0";};
        else 	      		    {; local eq1 = "`y_reg1'";   local eq2 = "`y_reg2'";  };

	// define initial values
	tempname I_reg1 I_reg2 I_prob TMP ll_0;
	tempvar	y_reg1_t y_reg2_t;

	generate double `y_reg1_t'=`y_reg1' if `y_prob'==1;	      // define subsamples for initial value regressions
	generate double `y_reg2_t'=`y_reg2' if `y_prob'==0;

	scalar `ll_0' = 0;

	noisily display _newline "Fitting initial values " _continue;

	probit  `y_prob'   `x_prob' `weight' if `touse', nocoef; sca `ll_0'=`ll_0'+e(ll); noisily display "." _continue;
 	regress `y_reg1_t' `x_reg1' `weight' if `touse', nohead; sca `ll_0'=`ll_0'+e(ll); noisily display "." _continue;
 	regress `y_reg2_t' `x_reg2' `weight' if `touse', nohead; sca `ll_0'=`ll_0'+e(ll); noisily display "." _continue;

	heckman `y_reg1_t' `x_reg1' if `touse', select(`x_prob') twostep; noisily display "." _continue;
	       	matrix `TMP'   = e(b);
	       	matrix `I_reg1'= `TMP'[1,1.."`y_reg1_t':_cons"];                    // initial values for regression 1
	       	matrix `I_prob'= `TMP'[1,"select:`px1'".."select:_cons"];           // initial values for probit
	       	local I_rho1  = atanh(e(rho));                       	            // initial values for rho1
	       	local I_sig1  = ln(e(sigma));                                       // initial values for sigma 1
	heckman `y_reg2_t' `x_reg2' if `touse', select(`x_prob') twostep; noisily display "." _continue;
	       	matrix `TMP'   = e(b);
	       	matrix `I_reg2'= `TMP'[1,1.."`y_reg2_t':_cons"];                    // initial values for regression 2
	       	local I_rho2  = - atanh(e(rho)); 			            // initial values for rho2
	       	local I_sig2  = ln(e(sigma));                                       // initial values for sigma2

	}; // end quietly

	// check if any variables were dropped by heckman in case of multicolinearity
	// do not use initial values if some variables were dropped
	local vars_x1 : word count(`x_reg1');
	local vars_x2 : word count(`x_reg2');
	local vars_pb : word count(`x_prob');

	if ((colsof(`I_reg1')~=`vars_x1'+1) |
	    (colsof(`I_reg2')~=`vars_x2'+1) |
	    (colsof(`I_prob')~=`vars_pb'+1))  local init_cond = "";
	else local init_cond init(`I_reg1' `I_reg2' `I_prob' `I_sig1' `I_sig2' `I_rho1' `I_rho2',copy);

	// maximization
	ml model d2 movestay_d2
	      	(`eq1'	   :`y_reg1' = `x_reg1')
	      	(`eq2'	   :`y_reg2' = `x_reg2')
	      	(`y_prob'  :`y_prob' = `x_prob')
	      	/lns1
	      	/lns2
	      	/r1
	      	/r2
	       `weight'
	      	if `touse'
	      	,
	      	title("Endogenous switching regression model")
	        search(quietly)
	      	`init_cond'
	      	maximize
	      	missing
	      	nopreserve
		`clopt'
		`robust'
		`log'
 	      	`mlopts'
	      	;

	ereturn scalar ll_0 = `ll_0';     	    	                              	// loglikelihood for non-correlated case
	ereturn scalar k_aux= 4;							// identify ancillary parameters
        ereturn local cmd    "movestay";    	                       	                // fill in e(cmd) ;
	ereturn local title2 "(endogenous switching regression model)" ;
	ereturn local ivar_reg1 `x_reg1';
	ereturn local ivar_reg2 `x_reg2';
	ereturn local ivar_prob `x_prob';
	ereturn local depvar "`eq1' `eq2' `y_prob'";

	if "`robust'"=="" {; eret local chi2_ct "LR";
			     eret scalar chi2_c = 2 * (e(ll) - `ll_0')  ;}; // end if
	else 		  {; eret local chi2_ct "Wald"; quietly test [r1]_b[_cons] [r2]_b[_cons];
			     eret scalar chi2_c = r(chi2);}; // end else
	eret scalar p_c = chiprob(1, e(chi2_c));

	movestay_replay;
end; // end program movestay_mx