/* 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 functions called during distrib main loop;
*/

proc iml;

	**Simulates random intercepts ('u' matrix) for each subject;
	start generate_u_matrix(use_mcmc_u_matrices,
													mcmc_u_matrix,
													u_standard_deviation,
													distrib_subjects,
													mcmc_subjects,
													num_subjects,
													num_variables);
					
		if symgetn("DEBUG") then print "begin module: generate_u_matrix";	
		alg = symgetn("v143code") + 1;
		
		**total number of episodic and daily variables;
		if upcase(use_mcmc_u_matrices) = "Y" then do;
		
			u_matrix = j(num_subjects, num_variables, .);
			
			num_stored = ncol(loc(element(distrib_subjects, mcmc_subjects)));
			num_sim = num_subjects - num_stored;
			
			**use the U matrix for the current replicate as a base;
			if num_stored > 0 then do;
				
				u_matrix[loc(element(distrib_subjects, mcmc_subjects)),] = mcmc_u_matrix[loc(element(mcmc_subjects, distrib_subjects)),];
			end;
			
			**simulate remaining observations if needed;
			if num_sim > 0 then do;
			
				normals = j(num_sim, num_variables, .);
				call randgen(normals, "Normal");
/* BW CHANGE 2025-10-20 */
				u_matrix[loc(^element(distrib_subjects, mcmc_subjects)),] = matMultAB(normals, u_standard_deviation, alg);
				*u_matrix[loc(^element(distrib_subjects, mcmc_subjects)),] = normals * u_standard_deviation;

			end;
		end;
		else do;
		
			**simulate all observations if not using MCMC U matrices;
			normals = j(num_subjects, num_variables, .);
			call randgen(normals, "Normal");
/* BW CHANGE 2025-10-20 */
			u_matrix = matMultAB(normals, u_standard_deviation, alg);
			*u_matrix = normals * u_standard_deviation;
		end;
		
		if symgetn("DEBUG") then print "end module: generate_u_matrix";
		return u_matrix;
	finish;
	
	**Calculates the probability of consumption for episodic variables;
	start calculate_probability(xbeta_u,
															episodic_variables,
															num_records,
															num_episodic,
															has_never_consumers,
															never_consumers);
															
		if symgetn("DEBUG") then print "begin module: calculate_probability";

		if num_episodic = 0 then do;
		
			consumption_probability = {};
			return consumption_probability;
		end;
		
		**Calculate consumption probability of episodic variables;
		consumption_probability = j(num_records, num_episodic, .);
		
		do i = 1 to num_episodic;
		
			consumption_probability[,i] = cdf("Normal", xbeta_u[,2*i - 1]);
		end;
		
		**if never-consumers are allowed, set consumption probabilities for never-consumers to zero for the first episodic food;
		if has_never_consumers = 1 then do;
		
			if any(never_consumers) = 1 then do;
			
				consumption_probability[loc(never_consumers),1] = 0;
			end;
		end;
		
		if symgetn("DEBUG") then print "end module: calculate_probability";
		return consumption_probability;
	finish;
	
	**Calculates the consumption amount on the original scale for episodic and daily variables;
	start calculate_amount(xbeta_u,
												 sigma_e_mean,
												 backtransformation_data,
												 episodic_variables,
												 daily_variables,
												 num_records,
												 num_episodic,
												 num_daily,
												 has_never_consumers,
												 never_consumers);
											
		if symgetn("DEBUG") then print "begin module: calculate_amount";
		variables = episodic_variables || daily_variables;
		
		**backtransformation parameters;
		rc = checkFile(backtransformation_data, "calculate_amount module");
		use (backtransformation_data);
			read all var {variable} into tran_variable;
			read all var {tran_lambda} into tran_lambda;
			read all var {minamount} into minamount;
			read all var {tran_center} into tran_center;
			read all var {tran_scale} into tran_scale;
			read all var {biomarker} into biomarker;
		close (backtransformation_data);
		
		**calculate consumption amount on the original scale for all variables;
		backtransformed_amount = j(num_records, num_episodic + num_daily, .);
		
		do i = 1 to num_episodic + num_daily;
		
			backtran_index = loc(upcase(tran_variable) = upcase(variables[i]));
			
			if i <= num_episodic then do;
				amount_index = 2*i;
			end;
			else do;
				amount_index = num_episodic + i;
			end;
			
			backtransformed_amount[,i] = backtransform(xbeta_u[,amount_index],
																								 sigma_e_mean[amount_index, amount_index],
																								 tran_lambda[backtran_index],
																								 tran_center[backtran_index],
																								 tran_scale[backtran_index],
																								 minamount[backtran_index],
																								 biomarker[backtran_index]);
		end;
		
		**if never-consumers are allowed, set amounts for never-consumers to missing for the first episodic food;
		if has_never_consumers = 1 & num_episodic > 0 then do;
		
			if any(never_consumers) = 1 then do;
			
				backtransformed_amount[loc(never_consumers), 1] = .;
			end;
		end;
		
		if symgetn("DEBUG") then print "end module: calculate_amount";
		return backtransformed_amount;
	finish;
	
	start calculate_usual_intake(consumption_probability,
															 backtransformed_amount,
															 episodic_variables,
															 daily_variables,
															 num_records,
															 num_episodic,
															 num_daily,
															 has_never_consumers,
															 never_consumers);
		
		if symgetn("DEBUG") then print "begin module: calculate_usual_intake";
		**calculate usual intake for episodic and daily variables;
		usual_intake = j(num_records, num_episodic + num_daily, .);
		
		if num_episodic > 0 then do;
			
			usual_intake[,1:num_episodic] = consumption_probability # backtransformed_amount[,1:num_episodic];
		end;
		
		if num_daily > 0 then do;
			
			usual_intake[,(num_episodic+1):(num_episodic+num_daily)] = backtransformed_amount[,(num_episodic+1):(num_episodic+num_daily)];
		end;
		
		**if never-consumers are allowed, set usual intake for never-consumers to zero;
		if has_never_consumers = 1 & num_episodic > 0 then do;
		
			if any(never_consumers) = 1 then do;
			
				usual_intake[loc(never_consumers), 1] = 0;
			end;
		end;
		
		if symgetn("DEBUG") then print "end module: calculate_usual_intake";
		return usual_intake;
	finish;
	
	start calculate_supplemented_intake(usual_intake,
																			dietary_supplement_data,
																			variables_to_supplement,
																			episodic_variables,
																			daily_variables,
																			num_records);
																			
		if symgetn("DEBUG") then print "begin module: calculate_supplemented_intake";
																			
		if ncol(variables_to_supplement) = 0 then do;
		
			supplemented_intake = {};
			return supplemented_intake;
		end;
												
		variables = episodic_variables || daily_variables;
		
		num_supplemented = ncol(dietary_supplement_data);
		supplemented_intake = j(num_records, num_supplemented, .);
		
		do i = 1 to num_supplemented;
			
			variable_index = loc(variables = variables_to_supplement[i]);
			supplemented_intake[,i] = usual_intake[,variable_index] + dietary_supplement_data[,i];
		end;
		
		if symgetn("DEBUG") then print "end module: calculate_supplemented_intake";
		return supplemented_intake;
	finish;
	
	**Calculates consumption amounts backtransformed to the original scale from a mixed model predictor of a Box-Cox transformed variable;
	start backtransform(xbeta_u_var,
											sigma_e_var,
											lambda,
											center,
											scale,
											minimum_amount,
											is_biomarker);
									
		if symgetn("DEBUG") then print "begin module: backtransform";							
		if is_biomarker = 1 then do;
		
			boxcox_amounts = center + scale # xbeta_u_var;
			backtransformed_amounts = inverse_box_cox(boxcox_amounts, lambda, minimum_amount);
		end;
		else do;
		
			if lambda = 0 then do;
			
				boxcox_amounts = center + scale # xbeta_u_var + (scale**2)*sigma_e_var/2;
				backtransformed_amounts = inverse_box_cox(boxcox_amounts, 0, minimum_amount);
			end;
			else do;
			
				**perform 9-point approximation for non-zero lambda;
				x9pt = {-2.1, -1.3, -0.8, -0.5, 0, 0.5, 0.8, 1.3, 2.1};
				w9pt = {0.063345, 0.080255, 0.070458, 0.159698, 0.252489, 0.159698, 0.070458, 0.080255, 0.063345};
				w9pt = w9pt/sum(w9pt);
				
				backtransformed_amounts = j(nrow(xbeta_u_var), 1, 0);
				do i = 1 to 9;
				
					boxcox_amounts = center + scale # (xbeta_u_var + x9pt[i]*sqrt(sigma_e_var));
					backtransformed_amounts = backtransformed_amounts + w9pt[i] # inverse_box_cox(boxcox_amounts, lambda, minimum_amount);
				end;
			end;
		end;
		
		if symgetn("DEBUG") then print "end module: backtransform";	
		return backtransformed_amounts;
	finish;
	
	**Converts Box-Cox transformed values back to the original scale;
	start inverse_box_cox(boxcox_values,
												lambda,
												minimum_amount);
				
		if symgetn("DEBUG") then print "begin module: inverse_box_cox";			
		if lambda = 0 then do;
		
			inverse_boxcox = exp(boxcox_values);
		end;
		else do;
		
			inverse_boxcox = ((1 + lambda # boxcox_values) <> 0) ## (1/lambda);
		end;
		
		inverse_boxcox = inverse_boxcox <> minimum_amount;
		
		if symgetn("DEBUG") then print "end module: inverse_box_cox";
		return inverse_boxcox;
	finish;
	
	reset storage=WORK.distrib_modules;
	
	store module=_all_;
quit;