/*
History:
2026-02-04 Update to allow for two-way interactions.
2026-01-06 Update macro to allow for categorical covariates.
           The options cat_covariates, cat_reference, and cat_maxn have been added.

Gelman-Rubin Test for MCMC Convergence

Description:

	Tests an MCMC model specification for convergence using the
	Gelman-Rubin test. A model is fit multiple times with the specified
	parameters using different random seeds each time.
	
Parameters:

	- num_chains: Integer number of MCMC chains to run. (default = 5)
	- initial_mcmc_seed: Numeric starting seed for the random number generator. 
											 If specified, the seed will be incremented for each MCMC replicate. 
											 If blank, uses a randomly generated integer from -10^7 to 10^7, exclusive for each replicate. (default = blank)
	- outlib: The SAS library to store output datasets. (default = WORK)
	- outname: The name of the output dataset.
	- cat_covariates: Space-delimited list of names of categorical covariates. Any character variable
	                  in covariates will be treated as categorical. For each categorical covariate,
					  dummy variables will be created. See cat_reference below for specifying
					  the reference categories.
	- cat_reference:  Space-delimited list of reference categories for the categorical variables
                      in cat_covariates. The length and order of cat_reference must match the 
                      length and order of cat_covariates.					  
	- cat_maxn: The maximum number of categories for a categorical variable. The default is 25.
	- DEBUG:  0 or 1 for debugging purposes. The default is 0.
	- v143code: Flag to run alternative code for IML version 14.3. (0/1, default is determined at run time)
	
	See %nci_multivar_mcmc() for other parameters.
	
Output:

	The following SAS data set is created in outlib:
	
		outname: A SAS data set with observation containing Gelman-Rubin statistics for each MCMC parameter.
	
Details:

	The Gelman-Rubin test works by creating multiple MCMC chains with
	different random seeds and calculating the within-chain and between-chain
	variation (Gelman and Rubin, 1992). If the model parameters converge, there
	should be little to no difference between different chains which will cause
	the between-chain variance to fall to zero. This means that convergence can
	be assessed by comparing the total variance of the model parameters to the
	within-chain variance.
	
	The Gelman-Rubin statistic is the square root of the ratio between the
	total variance and the within-chain variance of a parameter. If the
	Gelman-Rubin statistic is close to 1, then the model has converged for that
	parameter. A cutoff of 1.1 for convergence is suggested by Gelman, et al.
	(2004).
*/
%macro gelman_rubin(num_chains=5,
										initial_mcmc_seed=,
										pre_mcmc_lib=WORK,
										pre_mcmc_data=,
										id=,
										repeat_obs=,
										weight=,
										episodic_variables=,
										episodic_indicators=,
										episodic_amounts=,
										daily_variables=,
										daily_amounts=,
										never_consumer_variable=,
										never_consumer_indicator=,
										never_consumer_amount=,
										default_covariates=,
										default_intercept=Y,
										episodic_indicator_covariates=,
										episodic_indicator_intercept=,
										episodic_amount_covariates=,
										episodic_amount_intercept=,
										daily_amount_covariates=,
										daily_amount_intercept=,
										individual_covariates=,
										individual_intercept=,
										never_consumer_covariates=,
										never_consumer_intercept=Y,
										num_mcmc_iterations=12000,
										num_burn=2000,
										num_thin=25,
										sigma_u_prior=,
										sigma_u_constant=N,
										low_memory_mode=N,
										outlib=WORK,
										outname=,
										cat_covariates=,
										cat_reference=,
										cat_maxn=25,
										DEBUG=0,
										v143code=);
										
	%local mcmc_seed i dat allv catvdata intvdata;
										
	%let mcmc_seed = &initial_mcmc_seed.;
	
	**Set the value for v143code if needed;
	%let v143code = %left(%trim(&v143code));
	%if %length(&v143code) = 0 %then %do;
		%store_IML_version(_tmp_dat_);
		%get_v143code(_tmp_dat_);
		%let v143code = 0;
	%end;	
	
	** Create dummy variables for categorical covariates;
	%local dat allv catvdata;
	%let dat      = &pre_mcmc_lib..&pre_mcmc_data._mcmc_in; /* input data */
	%let catvdata = _catv_info_tmp; /* output data of categorical variables */
	%let allv     = &default_covariates &episodic_indicator_covariates &episodic_amount_covariates;
	%let allv     = &allv &daily_amount_covariates &never_consumer_covariates;
	%let allv     = %normStr(&allv, upcase=1, removeMultBlanks=1);
	%createDummyVariables(&catvdata, &dat, &cat_covariates, refcats=&cat_reference, 
	                      allvars=&allv, maxncats=&cat_maxn, suffix=_CAT, DEBUG=&DEBUG);		
	**Update covariate lists for categorical covariates;	
	%let default_covariates            = %updateListWithDummyVars(&default_covariates,            &catvdata, suffix=_CAT);	
	%let episodic_indicator_covariates = %updateListWithDummyVars(&episodic_indicator_covariates, &catvdata, suffix=_CAT);		
	%let episodic_amount_covariates    = %updateListWithDummyVars(&episodic_amount_covariates,    &catvdata, suffix=_CAT);	
	%let daily_amount_covariates       = %updateListWithDummyVars(&daily_amount_covariates,       &catvdata, suffix=_CAT);	
	%let never_consumer_covariates     = %updateListWithDummyVars(&never_consumer_covariates,     &catvdata, suffix=_CAT);	
	%updateDataWithDummyVars(&individual_covariates, &catvdata, var="covariates", suffix=_CAT);	
	
	**Update covariates list for interactions;
	%let intvdata = _intv_info_tmp; /* output data of interaction variables */	
	%createInteractions(&intvdata, &dat, &allv, &catvdata, sep=X);
	%let default_covariates            = %updateListWithInteractions(&default_covariates,            &intvdata);
	%let episodic_indicator_covariates = %updateListWithInteractions(&episodic_indicator_covariates, &intvdata);
	%let episodic_amount_covariates    = %updateListWithInteractions(&episodic_amount_covariates,    &intvdata);
	%let daily_amount_covariates       = %updateListWithInteractions(&daily_amount_covariates,       &intvdata);
	%let never_consumer_covariates     = %updateListWithInteractions(&never_consumer_covariates,     &intvdata);
	%updateDataWithInteractions(&individual_covariates, &intvdata, var="covariates");
	
	%do i = 1 %to &num_chains.;
	
		**Fit MCMC model;
		%nci_multivar_mcmc(pre_mcmc_lib=&pre_mcmc_lib.,
											 pre_mcmc_data=&pre_mcmc_data.,
											 id=&id.,
											 repeat_obs=&repeat_obs.,
											 weight=&weight.,
											 episodic_variables=&episodic_variables.,
											 episodic_indicators=&episodic_indicators.,
											 episodic_amounts=&episodic_amounts.,
											 daily_variables=&daily_variables.,
											 daily_amounts=&daily_amounts.,
											 default_covariates=&default_covariates.,
											 episodic_indicator_covariates=&episodic_indicator_covariates.,
											 episodic_amount_covariates=&episodic_amount_covariates.,
											 daily_amount_covariates=&daily_amount_covariates.,
											 individual_covariates=&individual_covariates.,
											 default_intercept=&default_intercept.,
											 episodic_indicator_intercept=&episodic_indicator_intercept.,
											 episodic_amount_intercept=&episodic_amount_intercept.,
											 daily_amount_intercept=&daily_amount_intercept.,
											 individual_intercept=&individual_intercept.,
											 never_consumer_variable=&never_consumer_variable.,
											 never_consumer_indicator=&never_consumer_indicator.,
											 never_consumer_amount=&never_consumer_amount.,
											 never_consumer_covariates=&never_consumer_covariates.,
											 never_consumer_intercept=&never_consumer_intercept.,
											 mcmc_seed=&mcmc_seed.,
											 num_mcmc_iterations=&num_mcmc_iterations.,
											 num_burn=&num_burn.,
											 num_thin=&num_thin.,
											 sigma_u_prior=&sigma_u_prior.,
											 sigma_u_constant=&sigma_u_constant.,
											 low_memory_mode=&low_memory_mode.,
											 outname=_model,
											 DEBUG=&DEBUG,
											 v143code=&v143code);
											 
		**Calculate parameter means and variances;
		%extract_parameters(multivar_mcmc_model=_model,
												do_mean=Y,
												do_variance=Y,
												do_covariance=N,
												outname=_model);
												
		data _mcmc_mean&i.;
			set _model_means;
		run;
		
		data _mcmc_var&i.;
			set _model_var;
		run;
		
		**increment seed if it is specified;
		%if &mcmc_seed. ^= %str() %then %do;
		
			%let mcmc_seed = %eval(&mcmc_seed. + 1);
		%end;
	%end;
	
	data _mcmc_means;
		set _mcmc_mean1-_mcmc_mean&num_chains.;
	run;
	
	data _mcmc_variances;
		set _mcmc_var1-_mcmc_var&num_chains.;
	run;
	
	proc iml;
	
		num_samples = int((&num_mcmc_iterations. - &num_burn.)/&num_thin.);
		
		use _mcmc_means;
			read all var _ALL_ into mcmc_means[colname=parameter_names];
		close _mcmc_means;
		
		use _mcmc_variances;
			read all var _ALL_ into mcmc_variances;
		close _mcmc_variances;
		
		num_parameters = ncol(mcmc_means);
		
		**calculate the joint mean of all chains;
		joint_mean = mcmc_means[:,];
		
		**calculate the within-chain, between-chain, and total variances for each parameter;
		within_variance = mcmc_variances[:,];
		between_variance = (num_samples/(&num_chains. - 1)) # ((mcmc_means - joint_mean) ## 2)[+,];
		total_variance = (num_samples/(num_samples - 1)) # within_variance + ((&num_chains. + 1)/(&num_chains. * num_samples)) # between_variance;
		
		**Calculate the Gelman-Rubin statistics for each parameter;
		**if both total variance and within-chain variance are zero, the statistic is one;
		gr_statistics = j(1, num_parameters, .);
		
		if any(total_variance > 0 & within_variance > 0) then do;
		
			nonzero = loc(total_variance > 0 & within_variance > 0);
			gr_statistics[nonzero] = sqrt(total_variance[nonzero]/within_variance[nonzero]);
		end;
		
		if(any(total_variance = 0 & within_variance = 0)) then do;
		
			gr_statistics[loc(total_variance = 0 & within_variance = 0)] = 1;
		end;
		
		create &outlib..&outname. from gr_statistics[colname=parameter_names];
			append from gr_statistics;
		close &outlib..&outname.;
	quit;
%mend gelman_rubin;