/* History:
2025-10-30 Add DEBUG, check the existence of files loaded, and add code to call modules in
           source_v143.sas for alternative v14.3 code.
*/

/*
Program Purpose: Store modules for initializing MCMC
*/

proc iml;

	**Loads variables and covariates from SAS datasets;
	start load_mcmc_variables(episodic_indicators,
														episodic_amounts,
														daily_amounts,
														covariates,
														intercepts,
														num_episodic,
														num_daily,
														num_variables,
														num_covariates,
														num_coefficients,
														has_never_consumers,
														never_consumer_covariates,
														never_consumer_intercept,
														num_never_consumer_covariates,
														num_never_consumer_coefficients,
														variable_data,
														covariate_data,
														never_consumer_covariate_data);
														
		use (variable_data);
		
			read all var {num_episodic} into num_episodic;
			read all var {num_daily} into num_daily;
			read all var {num_variables} into num_variables;
			
			if num_episodic > 0 then do;
			
				read all var (cats("episodic_ind", 1:num_episodic)) into episodic_indicators;
				read all var (cats("episodic_amt", 1:num_episodic)) into episodic_amounts;
			end;
			
			if num_daily > 0 then do;
			
				read all var (cats("daily_amt", 1:num_daily)) into daily_amounts;
			end;
			
			read all var {has_never_consumers} into has_never_consumers;
		close (variable_data);
		
		use (covariate_data);
		
			read all var {num_covariates} into num_covariates;
			read all var {intercept} into intercepts;
			read all var {num_coefficients} into num_coefficients;
			
			if max(num_covariates) > 0 then do;
			
				read all var (cats("covariate", 1:max(num_covariates))) into covariates;
			end;
		close (covariate_data);
		
		if has_never_consumers = 1 then do;
		
			use (never_consumer_covariate_data);
			
				read all var {num_covariates} into num_never_consumer_covariates;
				read all var {intercept} into never_consumer_intercept;
				read all var {num_coefficients} into num_never_consumer_coefficients;
				
				if num_never_consumer_covariates > 0 then do;
				
					read all var (cats("covariate", 1:num_never_consumer_covariates)) into never_consumer_covariates;
				end;
			close (never_consumer_covariate_data);
		end;
	finish;

	**Splits input dataset by repeated observations and extracts subject, observation, and weight data;
	start initialize_subject_data(input_data,
																id,
																repeat_obs,
																weight);
		if symgetn("DEBUG") then print "begin module: initialize_subject_data";			

		rc = checkFile(input_data, "initialize_subject_data module");
		use (input_data);
		
			read all var id into record_id;
			read point(uniqueby(record_id)) var id into subject_id;
		close (input_data);
		
		use(input_data);
		
			read all var repeat_obs into recall;
		close (input_data);
		
		recall_levels = unique(recall)`;
		
		num_records = nrow(record_id);
		num_subjects = nrow(subject_id);
		num_recalls = nrow(recall_levels);
		
		recall_availability = j(num_subjects, num_recalls, .);
		do day_k = 1 to num_recalls;
		
			recall[loc(recall = recall_levels[day_k])] = day_k;
			recall_availability[,day_k] = element(subject_id, record_id[loc(recall = day_k)]);
		end;
		
		if weight ^= "" then do;
		
			use (input_data);
			
				read point(uniqueby(record_id)) var weight into subject_weighting;
			close (input_data);
		end;
		else do;
		
			subject_weighting = j(num_subjects, 1, 1);
		end;
		subject_weighting = subject_weighting/mean(subject_weighting);
		
		subject_data = [#"record_id" = record_id,
										#"subject_id" = subject_id,
										#"recall" = recall,
										#"recall_availability" = recall_availability,
										#"subject_weighting" = subject_weighting,
										#"num_records" = num_records,
										#"num_subjects" = num_subjects,
										#"num_recalls" = num_recalls];
		if symgetn("DEBUG") then print "end module: initialize_subject_data";								
		return subject_data;
	finish;
	
	**Initializes model variable matrices for the MCMC function;
	start initialize_variable_matrices(input_data,
																		 episodic_indicators,
																		 episodic_amounts,
																		 daily_amounts,
																		 recall,
																		 recall_availability,
																		 num_variables,
																		 num_episodic,
																		 num_daily,
																		 num_subjects,
																		 num_recalls);
																		 
		if symgetn("DEBUG") then print "begin module: initialize_variable_matrices";	
		
		**model variable matrices for all recalls;
		rc = checkFile(input_data, "initialize_variable_matrices module");
		
		variable_names = j(num_variables, 1, BlankStr(32));
		
		use (input_data);
		
			if num_episodic > 0 then do;
		
				read all var (episodic_indicators) into episodic_indicator_matrix;
				variable_names[do(1, 2*num_episodic - 1, 2)] = episodic_indicators;
				
				read all var (episodic_amounts) into episodic_amount_matrix;
				episodic_amount_matrix[loc(^episodic_indicator_matrix)] = 0;
				variable_names[do(2, 2*num_episodic, 2)] = episodic_amounts;
			end;
			
			if num_daily > 0 then do;
			
				read all var (daily_amounts) into daily_amount_matrix;
				variable_names[do(2*num_episodic + 1, 2*num_episodic + num_daily, 1)] = daily_amounts;
			end;
		close (input_data);
									
		**variable matrices for each recall;									 
		episodic_indicator_matrices = ListCreate(num_recalls);
		episodic_amount_matrices = ListCreate(num_recalls);
		daily_amount_matrices = ListCreate(num_recalls);
		
		do day_k = 1 to num_recalls;
		
			records_in_recall = (recall = day_k);
			subjects_in_recall = (recall_availability[,day_k]);
			
			if num_episodic > 0 then do;
			
				**episodic indicators;
				episodic_indicator_matrix_k = j(num_subjects, num_episodic, 0);
				episodic_indicator_matrix_k[loc(subjects_in_recall),] = episodic_indicator_matrix[loc(records_in_recall),];
				
				episodic_indicator_matrices$day_k = episodic_indicator_matrix_k;
				
				**episodic amounts;
				episodic_amount_matrix_k = j(num_subjects, num_episodic, 0);
				episodic_amount_matrix_k[loc(subjects_in_recall),] = episodic_amount_matrix[loc(records_in_recall),];
				
				episodic_amount_matrices$day_k = episodic_amount_matrix_k;
			end;
			
			if num_daily > 0 then do;
			
				**daily amounts;
				daily_amount_matrix_k = j(num_subjects, num_daily, 0);
				daily_amount_matrix_k[loc(subjects_in_recall),] = daily_amount_matrix[loc(records_in_recall),];
				
				daily_amount_matrices$day_k = daily_amount_matrix_k;
			end;
		end;
		
		variable_matrix_data = [#"episodic_indicator_matrices" = episodic_indicator_matrices,
												 		#"episodic_amount_matrices" = episodic_amount_matrices,
												 		#"daily_amount_matrices" = daily_amount_matrices,
												 		#"variable_names" = variable_names];
														
		if symgetn("DEBUG") then print "end module: initialize_variable_matrices";										 		
		return variable_matrix_data;
	finish;
		
	**Initializes covariate matrices for the MCMC function;
	start initialize_covariate_matrices(input_data,
																			covariates,
																			intercepts,
																			record_id,
																			recall,
																			recall_availability,
																			subject_weighting,
																			num_variables,
																			num_covariates,
																			num_coefficients,
																			num_records,
																			num_subjects,
																			num_recalls,
																			has_never_consumers,
																			never_consumer_covariates,
																			never_consumer_intercept,
																			num_never_consumer_covariates,
																			num_never_consumer_coefficients);	
																			
		if symgetn("DEBUG") then print "begin module: initialize_covariate_matrices";		
    alg = symgetn("v143code") + 1;																
		
		**covariate matrices for model variables;
		covariate_matrices = ListCreate(num_variables);
		coefficient_names = ListCreate(num_variables);
		
		use (input_data);
		
			do var_j = 1 to num_variables;
			
				num_covariates_j = num_covariates[var_j];
				num_coefficients_j = num_coefficients[var_j];
			
				if num_covariates_j > 0 then do;
			
					covariates_j = covariates[var_j, 1:num_covariates_j];
				end;
				
				intercept_j = intercepts[var_j];
				
				**covariate matrix for all recalls;
				covariate_matrix_j = j(num_records, num_coefficients_j, .);
				coefficients_j = j(num_coefficients_j, 1, BlankStr(32));
				
				if intercept_j = 1 then do;
				
					covariate_matrix_j[,1] = 1;
					coefficients_j[1] = "intercept";
				end;
				
				if num_covariates_j > 0 then do;
				
					read all var (covariates_j) into covariate_data_j;
					covariate_matrix_j[,intercept_j + (1:num_covariates_j)] = covariate_data_j;
					coefficients_j[intercept_j + (1:num_covariates_j)] = covariates_j;
				end;
		
				**covariate matrices for each recall;
				covariate_matrices$var_j = ListCreate(num_recalls);
			
				do day_k = 1 to num_recalls;
				
					records_in_recall = (recall = day_k);
					subjects_in_recall = (recall_availability[,day_k]);
			
					covariate_matrix_j_k = j(num_subjects, num_coefficients_j, 0);
					covariate_matrix_j_k[loc(subjects_in_recall),] = covariate_matrix_j[loc(records_in_recall),];
					
					covariate_matrices$var_j$day_k = covariate_matrix_j_k;
				end;
				
				coefficient_names$var_j = coefficients_j;
			end;
		close (input_data);
		
		**weighted covariate matrices and squared sum;
		wt_covariate_matrices = ListCreate(num_variables);
		wt_covariate_sq_sums = ListCreate(num_variables);
		
		do var_j = 1 to num_variables;
		
			num_coefficients_j = num_coefficients[var_j];
			
			wt_covariate_matrices$var_j = ListCreate(num_recalls);
			wt_covariate_sq_sums$var_j = j(num_coefficients_j, num_coefficients_j, 0);
			
			do day_k = 1 to num_recalls;
			
				wt_covariate_matrices$var_j$day_k = (subject_weighting # recall_availability[,day_k] # covariate_matrices$var_j$day_k)`;
/* BW CHANGE 2025-10-22 */				
				*wt_covariate_sq_sums$var_j = wt_covariate_sq_sums$var_j + wt_covariate_matrices$var_j$day_k * covariate_matrices$var_j$day_k;
				wt_covariate_sq_sums$var_j = wt_covariate_sq_sums$var_j + matMultAB(wt_covariate_matrices$var_j$day_k, covariate_matrices$var_j$day_k, alg);
		
			end;
		end;
		
		**never-consumer covariate matrix;
		if has_never_consumers = 1 then do;
		
			use (input_data);
		
				never_consumer_covariate_matrix = j(num_subjects, num_never_consumer_coefficients, .);
				never_consumer_coefficient_names = j(num_never_consumer_coefficients, 1, BlankStr(32));
			
				if never_consumer_intercept = 1 then do;
			
					never_consumer_covariate_matrix[,1] = 1;
					never_consumer_coefficient_names[1] = "intercept";
				end;
	
				if num_never_consumer_covariates > 0 then do;
			
					read point(uniqueby(record_id)) var (never_consumer_covariates) into never_consumer_covariate_data;
					never_consumer_covariate_matrix[,never_consumer_intercept + (1:num_never_consumer_covariates)] = never_consumer_covariate_data;
					never_consumer_coefficient_names[never_consumer_intercept + (1:num_never_consumer_covariates)] = never_consumer_covariates;
				end;
			close (input_data);
		end;
		else do;
		
			never_consumer_covariate_matrix = {};
			never_consumer_coefficient_names = {};
		end;
		
		**output covariate data;
		covariate_matrix_data = [#"covariate_matrices" = covariate_matrices,
														 #"coefficient_names" = coefficient_names,
														 #"wt_covariate_matrices" = wt_covariate_matrices,
														 #"wt_covariate_sq_sums" = wt_covariate_sq_sums,
														 #"never_consumer_covariate_matrix" = never_consumer_covariate_matrix,
														 #"never_consumer_coefficient_names" = never_consumer_coefficient_names];
														 
		if symgetn("DEBUG") then print "end module: initialize_covariate_matrices";												 
		return covariate_matrix_data;
	finish initialize_covariate_matrices;
	
	**Initializes MCMC priors and starting values - beta, sigma-e, r, theta, and sigma-u, W, U, and never-consumer parmaeters (alpha1 and consumer probabilities);
	**Priors and starting values are defined in Appendices A.2 and A.3 of Zhang, et al. (2011) and Section 2.3.3 of Bhadra, et al. (2020) for never-consumers;
	start initialize_priors(episodic_indicator_matrices,
													episodic_amount_matrices,
													daily_amount_matrices,
													covariate_matrices,
													sigma_u_prior_data,
													recall_availability,
													num_episodic,
													num_daily,
													num_variables,
													num_coefficients,
													num_subjects,
													num_recalls,
													has_never_consumers,
													never_consumer_covariate_matrix,
													num_never_consumer_coefficients);
													
		if symgetn("DEBUG") then print "begin module: initialize_priors";
		alg = symgetn("v143code") + 1;
		
		**initialize beta prior;
		beta_mean_prior = ListCreate(num_variables);
		beta_covariance_prior = ListCreate(num_variables);
		
		do var_j = 1 to num_variables;
	
			num_coefficients_j = num_coefficients[var_j];
			
			**beta mean prior (number of coefficients x 1);
			beta_mean_prior_j = j(num_coefficients_j, 1, 0);
			beta_mean_prior$var_j = beta_mean_prior_j;
			
			**beta covariance prior (number of coefficients x number of coefficients);
			beta_covariance_prior_j = i(num_coefficients_j) # 100;
			beta_covariance_prior$var_j = beta_covariance_prior_j;
		end;
		beta_prior = [#"mean" = beta_mean_prior,
									#"covariance" = beta_covariance_prior];
		
		**initialize r matrix prior (number of episodically consumed foods - 1 x 1) and theta matrix prior ((number of episodically consumed foods - 1)^2 x 1);
		if num_episodic >= 2 then do;
		
			r_matrix_prior = j(num_episodic - 1, 1, 0);
			theta_matrix_prior = j((num_episodic - 1)##2, 1, 0);
		end;
		else do;
		
			r_matrix_prior = {};
			theta_matrix_prior = {};
		end;
	
		**Initialize V matrix prior (number of episodic and daily variables x number of episodic and daily variables);
		v_matrix_prior = i(num_variables);	/* identity matrix */
		
		**Initialize sigma-e prior (number of episodic and daily variables x number of episodic and daily variables);
/* BW CHANGE 2025-10-14 */
		/*sigma_e_prior = v_matrix_prior * v_matrix_prior`;*/
		/* No need to multiply an identity matrix */
		sigma_e_prior = v_matrix_prior;

		**Initialize sigma-u prior (number of episodic and daily variables x number of episodic and daily variables);
		if sigma_u_prior_data ^= "" then do;
		    rc = checkFile(sigma_u_prior_data, "initialize_priors module");
			use (sigma_u_prior_data);
			
				read all var _NUM_ into sigma_u_prior;
			close (sigma_u_prior_data);
		end;
		else do;
			sigma_u_prior = j(num_variables, num_variables, 0.5) + i(num_variables) # 0.5;
		end;
 
		**Initialize U matrix prior (number of subjects x number of episodic and daily variables);
/* BW CHANGE 2025-10-14 */
		*call eigen(sigma_u_eigvals, sigma_u_eigvecs, sigma_u_prior);

		tmp             = eigenDecomp(sigma_u_prior, alg); 
		sigma_u_eigvals = tmp$"eigenvalues";
		sigma_u_eigvecs = tmp$"eigenvectors";
		free tmp;

/* BW CHANGE 2025-10-14 */
		**sqrt_sigma_u = sigma_u_eigvecs * diag(sqrt(sigma_u_eigvals)) * sigma_u_eigvecs`;
		sqrt_sigma_u = matMultADAt(sigma_u_eigvecs, sqrt(sigma_u_eigvals), alg);
				
		normals = j(num_subjects, num_variables, .);
		call randgen(normals, "Normal", 0, 1);
		
/* BW CHANGE 2025-10-14 */
		*u_matrix_prior = normals * sqrt_sigma_u;
		u_matrix_prior = matMultAB(normals, sqrt_sigma_u, alg);
		
		**Initialize XBeta prior (number of subjects x number of episodic and daily variables for each recall);
		xbeta_prior = ListCreate(num_recalls);
		
		do day_k = 1 to num_recalls;
		
			xbeta_prior_k = j(num_subjects, num_variables, .);
			
			do var_j = 1 to num_variables;
/* BW CHANGE 2025-10-14 */			
				*xbeta_prior_k[,var_j] = covariate_matrices$var_j$day_k * beta_mean_prior$var_j;
				xbeta_prior_k[,var_j] = matMultAB(covariate_matrices$var_j$day_k, beta_mean_prior$var_j, alg);
			end;
			
			xbeta_prior$day_k = xbeta_prior_k;
		end;
		
		**Initialize XBeta_U prior (number of subjects x number of episodic and daily variables for each recall);
		xbeta_u_prior = ListCreate(num_recalls);
		do day_k = 1 to num_recalls;
		
			xbeta_u_prior_k = xbeta_prior$day_k + u_matrix_prior;
			
			xbeta_u_prior$day_k = xbeta_u_prior_k;
		end;
		
		**Initialize W matrix prior (number of subjects x number of episodic and daily variables for each recall);
		w_matrix_prior = ListCreate(num_recalls);
		
		do day_k = 1 to num_recalls;
		
			w_matrix_prior_k = j(num_subjects, num_variables, .);
			
			do episodic_j = 1 to num_episodic;
			
				normals = j(num_subjects, 1, .);
				call randgen(normals, "Normal", 0, 1);
			
				w_matrix_prior_k[,2*episodic_j - 1] = (2 # episodic_indicator_matrices$day_k[,episodic_j] - 1) #
													 										abs(xbeta_u_prior$day_k[,2*episodic_j - 1] + normals) #
													 										recall_availability[,day_k];					 
				w_matrix_prior_k[,2*episodic_j] = episodic_amount_matrices$day_k[,episodic_j];
			end;
			
			do daily_j = 1 to num_daily;
			
				w_matrix_prior_k[,2*num_episodic + daily_j] = daily_amount_matrices$day_k[,daily_j];
			end;
			
			w_matrix_prior$day_k = w_matrix_prior_k;
		end;
		
		w_matrix_prior = update_w_matrix(w_matrix_prior,
																		 xbeta_u_prior,
																		 sigma_e_prior,
																		 recall_availability,
																		 episodic_indicator_matrices,
																		 num_subjects,
																		 num_episodic,
																		 num_variables,
																		 num_recalls);
		
		**never-consumers priors;
		if has_never_consumers = 1 then do;
			
			**Initialize alpha1 mean (number of never-consumer coefficents x 1) and covariance (number of never-consumer coefficients x number of never-consumer coefficients) priors;
			alpha1_mean_prior = j(num_never_consumer_coefficients, 1, 0);
			
			alpha1_covariance_prior = i(num_never_consumer_coefficients);
			
			**Initialize consumer probabilities prior (number of subjects x 1);
/* BW CHANGE 2025-10-22 */
			*g_alpha_prior = never_consumer_covariate_matrix * alpha1_mean_prior;
			g_alpha_prior = matMultAB(never_consumer_covariate_matrix, alpha1_mean_prior, alg);
			consumer_probabilities_prior = cdf("Normal", g_alpha_prior);
		end;
		else do;
		
			alpha1_mean_prior = {};
			alpha1_covariance_prior = {};
			consumer_probabilities_prior = {};
		end;
		alpha1_prior = [#"mean" = alpha1_mean_prior,
										#"covariance" = alpha1_covariance_prior];
										
		mcmc_priors = [#"alpha1_prior" = alpha1_prior,
									 #"consumer_probabilities_prior" = consumer_probabilities_prior,
									 #"beta_prior" = beta_prior,
									 #"sigma_e_prior" = sigma_e_prior,
									 #"r_matrix_prior" = r_matrix_prior,
									 #"theta_matrix_prior" = theta_matrix_prior,
									 #"v_matrix_prior" = v_matrix_prior,
									 #"sigma_u_prior" = sigma_u_prior,
									 #"u_matrix_prior" = u_matrix_prior,
									 #"xbeta_prior" = xbeta_prior,
									 #"xbeta_u_prior" = xbeta_u_prior,
									 #"w_matrix_prior" = w_matrix_prior];
									 
		if symgetn("DEBUG") then print "end module: initialize_priors";							 
		return mcmc_priors;
	finish;
	
	start initialize_mcmc(input_data,
												id,
												repeat_obs,
												weight,
												episodic_indicators,
												episodic_amounts,
												daily_amounts,
												covariates,
												intercepts,
												sigma_u_prior_data,
												num_episodic,
												num_daily,
												num_variables,
												num_covariates,
												num_coefficients,
												has_never_consumers,
												never_consumer_covariates,
												never_consumer_intercept,
												num_never_consumer_covariates,
												num_never_consumer_coefficients,
												mcmc_seed=0);
			
		if symgetn("DEBUG") then print "begin module: initialize_mcmc";
		
		**1. Set seed for MCMC;
		if mcmc_seed = 0 then do;
		
			mcmc_seed = j(1, 1, .);
			call randgen(mcmc_seed, "Uniform", 0, 1);
			mcmc_seed = int(mcmc_seed * 20000000 + 1);
		end;
		call randseed(mcmc_seed);
		
		**2. Define subject-level data;
		subject_data = initialize_subject_data(input_data,
																					 id,
																					 repeat_obs,
																					 weight);
																 
		**3. Define variable matrices;		
		variable_matrix_data = initialize_variable_matrices(input_data,
																										 		episodic_indicators,
																										 		episodic_amounts,
																										 		daily_amounts,
																										 		subject_data$"recall",
																										 		subject_data$"recall_availability",
																										 		num_variables,
																										 		num_episodic,
																										 		num_daily,
																										 		subject_data$"num_subjects",
																										 		subject_data$"num_recalls");
																									 	 
		**4. Define covariate matrices;
		covariate_matrix_data = initialize_covariate_matrices(input_data,
																													covariates,
																													intercepts,
																										 	 		subject_data$"record_id",
																										 	 		subject_data$"recall",
																										 	 		subject_data$"recall_availability",
																										 	 		subject_data$"subject_weighting",
																										 	 		num_variables,
																										 	 		num_covariates,
																										 	 		num_coefficients,
																										 	 		subject_data$"num_records",
																										 	 		subject_data$"num_subjects",
																										 	 		subject_data$"num_recalls",
																										 	 		has_never_consumers,
																										 	 		never_consumer_covariates,
																										 	 		never_consumer_intercept,
																										 	 		num_never_consumer_covariates,
																										 	 		num_never_consumer_coefficients);
																										 	 
		**5. Initialize MCMC priors;
		mcmc_priors = initialize_priors(variable_matrix_data$"episodic_indicator_matrices",
																		variable_matrix_data$"episodic_amount_matrices",
																		variable_matrix_data$"daily_amount_matrices",
																		covariate_matrix_data$"covariate_matrices",
																		sigma_u_prior_data,
																		subject_data$"recall_availability",
																		num_episodic,
																		num_daily,
																		num_variables,
																		num_coefficients,
																		subject_data$"num_subjects",
																		subject_data$"num_recalls",
																		has_never_consumers,
																		covariate_matrix_data$"never_consumer_covariate_matrix",
																		num_never_consumer_coefficients);
																		
		mcmc_parameters = [#"subject_data" = subject_data,
											 #"variable_matrix_data" = variable_matrix_data,
											 #"covariate_matrix_data" = covariate_matrix_data,
											 #"priors" = mcmc_priors,
											 #"mcmc_seed" = mcmc_seed];
											 
		if symgetn("DEBUG") then print "end module: initialize_mcmc";									 
		return mcmc_parameters;
	finish;
	
	reset storage=WORK.mcmc_modules;
	
	store module=_all_;
quit;