*! pystacked v0.8.0
*! last edited: 14july2026
*! authors: aa/ms
*! pystacked2_p = uses python code appended to the file

// do not edit this file manually; edit pystacked1_p.ado, then run R script
// parent program
program define pystacked2_p, rclass
    version 16.0
    syntax namelist(min=1 max=2) [if] [in], [ ///
                                                            pr /// 
                                                            xb /// 
                                                            Resid /// 
                                                            class /// 
                                                            TRANSForm /// legacy option, equivalent to basexb
                                                            BASExb ///
                                                            force ///
                                                            CValid ///
                                                            ]

    // python: from pystacked_p import *

    if ("`force'"=="") {
        // check datasignature based on depvar, xvars and foldvar
        fvrevar `e(depvar)' `e(allxvars_o)' `e(foldvar)', list
        local dslist `r(varlist)'
        local dslist : list uniq dslist
        qui _datasignature `dslist'
        if "`e(datasignature)'"~="`r(datasignature)'" {
            di as err "error: data in memory has changed since last -pystacked- call"
            di as err "you are not allowed to change data in memory between -pystacked- fit and -predict-"
            di as err "use the -force- option to override"
            exit 198
        }
        // check sort order
        local sortvars : sortedby
        if "`sortvars'"~="`e(sortvars)'" {
        	di as err "error: sort order of data in memory has changed since last -pystacked- call"
            di as err "you are not allowed to change data in memory between -pystacked- fit and -predict-"
            exit 198
        }
    } 
    
    * legacy option
    if "`transform'"~="" {
        local basexb basexb
        local transform
        di as err "transform option is deprecated; use 'basexb'"
    }

    * only 1 option max
    local optcount : word count `resid' `pr' `xb' `class'
    if `optcount'>1 {
        di as err "only one of options 'pr xb resid class' allowed"
        exit 198
    }

    if "`resid'"!="" & "`basexb'`cvalid'"!="" {
        di as err "resid not allowed with: `basexb' `cvalid'"
        exit 198
    }
    
    if "`cvalid'"~="" & "`basexb'"=="" {
        di as err "error - option cvalid currently supported only with option basexb"
        exit 198
    }

    local command=e(cmd)
    if ("`command'"~="pystacked") {
        di as err "error: -pystacked_p- supports only the -pystacked- command"
        exit 198
    }
    *

    local depvar `e(depvar)'

    tokenize `namelist'
    if "`2'"=="" {                    //  only new varname provided
        local predictvar `1'
    }
    else {                            //  datatype also provided
        local vtype `1'
        local predictvar `2'
    }
    *

    marksample touse, novarlist

    if "`basexb'"=="" {
        qui gen `vtype' `predictvar' = .
    }
    
    * Get predictions
    python: post_prediction("`predictvar'","`basexb'","`cvalid'","`vtype'","`touse'","`pr'`xb'`class'")
    
    if "`resid'"!="" {
        replace `predictvar' = `depvar' - `predictvar' if `touse'
    }
    
end

* --- Python code appended below ---
python:
#! pystacked v0.8.0
#! last edited: 14july2026
#! authors: aa/ms

# Import SFI, always with stata 16
from sfi import Data,Matrix,Scalar,Macro
from sfi import SFIToolkit
import numpy as np

def post_prediction(pred_var,basexb,cvalid,var_type,touse,pred_type):

	# Start with a working flag
	Scalar.setValue("r(import_success)", 1, vtype='visible')

	# Import model from Python namespace
	try:
		from __main__ import model_object as model
		from __main__ import model_xvars as xvars
		from __main__ import model_methods as methods
		from __main__ import id as id
		from __main__ import type as type
	except ImportError:
		print("Error: Could not find pystacked estimation results.")
		Scalar.setValue("r(import_success)", 0, vtype='visible')
		return

	touse = np.array(Data.get(touse))

	if type=="class" and pred_type=="":
		pred_type="pr"
	elif type=="reg" and pred_type=="":
		pred_type="xb"
	elif type=="class" and pred_type=="xb":
		SFIToolkit.stata('di as err "xb/resid not supported with classification"')
		#"
		SFIToolkit.error(198)
	elif type=="reg" and pred_type=="class":
		SFIToolkit.stata('di as err "class not supported with regression"')
		#"
		SFIToolkit.error(198)		
	elif type=="reg" and pred_type=="pr":
		SFIToolkit.stata('di as err "pr not supported with regression"')
		#"
		SFIToolkit.error(198)	

	if basexb=="" and cvalid=="":
		# stacked prediction
		if type=="class" and pred_type == "pr":
			from __main__ import predict_proba as pred
			if pred.ndim>1:
				pred=pred[:,1]
		elif type=="class" and pred_type == "class":
			from __main__ import predict as pred
			if pred.ndim>1:
				pred=pred[:,1]
		else: 
			from __main__ import predict as pred
		pred[touse==0] = np.nan
		Data.store(var=pred_var,val=pred,obs=None)

	elif basexb!="" and cvalid=="":
		# learner predictions
		from __main__ import transform as transf
		ncol = transf.shape[1]
		for j in range(ncol):
			if var_type == "double":
				Data.addVarDouble(pred_var+str(j+1))
			else: 
				Data.addVarFloat(pred_var+str(j+1))
			pred=transf[:,j]
			predna =np.isnan(pred)
			if pred_type=="class":
				pred=(pred>0.5)*1 
				pred=pred.astype(float)
				pred[predna]=np.nan
				Data.setVarLabel(pred_var+str(j+1),"Predicted class "+" "+methods[j])
			elif type=="class":
				Data.setVarLabel(pred_var+str(j+1),"Predicted probability "+" "+methods[j])
			else:
				Data.setVarLabel(pred_var+str(j+1),"Predicted value"+" "+methods[j])
			pred[touse==0] = np.nan
			Data.store(var=pred_var+str(j+1),val=pred,obs=None)

	elif basexb!="" and cvalid!="":
		try:  
			from __main__ import cvalid as transf
			if np.isnan(transf).all():
				SFIToolkit.stata('di as res "Warning: cvalid option not available with selected final estimator;"')
				SFIToolkit.stata('di as res "         cross-validated predicted values are set to missing"')
		except ImportError:
			SFIToolkit.stata('di as err "Error: Could not find cross-validated predicted values."')
			#"
			SFIToolkit.error(198)
			return
		id = id -1
		id = id.tolist()
		ncol = transf.shape[1]
		touse =touse[id] 
		for j in range(ncol):
			if var_type == "double":
				Data.addVarDouble(pred_var+str(j+1))
			else: 
				Data.addVarFloat(pred_var+str(j+1))
			pred=transf[:,j]
			predna =np.isnan(pred)
			if pred_type=="class":
				pred=(pred>0.5)*1 
				pred=pred.astype(float)
				pred[predna]=np.nan
				Data.setVarLabel(pred_var+str(j+1),"Cross-validated predicted class "+" "+methods[j])
			elif type=="class":
				Data.setVarLabel(pred_var+str(j+1),"Cross-validated predicted probability "+" "+methods[j])
			else:
				Data.setVarLabel(pred_var+str(j+1),"Cross-validated predicted value"+" "+methods[j])
			pred[touse==0] = np.nan
			Data.store(var=pred_var+str(j+1),val=pred,obs=id)
end