// two-stage estimation procedure from Byzalov and Basu, "Modeling the Determinants of Meet-or-Just-Beat Behavior in Distribution Discontinuity Tests",
// https://ssrn.com/abstract=3117088

program define kinkyX, eclass byable(recall) sortpreserve
	version 15
	syntax varlist [if], binwidth(real) est_bins(int) em_bins(int) em_type(string) degree(int) cluster(varname) [em_binsMinus(int 0) smooth(name) kinky(name) sloppy predict(name) graph(name) graphtitle(string)]
		// binwidth(real) = bin width for the discretization of earnings
		// est_bins(int) = number of estimation bins on each side of zero
		// em_bins(int) = number of bins affected by earnings management
		// em_binsMinus(int) = (optional) a different number of bins affected by earnings management below zero (if not specified, em_bins is used both below and above zero)
		// em_type(string) = use em_type(i) for Model I (flat increments), em_type(ii) for Model II (triangular increments). For backward compatibility, you can also use em_type(0) for Model I and em_type(1) for Model II
		// degree(int) = degree of the polynomial approximation for the pre-managed distribution (sould be less than 2*(est_bins-em_bins) to avoid perfect overfitting)
		// cluster(varname) = the cluster variable (required -- why would you ever want to use unclustered SE?); for now, only one-way clustering is supported
		// smooth(name) = stem for the new variable names in the estimates of the smooth pre-managed distribution
		// kinky(name) = stem for the new variable names in the estimates of the distribution discontinuity
		//   CAREFUL -- the code will overwrite any existing `smooth'* and `kinky'* variables in the dataset
		// sloppy = use sloppy single-stage SE instead of the properly adjusted (and much more time-consuming) two-stage SE -- use this option only to save time in quick-and-dirty exploratory analysis
		// predict(varname) = computes the predicted probability of the observed bin
		// graph(name) = graph with the actual histogram vs the predicted frequency; must specify graph name
		// graphtitle(string) = graph title
		// By default, the estimation procedure restricts the estimation sample to observations with earnings within +- est_bins*bin_width; use if to impose additional requirements

// See the paper for usage examples. If you find any bugs, please let me know at dbyzalov@temple.edu

// Change log
// July 13, 2018 -- redefined the em_type option as em_type(i) and em_type(ii) for consistency with the model labels I and II (previously, this option was em_type(0) for Model I and em_type(1) for Model II, which is a recipe for confusion)
//               -- incorporated em_binsMinus as an optional parameter for asymmetric small-loss and small-profit intervals in the main version of the code
// June 29, 2018 -- fixed a bug to address "name conflict" error message in newer versions of Stata. This bug did not affect any of the estimates, and it was driven by a minor compatibility difference between Stata 10 and Stata 15 in the treatment of matrix labels.
// January 2018 -- original version

	qui{
		noisily dis in green "Start... " _continue
		marksample touse

		if "`smooth'"=="" local smooth="smooth"
		if "`kinky'"=="" local kinky="kinky"

		if "`em_type'"=="i" local em_type = 0
		if "`em_type'"=="ii" local em_type = 1

		if `em_binsMinus'==0 local em_binsMinus=`em_bins'

		tokenize `varlist'  		// `Y' is the dep var, `Xvars' are the original Xs (before interacting them with bin-specific vars and discontinuity-specific vars)
		local Y "`1'"
		macro shift 1
		local Xvars "`*'"

		tempvar isEstSample  // full estimation sample for both stages
		gen `isEstSample' = `touse'*(`Y'>=-`binwidth'*`est_bins' & `Y'<`binwidth'*`est_bins')     // inclusive (exclusive) of lower (upper) bound for consistency with the bin defs below

		tempvar isPredSample // the prediction sample comprises both the estimation sample and the cross-validation sample (all other obs in the same bins, if the est sample is restricted by IF)
		if "`predict'" != ""{
  			gen `isPredSample' = (`Y'>=-`binwidth'*`est_bins' & `Y'<`binwidth'*`est_bins')
		}
		else gen `isPredSample' = `isEstSample'

		// generate a separate row for each possible bin*observation combination (binID*obsID)
		tempvar obsID binID   // obsID indexes observations in the original data, binID indexes bins within a given obsID
		gen `obsID' = _n
		expand 2*`est_bins' if `isPredSample'
		sort `obsID'
		by `obsID': gen `binID' = _n - `est_bins' if `isPredSample' // binID=0 is just below zero, =1 is just above zero

		tempvar dumYinBin  // =1 if Y falls in bin #binID for the current row
		gen `dumYinBin' = (`Y'>=(`binID'-1)*`binwidth')*(`Y'<`binID'*`binwidth') if `isPredSample' // compute for full predSample to be able to identify relevant bin X-s for predictions

		tempvar isJustBelowBins isJustAboveBins isEMbins  // all 3 are computed for the full sample for more convenient sample selection later on
		gen `isJustBelowBins' = (`binID'>-`em_binsMinus' & `binID'<=0)*`isPredSample' // XXX EM bins just below zero
		gen `isJustAboveBins' = (`binID'>0 & `binID'<=`em_bins')*`isPredSample'  // XXX EM bins just above zero
		gen `isEMbins' = `isJustBelowBins' + `isJustAboveBins'

		tempvar isEstSample_stage1 isEstSample_stage2  // estimation sample for the respective stage
		gen `isEstSample_stage1' = (!`isEMbins' & `isEstSample')
		gen `isEstSample_stage2' = (`isEMbins' & `isEstSample')

		/********** prepare bin-related variables for stage 1 ***********/
		noisily dis  in green "stage 1... " _continue

		capture drop `smooth'0 `smooth'1
		gen `smooth'0 = 1 if `isPredSample'  // intercept -- using consistent notation to relate it to the smooth pdf
		gen `smooth'1 = (`binID'-0.5)*`binwidth' if `isPredSample'  // value of Y at the bin midpoint
		forval i=2/`degree'{
			capture drop `smooth'`i'
			gen `smooth'`i' = (10^(`i'-1))*(`smooth'1)^`i' if `isPredSample'
		}

		local lastVarName = "`smooth'`degree'"
		if "`Xvars'"!=""{
			foreach binvar of varlist `smooth'0-`smooth'`degree'{
				foreach xvar of varlist `Xvars'{
					capture drop `binvar'_`xvar'
					gen `binvar'_`xvar' = `binvar'*`xvar' if `isPredSample'
					local lastVarName = "`binvar'_`xvar'"
				}
			}
		}

		local Z_stage1 "`smooth'0-`lastVarName'"

		/************** stage 1 regression *************/
		tempvar bigEps predY_stage1
		tempname b1 e_N1 df_r1 df_m1
		reg `dumYinBin' `Z_stage1' if `isEstSample_stage1', noconst
		scalar `e_N1'=e(N)
		matrix `b1' = e(b)
		sca `df_r1' = e(df_r)
		sca `df_m1' = e(df_m)
		predict `bigEps' if `isEstSample_stage1', resid  // I store the residuals from both stages in the same variable for the computation of cross-stage covariances
		predict `predY_stage1' if `isPredSample', xb

		/************** prepare variables for stage 2 *****************/
		noisily dis in green "stage 2... " _continue

		tempvar Y_stage2
		gen `Y_stage2' = `dumYinBin'-`predY_stage1' if `isEstSample_stage2'

		tempvar D  // D determines the shape of discontinuity; by construction, D has mean -1 in the EM bins below zero; and XXX +1*em_binsMinus/em_bins in the EM bins above zero
		gen `D' = ((`em_binsMinus'/`em_bins')*`isJustAboveBins'-`isJustBelowBins') if `isEMbins'  // flat increment -- compute for all EM bin obs in isPredSample
		if `em_type'==1 {
			replace `D' = `D'*2*(1 - abs(`binID'-0.5)/`em_binsMinus') if `isJustBelowBins' // XXX triangular increment
			replace `D' = `D'*2*(1 - abs(`binID'-0.5)/`em_bins') if `isJustAboveBins' // XXX
		}

		capture drop `kinky'
		gen `kinky' = `predY_stage1'*`D' if `isJustBelowBins'
		sort `obsID' `binID'

		tempvar selFirstBelow selOtherBelow selFirstAbove selOtherAbove
		gen `selFirstBelow' = (`isJustBelowBins' & !`isJustBelowBins'[_n-1])
		gen `selOtherBelow' = (`isJustBelowBins' & !`selFirstBelow')
		gen `selFirstAbove' = (`isJustAboveBins' & !`isJustAboveBins'[_n-1])
		gen `selOtherAbove' = (`isJustAboveBins' & !`selFirstAbove')

		tempvar tmpMeanBelow meanBelow
//		by `obsID': egen `tmpMeanBelow' = mean(`kinky')           // mean can vary across obsID due to X
//		by `obsID': egen `meanBelow' = mean(`tmpMeanBelow')
         // a quicker way to do the commented-out egen above (egen takes a lot of time) -- MUST BE SORTED BY `obsID' `binID' before this computation!!!
		gen `meanBelow' = `kinky' if `selFirstBelow'
		replace `meanBelow' = `kinky'+`meanBelow'[_n-1] if `selOtherBelow'
		replace `meanBelow' = `meanBelow'[_n-1]/`em_binsMinus' if `selFirstAbove'
		replace `meanBelow' = `meanBelow'[_n-1] if `selOtherAbove'

		replace `kinky' = -`meanBelow'*`D' if `isJustAboveBins'   // by construction, D has the offsetting mean in isJustAboveBins

		if "`Xvars'"!=""{
			local firstVarName = ""
			foreach var of varlist `Xvars'{
				capture drop `kinky'_`var'
				gen `kinky'_`var' = `kinky'*`var' if `isEMbins'
				if "`firstVarName'"==""  local firstVarName = "`kinky'_`var'"
				local lastVarName = "`kinky'_`var'"
			}
			local W_stage2 = "`kinky' `firstVarName'-`lastVarName'"
		}
		else{
			local W_stage2 = "`kinky'"
		}


		/************** stage 2 regression *************/
		tempvar predY_stage2 predY
		tempname b2 e_N2 df_r2 df_m2 N_clust v2
		reg `Y_stage2' `W_stage2' if `isEstSample_stage2', noconst cluster(`cluster')  // clustered only to get # clusters for final report
		scalar `e_N2'=e(N)
		matrix `b2' = e(b)
		matrix `v2' = e(V)
		sca `df_r2' = e(df_r)
		sca `df_m2' = e(df_m)
		sca `N_clust' = e(N_clust)

		predict `predY_stage2' if `isEMbins', xb
		replace `bigEps' = `Y_stage2' - `predY_stage2' if `isEstSample_stage2'

		gen `predY' = `predY_stage1' if `isPredSample'
		replace `predY' = `predY' + `predY_stage2' if `isEMbins'

		/************ compute predicted probability for the observed bin *************/
		if "`predict'"!=""{
			tempvar tmpPredY sumPredY
			by `obsID': egen `sumPredY' = total(`predY') if `isPredSample'
			gen `tmpPredY' = `predY'/`sumPredY' if `dumYinBin'==1  // it is essential to scale by `sumPredY' -- otherwise, a model with upward-biased prob in the EM interval will seem to have better KLIC (note that we can get such upward bias in two-stage estimation)
			capture drop `predict'
			by `obsID': egen `predict' = mean(`tmpPredY') if `isPredSample'
		}

		/************** predicted vs actual histogram plot ****************/
		if "`graph'" != ""{
			capture graph drop `graph'
			tempvar actFreq predFreq y
			sort `binID'
			by `binID': egen `actFreq' = mean(`dumYinBin') if `isEstSample'
			by `binID': egen `predFreq' = mean(`predY') if `isEstSample'

			// compute continuous predicted pdf for the graph
			tempvar sm predContY contD
			tempname coef
			gen `predContY' = 0
			gen `y' = 1 if `isEstSample'

			forval i=0/`degree'{
				mat `coef' = `b1'["y1", "`smooth'`i'"]
				replace `predContY' = `predContY' + `coef'[1,1]*`y' if `isEstSample'
				if "`Xvars'"!=""{
					foreach xvar of varlist `Xvars'{
						mat `coef' = `b1'["y1", "`smooth'`i'_`xvar'"]
						replace `predContY' = `predContY' + `coef'[1,1]*`y'*`xvar' if `isEstSample'
					}
				}
				if `i'==0{
					replace `y'=`Y' if `isEstSample'
				}
				else{
					replace `y'=`y'*`Y'*10 if `isEstSample'
				}
			}
			gen `contD' = (`Y'>=0 & `Y'<`em_bins'*`binwidth')*(`em_binsMinus'/`em_bins') - (`Y'>=-`em_binsMinus'*`binwidth' & `Y'<0) if `isEstSample' // XXX
			if `em_type'==1 {
				replace `contD' = `contD'*2*(1 - abs(`Y')/(`em_bins'*`binwidth') ) if `isEstSample' & `Y'>=0  // XXX
				replace `contD' = `contD'*2*(1 - abs(`Y')/(`em_binsMinus'*`binwidth') ) if `isEstSample' & `Y'<0 // XXX
			}

			tempvar contKinky
			gen `contKinky' = 0
			replace `contKinky' =  `predContY'*`contD' if `Y'>=-`em_binsMinus'*`binwidth' & `Y'<0
			replace `contKinky' = -`meanBelow'*`contD' if `Y'>=0 & `Y'<`em_bins'*`binwidth'

			mat `coef' = `b2'["y1", "`kinky'"]
			replace `predContY' = `predContY' + `coef'[1,1]*`contKinky' if `isEstSample'
			if "`Xvars'"!=""{
				foreach xvar of varlist `Xvars'{
					mat `coef' = `b2'["y1", "`kinky'_`xvar'"]
					replace `predContY' = `predContY' + `coef'[1,1]*`contKinky'*`xvar' if `isEstSample'
				}
			}

			tempvar selContGraph
			gen `selContGraph' = (`isEstSample' & `dumYinBin'==1)
			local selProb = 10000 / ( (`e_N1'+`e_N2') / (2*`est_bins') )  // reduce #obs on the continuous graph to ~2000
			replace `selContGraph'= `selContGraph'*(runiform()<`selProb')

			sum `obsID' if `isEstSample'
			local firstObsID = r(min)  // identify first obs in the estimation sample
			twoway (bar `actFreq' `smooth'1 if `obsID'==`firstObsID', barwidth(`binwidth') ysc(r(0) noline) color(gs12) lcolor(black) legend(off) ylabel(none) ytick(none))  || (connected `predContY' `Y' if `selContGraph', ysc(r(0) noline) lcolor(blue) lwidth(vthick) msymbol(none) mcolor(blue) legend(off) sort(`Y') ylabel(none) ytick(none)), name(`graph') title("`graphtitle'") plotregion(fcolor(white) margin(zero)) graphregion(fcolor(white))

		}



		// prepare the Jacobian matrices
		// after multiple re-writes to speed up the two-stage SE computation, the code below might be unreadable
		tempname A B C BC1 tmpZeros bigM
		matrix accum `A' = `Z_stage1' if `isEstSample_stage1', noconst // -N*dMomentCond1/dTheta1

		matrix accum `BC1' = `Z_stage1' `W_stage2' if `isEstSample_stage2', noconst
		matrix `C' = `BC1'[(`df_m1'+1)..., (`df_m1'+1)...]

		// -N*dMomentCond2/dTheta1 -- this one is messier

		if "`sloppy'"==""{ // sloppy single-stage derivatives -- use this only when you want to save time in quick-and-dirty preliminary analysis
			noisily dis in green "cross-derivatives..." _continue

			tempname BC2 tmpB

			tempvar dG0_dkinky tmpOne tmpSum
			gen `dG0_dkinky' = (`bigEps' - `predY_stage2')*`D' if `isEstSample_stage2'   // dMomentCond2[0]/dkinky (the derivs for other stage-2 moments are just scaled up by X)
			gen `tmpOne' = 1 if `isEstSample_stage2'

			foreach var of varlist `tmpOne' `Xvars'{
				noisily dis in green "." _continue

				gen tmpW_`var' = `var'*`dG0_dkinky' if `isEstSample_stage2'

				gen `tmpSum' = 0 if `selFirstAbove'
				forval i=1/`em_bins'{ // XXX
					replace `tmpSum' = `tmpSum' + tmpW_`var'[_n+`i'-1] if `selFirstAbove' & `isEstSample_stage2'
				}
				forval i=1/`em_binsMinus'{ // XXX
					replace tmpW_`var'=tmpW_`var'-`D'*`tmpSum'[_n+`i']/`em_binsMinus' if `selFirstAbove'[_n+`i'] & `isEstSample_stage2'
				}
				drop `tmpSum'
			}

			matrix accum `BC2' = `Z_stage1' tmpW_* if `isJustBelowBins' & `isEstSample_stage2', noconst
			drop tmpW_*

			matrix `tmpB' = `BC1'-`BC2'
			matrix `B' = `tmpB'[(`df_m1'+1)..., 1..`df_m1']
		}
		else{
			mat `B' = J(`df_m2',`df_m1',0)
		}

		mat `tmpZeros' = J( rowsof(`A'), colsof(`C'), 0)

		mat `bigM' = (`A',`tmpZeros')\(`B',`C')    // the combined Jacobian for both stages

		/*************** compute the clustered covariance matrix of the moments from both steps ****************/

		tempname bigX bigOmega bigV // combines explanatory variables from both stages (filled in with zeros for irrelevant observations for the respective stage)
					  // for the computation of appropriately clustered covariance matrix of the moment conditions

		if "`sloppy'"==""{
			noisily dis in green " valid two-stage cov(Theta)... " _continue

			foreach var of varlist `Z_stage1'{
				replace `var' = 0     if `isEstSample_stage2'  // I can start messing with these vars like this here for the computation of Omega -- I am not using them for anything else below
			}
			foreach var of varlist `W_stage2'{
				replace `var' = 0     if `isEstSample_stage1'
			}

			sort `cluster'
			matrix opaccum `bigOmega' = `Z_stage1' `W_stage2' if `isEstSample', group(`cluster') opvar(`bigEps') noconst
		}
		else{
			noisily dis in green " sloppy single-stage cov(Theta)... " _continue

			tempname omega1 omega2
			sort `cluster'

			matrix opaccum `omega1' = `Z_stage1' if `isEstSample_stage1', group(`cluster') opvar(`bigEps') noconst

			matrix opaccum `omega2' = `W_stage2' if `isEstSample_stage2', group(`cluster') opvar(`bigEps') noconst
			matrix `bigOmega' = (`omega1',`tmpZeros')\(`tmpZeros'',`omega2')
			local rname1 : rowfullnames `omega1'		// copy correct row/column names for bigOmega -- otherwise, newer versions of Stata complain about "name conflict"
			local cname1 : colfullnames `omega1'
			local rname2 : rowfullnames `omega2'
			local cname2 : colfullnames `omega2'

			matrix rownames `bigOmega' = `rname1' `rname2'
			matrix colnames `bigOmega' = `cname1' `cname2'
		}
		matrix `bigV' = inv(`bigM')*`bigOmega'*inv(`bigM'')  // Correct Cov(theta) for two-stage estimation
		local rname : rowfullnames `bigOmega'		// copy row/column names from bigOmega to bigV --- otherwise, newer versions of Stata complain about "name conflict"
		local cname : colfullnames `bigOmega'
		matrix rownames `bigV' = `rname'
		matrix colnames `bigV' = `cname'

		// compute ~R2 for the full estimation sample
		tempvar dev2
		tempname varR varY tmpN total_R2 total_adjR2
		gen `dev2' = `bigEps'^2 if `isEstSample'
		sum `dev2'
		sca `varR' = r(mean)  // cannot just use canned Var(bigEps) -- mean(bigEps) may be non-zero in the EM bins
		sum `dumYinBin' if `isEstSample'
		sca `varY' = r(Var)
		sca `tmpN' = r(N)
		sca `varY' = `varY'*(`tmpN'-1)/`tmpN'

		sca `total_R2' = 1 - `varR'/`varY'
		sca `total_adjR2' = 1 - (1-`total_R2')*(`tmpN'-1)/(`tmpN'-`df_m1'-`df_m2')


		drop if `isPredSample' & `binID'!=0  // these are the duplicate copies of the original observations

		ereturn clear
		tempname b v
		mat `b' = `b1',`b2'
		mat `v' = `bigV'
		ereturn post `b' `v', esample(`isEstSample')
		ereturn scalar N = ( (`e_N1'+`e_N2') / (2*`est_bins') )
		ereturn scalar N_clust = `N_clust'
		ereturn scalar df_m = (`df_m1'+`df_m2')
		ereturn scalar df_r = `N_clust' - (`df_m1'+`df_m2')
		ereturn scalar r2 = `total_R2'
		ereturn scalar r2_a = `total_adjR2'
		ereturn local cmdline `"`0'"'
		ereturn local cmd "cmdname"
	} //end qui

	dis ""
	dis ""
	dis in yellow "Bin width = `binwidth'"                 _column(56) "Number of observations = " e(N)
	dis in yellow "Number of bins (on each side of zero)"  _column(56) "Number of clusters     = " e(N_clust)
	dis in yellow "  estimation interval    = `est_bins'"  _column(56) "Cluster variable = `cluster'"
	dis in yellow "  discontinuity interval(-/+) = `em_binsMinus' `em_bins'"   _column(56) "R2     = " %7.6f `total_R2' // XXX
	dis in yellow "Polynomial degree = `degree'"           _column(56) "adj R2 = " %7.6f `total_adjR2'

	if `em_type'==1  dis in yellow "Increment type = TRIANGULAR"
	else  dis in yellow "Increment type = FLAT"

	if "`sloppy'"=="" dis in yellow "Valid two-stage SE"
	else dis in yellow "Sloppy single-stage SE"

	ereturn display
end


/* a meaningless comment at the end of file -- without it, Stata can produce cryptic error messages because it is confused by end of file for some reason */