/* 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:
IML modules for initializing distrib macro
*/

proc iml;

	**Load variables, covariates, iterations, and subjects from MCMC model;
	start load_distrib_variables(episodic_variables,
															 daily_variables,
															 covariates,
															 intercepts,
															 beta_mean,
															 sigma_u_mean,
															 sigma_e_mean,
															 mcmc_subjects,
															 num_episodic,
															 num_daily,
															 num_variables,
															 num_covariates,
															 num_coefficients,
															 num_mcmc_subjects,
															 num_post,
															 has_never_consumers,
															 never_consumer_covariates,
															 never_consumer_intercept,
															 alpha1_mean,
															 num_never_consumer_covariates,
															 num_never_consumer_coefficients,
															 multivar_mcmc_model);
	
		**variables;
		use (cats(multivar_mcmc_model, "_vars"));
		
			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_var", 1:num_episodic)) into episodic_variables;
			end;
			
			if num_daily > 0 then do;
			
				read all var (cats("daily_var", 1:num_daily)) into daily_variables;
			end;
			
			read all var {has_never_consumers} into has_never_consumers;
		close (cats(multivar_mcmc_model, "_vars"));		
		
		**covariates;
		use (cats(multivar_mcmc_model, "_covars"));
		
			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 (cats(multivar_mcmc_model, "_covars"));
		
		if has_never_consumers = 1 then do;
		
			use (cats(multivar_mcmc_model, "_covarsnc"));
			
				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 (cats(multivar_mcmc_model, "_covarsnc"));
		end;
		
		**posterior means;
		beta_mean = ListCreate(num_variables);
		
		do var_j = 1 to num_variables;
		
			use (cats(multivar_mcmc_model, "_beta", var_j, "_mean"));
			
				read all var _NUM_ into beta_j_mean;
				beta_mean$var_j = beta_j_mean`;
			close (cats(multivar_mcmc_model, "_beta", var_j, "_mean"));
		end;
		
		use (cats(multivar_mcmc_model, "_sigma_u_mean"));
		
			read all var _NUM_ into sigma_u_mean;
		close (cats(multivar_mcmc_model, "_sigma_u_mean"));
		
		use (cats(multivar_mcmc_model, "_sigma_e_mean"));
		
			read all var _NUM_ into sigma_e_mean;
		close (cats(multivar_mcmc_model, "_sigma_e_mean"));
		
		if has_never_consumers = 1 then do;
		
			use (cats(multivar_mcmc_model, "_alpha1_mean"));
				
				read all var _NUM_ into alpha1_mean;
				alpha1_mean = alpha1_mean`;
			close (cats(multivar_mcmc_model, "_alpha1_mean"));
		end;
		
		**MCMC subjects;
		use (cats(multivar_mcmc_model, "_subjects")) nobs num_mcmc_subjects;
		
			read all var {subject_id} into mcmc_subjects;
		close (cats(multivar_mcmc_model, "_subjects"));
		
		**Post-MCMC iterations;
		use (cats(multivar_mcmc_model, "_iters"));
		
			read all var {num_post} into num_post;
		close (cats(multivar_mcmc_model, "_iters"));								 
	finish;

	**Calculate population variables (sample sizes, subject IDs, dietary supplements);
	start calculate_population_variables(distrib_population,
																			 id,
																			 nuisance_weight,
																			 dietary_supplements);
																			 
		if symgetn("DEBUG") then print "begin module: calculate_population_variables";																	 
		**Extract subjects;
		rc = checkFile(distrib_population, "calculate_population_variables module");
		use (distrib_population);
			read all var (id) into record_id;
		close (distrib_population);
		
		**calculate number of subjects and nuisance variable levels per subject;
		call tabulate(subject_id, num_nuisance_levels, record_id);
		subject_id = subject_id`;
		
		num_records = nrow(record_id);
		num_subjects = nrow(subject_id);
		
		**extract weights for nuisance variable levels and normalize sum of weights per subject to 1;
		if nuisance_weight ^= "" then do;
		
			use (distrib_population);
			
				read all var (nuisance_weight) into nuisance_level_weights;
			close (distrib_population);
		end;
		else do;
		
			nuisance_level_weights = j(num_records, 1, 1);
		end;
		
		do i = 1 to num_subjects;
		
			subject_index = loc(record_id = subject_id[i]);
		
			nuisance_level_weights[subject_index] = nuisance_level_weights[subject_index]/sum(nuisance_level_weights[subject_index]);
		end;
		
		**extract dietary supplement data;
		if dietary_supplements ^= "" then do;
		
			rc = checkFile(dietary_supplements, "calculate_population_variables module");
			use (dietary_supplements);
				read all var {variable} into variables_to_supplement;
				variables_to_supplement = variables_to_supplement`;
				
				read all var {supplement} into supplements;
			close (dietary_supplements);
			
			use (distrib_population);
				read all var (supplements) into dietary_supplement_data;
			close (distrib_population);
		end;
		else do;
		
			variables_to_supplement = {};
			dietary_supplement_data = {};
		end;
		
		**return population variables in list;
		population_variables = [#"subject_id" = subject_id,
														#"nuisance_level_weights" = nuisance_level_weights,
														#"num_records" = num_records,
														#"num_subjects" = num_subjects,
														#"num_nuisance_levels" = num_nuisance_levels,
														#"dietary_supplement_data" = dietary_supplement_data,
														#"variables_to_supplement" = variables_to_supplement];

		if symgetn("DEBUG") then print "end module: calculate_population_variables";
		return population_variables;
	finish;
	
	**Extract covariate data for each variable from input dataset;
	start extract_covariate_matrices(distrib_population,
																	 covariates,
																	 intercepts,
																	 num_variables,
																	 num_covariates,
																	 num_coefficients,
																	 num_records,
																	 has_never_consumers,
																	 never_consumer_covariates,
																	 never_consumer_intercept,
																	 num_never_consumer_covariates,
																	 num_never_consumer_coefficients);
																	 
		if symgetn("DEBUG") then print "begin module: extract_covariate_matrices";
																	 
		**Extract covariate matrices for each model component;
		covariate_matrices = ListCreate(num_variables);
		
		rc = checkFile(distrib_population, "extract_covariate_matrices module");
		use (distrib_population);
		
			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_j = j(num_records, num_coefficients_j, .);
			
				if intercept_j = 1 then do;
				
					covariate_matrix_j[,1] = 1;
				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;
				end;
				
				covariate_matrices$var_j = covariate_matrix_j;
			end;
		close (distrib_population);
		
		**extract never-consumer covariate matrix;		
		if has_never_consumers = 1 then do;
		
			use (distrib_population);
			
				never_consumer_covariate_matrix = j(num_records, num_never_consumer_coefficients, .);
		
				if never_consumer_intercept = 1 then do;
				
					never_consumer_covariate_matrix[,1] = 1;
				end;
			
				if num_never_consumer_covariates > 0 then do;
				
					read all 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;
				end;
			close (distrib_population);
		end;
		else do;
		
			never_consumer_covariate_matrix = {};
		end;
		
		covariate_matrix_data = [#"covariate_matrices" = covariate_matrices,
														 #"never_consumer_covariate_matrix" = never_consumer_covariate_matrix];

		if symgetn("DEBUG") then print "end module: extract_covariate_matrices";													
		return covariate_matrix_data;
	finish;
	
	**Calculates XBeta, U standard deviation (square root of U), sample sizes, and consumer probabilities;
	start initialize_distrib_parameters(covariate_matrices,
																			beta_mean,
																			sigma_u_mean,
																			sigma_e_mean,
																			num_variables,
																			num_records,
																			has_never_consumers,
																			never_consumer_covariate_matrix,
																			alpha1_mean);

		if symgetn("DEBUG") then print "begin module: initialize_distrib_parameters";		
		alg = symgetn("v143code") + 1;

		**calculate XBeta;
		xbeta = j(num_records, num_variables, .);

		do var_j = 1 to num_variables;
/* BW CHANGE 2025-10-20 */		
			*xbeta[,var_j] = covariate_matrices$var_j * beta_mean$var_j;
			xbeta[,var_j] = matMultAB(covariate_matrices$var_j, beta_mean$var_j, alg);
		end;

/* BW CHANGE 2025-10-20 */
		**calculate the standard deviation of U;
		*call eigen(sigma_u_eigvals, sigma_u_eigvecs, sigma_u_mean);
		tmp = eigenDecomp(sigma_u_mean, alg);
		sigma_u_eigvals = tmp$"eigenvalues";
		sigma_u_eigvecs = tmp$"eigenvectors";
		free tmp;
		
		*u_standard_deviation = sigma_u_eigvecs * diag(sqrt(sigma_u_eigvals)) * sigma_u_eigvecs`;
		u_standard_deviation = matMultADAt(sigma_u_eigvecs, sqrt(sigma_u_eigvals), alg);
	
		**if never-consumers are allowed, calculate consumer probabilities;
		if has_never_consumers = 1 then do;
			*g_alpha = never_consumer_covariate_matrix * alpha1_mean;
			g_alpha = matMultAB(never_consumer_covariate_matrix, alpha1_mean, alg);
			
			consumer_probabilities = cdf("Normal", g_alpha);
		end;
		else do;
		
			consumer_probabilities = {};
		end;
		
		**output list of distrib parameters;
		distrib_parameters = [#"xbeta" = xbeta,
													#"u_standard_deviation" = u_standard_deviation,
													#"sigma_e_mean" = sigma_e_mean,
													#"consumer_probabilities" = consumer_probabilities];

		if symgetn("DEBUG") then print "end module: initialize_distrib_parameters";												
		return distrib_parameters;
	finish;
	
	**Initializes parameters for multivar distrib main loop;
	start initialize_distrib(distrib_population,
													 id,
													 nuisance_weight,
													 covariates,
													 intercepts,
													 beta_mean,
													 sigma_u_mean,
													 sigma_e_mean,
													 num_variables,
													 num_covariates,
													 num_coefficients,
													 dietary_supplements,
													 has_never_consumers,
													 never_consumer_covariates,
													 never_consumer_intercept,
													 alpha1_mean,
													 num_never_consumer_covariates,
													 num_never_consumer_coefficients,
													 distrib_seed=0);
													 
		if symgetn("DEBUG") then print "begin module: initialize_distrib";		
													 
		**1. Set seed;
		if distrib_seed = 0 then do;
		
			distrib_seed = j(1, 1, .);
			call randgen(distrib_seed, "Uniform", 0, 1);
			distrib_seed = int(distrib_seed * 20000000 + 1);
		end;
		call randseed(distrib_seed);
		
		**2. Calculate population variables (sample sizes, number of covariate lists, number of variables);
		population_variables = calculate_population_variables(distrib_population,
																			 										id,
																			 										nuisance_weight,
																			 										dietary_supplements);
																													
		**3. Extract covariate matrices;
		covariate_matrix_data = extract_covariate_matrices(distrib_population,
																	 										 covariates,
																	 										 intercepts,
																	 										 num_variables,
																	 										 num_covariates,
																	 										 num_coefficients,
																	 										 population_variables$"num_records",
																	 										 has_never_consumers,
																	 										 never_consumer_covariates,
																	 										 never_consumer_intercept,
																	 										 num_never_consumer_covariates,
																	 										 num_never_consumer_coefficients);
																								
		**4. Initialize distrib parameters;
		distrib_parameters = initialize_distrib_parameters(covariate_matrix_data$"covariate_matrices",
																											 beta_mean,
																											 sigma_u_mean,
																											 sigma_e_mean,
																											 num_variables,
																											 population_variables$"num_records",
																											 has_never_consumers,
																											 covariate_matrix_data$"never_consumer_covariate_matrix",
																											 alpha1_mean);
																											 
		**5. Output distrib parameters;
		output_parameters = [#"xbeta" = distrib_parameters$"xbeta",
												 #"u_standard_deviation" = distrib_parameters$"u_standard_deviation",
												 #"sigma_e_mean" = distrib_parameters$"sigma_e_mean",
												 #"consumer_probabilities" = distrib_parameters$"consumer_probabilities",
												 #"subject_id" = population_variables$"subject_id",
												 #"nuisance_level_weights" = population_variables$"nuisance_level_weights",
												 #"num_records" = population_variables$"num_records",
												 #"num_subjects" = population_variables$"num_subjects",
												 #"num_nuisance_levels" = population_variables$"num_nuisance_levels",
												 #"dietary_supplement_data" = population_variables$"dietary_supplement_data",
												 #"variables_to_supplement" = population_variables$"variables_to_supplement",
												 #"distrib_seed" = distrib_seed];

		if symgetn("DEBUG") then print "end module: initialize_distrib";							 
		return output_parameters;
	finish;
	
	reset storage=WORK.distrib_modules;
	
	store module=_all_;
quit;