/* 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 distrib main loop
*/

proc iml;

	**Produces simulated Monte Carlo datasets using MCMC multivar parameters;
	start distrib_main_loop(xbeta,
													u_standard_deviation,
													sigma_e_mean,
													backtransformation_data,
													nuisance_level_weights,
													use_mcmc_u_matrices,
													mcmc_u_matrices,
													mcmc_subjects,
													distrib_subjects,
													episodic_variables,
													daily_variables,
													num_episodic,
													num_daily,
													num_variables,
													num_records,
													num_subjects,
													num_nuisance_levels,
													num_mcmc_subjects,
													num_post,
													num_simulated_u,
													dietary_supplement_data,
													variables_to_supplement,
													has_never_consumers,
													consumer_probabilities,
													low_memory_mode);
													
		if symgetn("DEBUG") then print "begin module: distrib_main_loop";		
													
		if upcase(use_mcmc_u_matrices) = "Y" then do;
		
			num_replicates = num_post;
			
			use (mcmc_u_matrices);
			u_matrix_names = cats("u_col", 1:num_variables);
		end;
		else num_replicates = num_simulated_u;
		
		**create distrib column names;
		distrib_names = cats("usl_", (episodic_variables || daily_variables));
		if ncol(variables_to_supplement) > 0 then distrib_names = distrib_names || cats("sup_", variables_to_supplement);
		if num_episodic > 0 then distrib_names = distrib_names || cats("prb_", episodic_variables) || cats("amt_", episodic_variables);
		
		num_distrib_cols = ncol(distrib_names);
													
		**main loop of replicates;
		if upcase(low_memory_mode) = "Y" then do;
		
			create _distrib_replicates var (distrib_names);
		end;
		else do;
		
			distrib_replicates = j(num_subjects * num_replicates, num_distrib_cols, .);
		end;
		
		do replicate = 1 to num_replicates;
		
			**generate U matrix from MCMC output and simulated data (if needed);
			if upcase(use_mcmc_u_matrices) = "Y" then do;
			
				setin (mcmc_u_matrices);
				read next (num_subjects) var (u_matrix_names) into mcmc_u_matrix;
			end;
			
			u_matrix = generate_u_matrix(use_mcmc_u_matrices,
																	 mcmc_u_matrix,
																	 u_standard_deviation,
																	 distrib_subjects,
																	 mcmc_subjects,
																	 num_subjects,
																	 num_variables);
																	 
			**calculate XBeta + U;
			xbeta_u = xbeta + u_matrix[repeat(1:num_subjects, num_nuisance_levels),];
			
			**find subjects that are never-consumers of the first episodic food;
			if has_never_consumers = 1 then do;
		
				random_uniform = j(num_subjects, 1, .);
				call randgen(random_uniform, "Uniform");
				
				never_consumers = (random_uniform[repeat(1:num_subjects, num_nuisance_levels)] > consumer_probabilities);
			end;
																								 
			**calculate consumption probabilities of each episodic food for each subject;
			consumption_probability = calculate_probability(xbeta_u,
																											episodic_variables,
																											num_records,
																											num_episodic,
																											has_never_consumers,
																											never_consumers);
			
			**calculate backtransformed amount of each episodic food for each subject;
			backtransformed_amount = calculate_amount(xbeta_u,
													 											sigma_e_mean,
													 											backtransformation_data,
													 											episodic_variables,
													 											daily_variables,
													 											num_records,
													 											num_episodic,
													 											num_daily,
													 											has_never_consumers,
													 											never_consumers);
			
			**calculate usual intake (t values);
			usual_intake = calculate_usual_intake(consumption_probability,
																 						backtransformed_amount,
																 						episodic_variables,
																 						daily_variables,
																 						num_records,
																 						num_episodic,
																 						num_daily,
																 						has_never_consumers,
																 						never_consumers);
			
			**calculate supplemented intake;
			supplemented_intake = calculate_supplemented_intake(usual_intake,
																													dietary_supplement_data,
																													variables_to_supplement,
																													episodic_variables,
																													daily_variables,
																													num_records);
			
			**combine intakes, probabilities, and amounts into dataset;
			replicate_matrix = usual_intake;
			if ncol(variables_to_supplement) > 0 then replicate_matrix = replicate_matrix || supplemented_intake;
			if num_episodic > 0 then replicate_matrix = replicate_matrix || consumption_probability || backtransformed_amount[,1:num_episodic];
			
			**sum over nuisance variable levels;
			replicate_sums = j(num_subjects, num_distrib_cols, 0);
			
			start_index = 0;
			
			do i = 1 to num_subjects;
			
				index = start_index + (1:num_nuisance_levels[i]);
				
				wt_replicate_i = nuisance_level_weights[index] # replicate_matrix[index,];
				replicate_sums[i,] = wt_replicate_i[+,];
				
				start_index = start_index + num_nuisance_levels[i];
			end;
			
			**save distrib replicate;
			if upcase(low_memory_mode) = "Y" then do;
			
				setout _distrib_replicates;
				append from replicate_sums;
			end;
			else do;
			
				replicate_index = (replicate - 1)*num_subjects + (1:num_subjects);
				distrib_replicates[replicate_index,] = replicate_sums;
			end;
		end;
		
		if upcase(use_mcmc_u_matrices) = "Y" then do;
		
			close (mcmc_u_matrices);
		end;
		
		if upcase(low_memory_mode) = "Y" then do;
		
			close _distrib_replicates;
		end;
		else do;
		
			create _distrib_replicates var (distrib_names);
				append from distrib_replicates;
			close _distrib_replicates;
		end;
		
		call symputx("num_replicates", num_replicates);
		if symgetn("DEBUG") then print "end module: distrib_main_loop";
	finish;
	
	reset storage=WORK.distrib_modules;
	
	store module=_all_;
quit;