/*******************************************************************************
*                                                                              *
*     Handles classification results to ensure classes are returned            *
*                                                                              *
*******************************************************************************/

*! classify
*! v 0.0.6
*! 08mar2024

// Drop program from memory if already loaded
cap prog drop classify

// Define program
prog def classify, 

	// Version statement 
	version 15
	
	// Syntax
	syntax anything(name = classes id = "Number of Classes") [if] ,			 ///   
			PStub(string asis) [ THReshold(real 0.5) PMethod(string asis)	 ///   
			POpts(string asis)]

	// Mark the sample that will be used
	marksample touse, strok
	
	// Set default prediction method if missing
	if mi(`"`pmethod'"') loc pmethod pr
	
	// Ensure classes is an integer
	if mod(`classes', 1) != 0 {
		
		// Display an error message
		di as err "The number of classes for a classification model must be" ///
		" an integer.  You specified `classes' number of classes."
		
		// Return error code and exit 
		error 126
		
	} // End of IF Block to handle invalid number of classes
	
	// Ensure classes has a value that is >= 2
	if `classes' < 2 {
		
		// Display an error message
		di as err "The number of classes specified for your model is < 2."
		
		// Return error code and exit
		error 125		
		
	} // End IF Block for invalid number of classes
	
	// Test that the threshold value is valid
	if (`threshold' >= 1 | `threshold' <= 0) {
		
		// Display an error message
		di as err "The classification threshold must be in (0, 1).  You "	 ///   
		"specified a value of `threshold'."
		
		// Return error code and exit
		error 125
		
	} // End IF Block to handle invalid threshold values
	
	// Test the number of classes:
	if `classes' == 2 {
		
		// Generate predicted values
		predict double `pstub' if `touse', `pmethod' `popts'
		
		// Replace predicted values with classes
		replace `pstub' = cond(`pstub' < `threshold' & !mi(`pstub'), 0,	 ///   
						  cond(`pstub' >= `threshold' & !mi(`pstub'), 1, .))
		
	} // End IF Block for binary classification problems
	
	// For values > 2
	else {
		
		// Generate predicted values
		predict double `pstub'_* if `touse', `pmethod' `popts'
		
		// Get the names of the variables that were just generated
		qui: ds `pstub'_*
		
		// Store varlist
		loc pvars `r(varlist)'
		
		// Identifies the highest probability value
		egen double `pstub' = rowmax(`pvars')
		
		// Loop over the predicted variable names
		foreach v in `pvars' {
			
			// Remove all letters from the variable name
			loc clsval `= ustrregexra("`v'", "^.*_", "")'
			
			// Replace the value of pstub with the class value if that is the 
			// highest probability
			replace `pstub' = `clsval' if `pstub' == `v' & !mi(`pstub')
			
		} // End of Loop over the predicted variables
		
		// Remove the individual predicted variables
		drop `pvars'
		
	} // End ELSE Block for multinomial/ordinal classification problems
	
	// Compress the classification variable
	qui: compress `pstub'
	
// End definition of the command
end