*! pheatplot 1.1.0 02 July, 2024 Nicolai Topstad Borgen (n.t.borgen@isp.uio.no)
*! pheatplot 1.0.0 27 September, 2023 Nicolai Topstad Borgen (n.t.borgen@isp.uio.no)
		
program define pheatplot, rclass

	version 17 
	
	syntax varlist(max=1), 	 									///
		[interaction(varname fv) frame(string) 					///
		BName(string asis) PName(string asis) 					///
		heatoptsp(string asis) heatoptsb(string asis)			///
		pvalues(string) SAVETable(string asis) 					///
		SAVEGraph(string asis)									///
		hexplot differences mono threshold(real .10)] 
	
	local xvar `varlist'
	
	if "`hexplot'"=="hexplot" local plot hexplot 
	if "`hexplot'"==""        local plot heatplot

		
	* Package dependency: heatplot *
	
	capture which heatplot 
	if _rc {
		di as result in smcl "Please install user-contributed package {it:heatplot} from SSC;"		///
			_newline `"click the following link to install package: {stata "ssc install heatplot":ssc install heatplot}"'
			exit 498
	}
	
	capture which colorpalette
	if _rc {
		di as result in smcl "Please install user-contributed package {it:colorpalette} from SSC;"		///
			_newline `"click the following link to install package: {stata "ssc install palettes":ssc install palettes}"'
			exit 498
	}	
	
	tempname errormata errorstata
	mata: `errormata'=findexternal("ColrSpace()") != J(1,1,NULL)
	mata: st_matrix("`errorstata'",`errormata')
	if `errorstata'[1,1]==0 {
		di as result in smcl "Please install user-contributed package {it:ColrSpace} from SSC;"		///
			_newline `"click the following link to install package: {stata "ssc install colrspace":ssc install colrspace}"'
			exit 498
	}	
	
	
	* Error messages *	
	if "`differences'"=="" & "`heatoptsb'"!="" {
		di as error in smcl "Option {it:heatoptsb()} is not allowed without option {it:differences}." 
		exit 498
	}
	
	if "`differences'"=="" & "`bname'"!="" {
		di as error in smcl "Option {it:bname()} is not allowed without option {it:differences}." 
		exit 498
	}
	
	if !inrange(`threshold',0.01,0.99) {
		di as error in smcl "Option {it:threshold()} must be a real number between 0.01 and 0.99." 
		exit 498
	}
		
	
	* Identify type of factor variables *  	
		
	if substr("`interaction'",2,1)=="." {
		local pre=substr("`interaction'",1,1)
		local ivar=substr("`interaction'",3,.)
	}
	else {
		local ivar `interaction'
	}
	
	if "`interaction'"!="" {
		qui margins i.`xvar', dydx(`ivar') post
	}
			
	
	* Names of matrices, frame, and graphs *
	
	tempname p coefficients frame pcolorclr bcolorclr pmonoclr bmonoclr
	
	local pmatrix `p'
	local bmatrix `coefficients'
	
	if "`frame'"==""		local frame `frame'
	if "`pname'"==""		local pname pvalues, replace
	if "`bname'"==""		local bname differences, replace

	
	* Create frame and matrices (to be filled below) * 
		
	frame create `frame' y x str50 Y str50 X pvalue estimate se 
	
	qui levelsof `xvar' if e(sample), local(l)
	matrix `pmatrix'=J(r(r),r(r),.)
	matrix `bmatrix'=J(r(r),r(r),.)
	local num=r(r)

	
	* Fill frame and matrices * 
	
	local labs 
	
	local r=0
	foreach i of local l {
		local ++r
		
		* Error message if use has not used factor variables *
		if `r'==1 {
			capture di _b[`i'.`xvar']
			if _rc {
				di as error in smcl "Variable `xvar' was not included using factor variable notation in the previous regression model." 	///
				_newline "Run regression again using factor variable notation (i.e., i.`xvar'), and then use {it:pheatplot}." 								///
				_newline `"For information about factor variables, see {stata "help fvvarlist":help fvvarlist}."'
				exit 498
			}
		}
		
		local labs `labs' `r' "`: label (`xvar') `i''"
		if `r'!=1 local ylabs `ylabs' `r' "`: label (`xvar') `i''"
		if `r'!=`num' local xlabs `xlabs' `r' "`: label (`xvar') `i''"
		
		local c=0
		foreach j of local l {
			local ++c
			if `i'==`j' continue
			local lab1 `: label (`xvar') `i''
			local lab2 `: label (`xvar') `j''
			if "`interaction'"=="" {
				qui lincom `i'.`xvar'-`j'.`xvar' 
			}
			if "`interaction'"!="" {
				if "`pre'"=="i" {
					qui lincom [1.`ivar']`i'.`xvar'-[1.`ivar']`j'.`xvar' 
				}
				if "`pre'"=="" {
					qui lincom [`ivar']`i'.`xvar'-[`ivar']`j'.`xvar' 
				}
			}
			frame post `frame' (`i') (`j') ("`lab1'") ("`lab2'")	///
				(r(p)) (r(estimate)) (r(se))
			matrix `pmatrix'[`r',`c']=r(p)
			matrix `bmatrix'[`r',`c']=r(estimate)
		}
	}
	
	* Default colors *
	
	mata: `pmonoclr'=ColrSpace()
	mata: `pmonoclr'.colors("#FFFFFF #C7C7C7 #A6A6A6 #242424")
	mata: `pmonoclr'.ipolate(100, "", ., ., (0, `threshold'-.01, `threshold', 1))
	
	mata: `pcolorclr'=ColrSpace()
	mata: `pcolorclr'.colors("#3A90FE #EDEDED #A89008")
	mata: `pcolorclr'.ipolate(100, "", ., ., (0, `threshold', 1))
	
	mata: `bmonoclr'=ColrSpace()
	mata: `bmonoclr'.colors("#FFFFFF #242424")
	mata: `bmonoclr'.ipolate(100)
	
	mata: `bcolorclr'=ColrSpace()
	mata: `bcolorclr'.colors("#3A90FE #EDEDED #A89008")
	mata: `bcolorclr'.ipolate(100)
	
	if "`mono'"=="" {
		local pcolor color(mata(`pcolorclr'), intensity(.90))
		local bcolor color(mata(`bcolorclr'))
	}
	if "`mono'"=="mono" {
		local pcolor color(mata(`pmonoclr'), intensity(.8))
		local bcolor color(mata(`bmonoclr'))
	}
	
	
	* Default options * 
	
	local pramp ramp(right space(20) label(0(.1)1) subtitle(P-value) format(%9.2f) length(55))
	local paspectratio aspectratio(1)
	local pylabel ylabel(`ylabs',nogrid)
	local pxlabel xlabel(`xlabs', nogrid angle(25))
	local pxtitle xtitle("") 
	local pytitle ytitle("") 
	local plower lower 
	local pdiagonal nodiagonal
	local pcuts cuts(0(.01)1)
	if "`pvalues'"=="off" local values 
	if "`pvalues'"=="" local values values(format(%9.2f)) 
		
	local bramp ramp(right space(15) label(#10) subtitle(Diff) length(55))
	local baspectratio aspectratio(1)
	local bylabel ylabel(`ylabs',nogrid)
	local bxlabel xlabel(`xlabs', nogrid angle(25))
	local bxtitle xtitle("") 
	local bytitle ytitle("") 
	local blower lower 
	local bdiagonal nodiagonal
		
	* Parse user-written options (takes precedence over default options) * 
	
	if "`heatoptsb'"!="" local useroptsb: subinstr local heatoptsb "," "", all
	if "`heatoptsp'"!="" local useroptsp: subinstr local heatoptsp "," "", all
	
	tokenize `useroptsp'
	local c=0 
	while `c'!=-1 {
		local ++c 
		
		local opt=substr("``c''",1,strpos("``c''","(")-1)
		if substr("`opt'",1,1)=="" local opt "``c''"
			
		if "`opt'"=="ramp" 					local pramp
		if substr("`opt'",1,6)=="aspect" 	local paspectratio
		if substr("`opt'",1,4)=="ylab"		local pylabel 
		if substr("`opt'",1,4)=="xlab"		local pxlabel 
		if substr("`opt'",1,1)=="c"			local pcolor 
		if substr("`opt'",1,3)=="xti"		local pxtitle 
		if substr("`opt'",1,3)=="yti"		local pytitle 
		if "`opt'"=="upper"					local plower 
		if "`opt'"=="diagonal"				local pdiagonal 
		if substr("`opt'",1,3)=="val"		local pvalues
		if substr("`opt'",1,3)=="cut"		local pcuts
	
		if "``c''"=="" local c -1	
		
	}
	
	tokenize `useroptsb'
	local c=0 
	while `c'!=-1 {
		local ++c 
		
		local opt=substr("``c''",1,strpos("``c''","(")-1)
		if substr("`opt'",1,1)=="" local opt "``c''"
		
		if "`opt'"=="ramp" 					local bramp
		if substr("`opt'",1,6)=="aspect" 	local baspectratio
		if substr("`opt'",1,4)=="ylab"		local bylabel 
		if substr("`opt'",1,4)=="xlab"		local bxlabel 
		if substr("`opt'",1,1)=="c"			local bcolor 
		if substr("`opt'",1,3)=="xti"		local bxtitle 
		if substr("`opt'",1,3)=="yti"		local bytitle 
		if "`opt'"=="upper"					local blower 
		if "`opt'"=="diagonal"				local bdiagonal 
		if substr("`opt'",1,3)=="val"		local bvalues
	
		if "``c''"=="" local c -1	
		
	}
	
	
	* Heatplots * 
	
	if "`differences'"=="differences" {
		`plot' `bmatrix', `heatoptsb'  `bramp' `baspectratio'				///
			`bylabel' `bxlabel' `bcolor' `bxtitle' `bytitle' `blower'		///
			`bdiagonal' name(`bname')	
	}
	
	`plot' `pmatrix', `heatoptsp' `values' `pramp'  `paspectratio' 			///
		`pylabel' `pxlabel' `pcolor' `pxtitle' `pytitle' `plower'			///
		`pdiagonal' `pcuts' name(`pname') 

	
	* Table (optional) *
		
	if `"`savetable'"'!="" {
		frame `frame' {
			putdocx clear
			putdocx begin, font(,10) 
			putdocx table tbl=data(Y X estimate se pvalue), border(all,nil)	///
				layout(autofitcontents) varnames							///
				note("Note: The difference is calculated as the coefficient for group 1 minus the coefficient for group 2 using lincom, while the standard error is the standard error of this difference.")
			putdocx table tbl(1,1) = ("Group 1")
			putdocx table tbl(1,2) = ("Group 2")
			putdocx table tbl(1,3) = ("Difference")
			putdocx table tbl(1,4) = ("Standard error")
			putdocx table tbl(1,5) = ("P-value")
			putdocx table tbl(1/1,.), bold border(top) border(bottom)
			putdocx table tbl(.,3), nformat(%9.4f)
			putdocx table tbl(.,4), nformat(%9.4f)
			putdocx table tbl(.,5), nformat(%9.3f)
			local last=_N+1 
			putdocx table tbl(`last',.), border(bottom)
			putdocx save `savetable'
		}
	}

	* Graph (optional) *
	if `"`savegraph'"'!="" {
		graph export `savegraph'
	}
	
		
	* Return matrices *
	
	return matrix differences=`bmatrix'
	return matrix pvalues=`pmatrix'
		
end