/* 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.
2025-11-25 Move update_logger module to v143_modules.
*/

/*
Program Purpose: Store modules for MCMC error-checking in IML
*/

proc iml;

	**Format memory size using appropriate units;
	start format_memsize(size);
	
		if size < 1024 then result = cat(size, " Bytes");
		else if size < 1024**2 then result = cat(size/1024, " kB");
		else if size < 1024**3 then result = cat(size/(1024**2), " MB");
		else result = cat(size/(1024**3), " GB");
		
		return result;
	finish;
	
	**Checks Sigma-u prior;
	start check_sigma_u(sigma_u_prior_data,
											variable_data);
					
		if symgetn("DEBUG") then print "begin module: check_sigma_u";
		alg = symgetn("v143code") + 1;
					
		rc = checkFile(variable_data, "check_sigma_u");
		use (variable_data);
			
			read all var {num_episodic} into num_episodic;
			read all var {num_daily} into num_daily;
		close (variable_data);
		
		num_variables = 2*num_episodic + num_daily;
		
		rc = checkFile(sigma_u_prior_data, "check_sigma_u");
		use (sigma_u_prior_data);
		
			read all var _NUM_ into sigma_u_prior;
		close (sigma_u_prior_data);
		
		**Check if Sigma-u prior is square;
		if nrow(sigma_u_prior) = ncol(sigma_u_prior) then square = 1;
		else square = 0;
		
		if square = 1 then do;
		
			**Check if size of Sigma-u prior matches number of indicators and amounts;
			if square = 1 then do;
				
				if nrow(sigma_u_prior) ^= num_variables then do;
				
					run update_logger(3, "Sigma-U", cat("Size of Sigma-u prior (", nrow(sigma_u_prior), ") does not match number of model variables (", num_variables, ")"));
				end;
			end;
		
			**Check if Sigma-u prior is symmetric;
			if all(abs(sigma_u_prior - sigma_u_prior`) < 10**-8) then symmetric = 1;
			else symmetric = 0;
			
			if symmetric = 0 then do;
			
				run update_logger(3, "Sigma-U", "Sigma-u prior is not symmetric");
			end;
			
			**Check if Sigma-u prior is positive-definite;
			if symmetric = 1 then do;
/* BW CHANGE 2025-1022 */			
				*call eigen(eigvals, eigvecs, sigma_u_prior);
				tmp     = eigenDecomp(sigma_u_prior, alg, 1); 
				eigvals = tmp$"eigenvalues";
				free tmp;
			
				if any(eigvals < 10**-8) then do;
				
					run update_logger(3, "Sigma-U", "Sigma-u prior is not positive-definite");
				end;
			end;
		end;
		else do;
		
			run update_logger(3, "Sigma-U", "Sigma-u prior is not square");
		end;		
		
		if symgetn("DEBUG") then print "end module: check_sigma_u";
	finish;
	
	**Checks design matrices for each model variable;
	start check_design(covariate_data,
										 variable_data,
										 pre_mcmc_data,
										 id);
										 
		if symgetn("DEBUG") then print "begin module: check_design";
		alg = symgetn("v143code") + 1;								 
										 
		input_data = cat(pre_mcmc_data, "_mcmc_in");
		
		submit input_data id;
		
			proc sort data=&input_data out=_input; by &id; run;
		endsubmit;
		
		rc = checkFile("_input", "check_design");
		use _input;
		
			read all var id into record_id;
		close _input;
		
		num_records = nrow(record_id);
		
		**extract variable data;
		rc = checkFile(variable_data, "check_design");
		use (variable_data);
		
			read all var {num_episodic} into num_episodic;
			read all var {num_daily} into num_daily;
			
			if num_episodic > 0 then do;
			
				read all var (cat("episodic_ind", 1:num_episodic)) into episodic_indicators;
				read all var (cat("episodic_amt", 1:num_episodic)) into episodic_amounts;
			end;
			
			if num_daily > 0 then do;
			
				read all var (cat("daily_amt", 1:num_daily)) into daily_amounts;
			end;
			
			read all var {has_never_consumers} into has_never_consumers;
		close (variable_data);
		
		num_variables = 2*num_episodic + num_daily;
		
		model_variables = {};
		if num_episodic > 0 then model_variables = model_variables || rowvec(episodic_indicators` || episodic_amounts`);
		if num_daily > 0 then model_variables = model_variables || daily_amounts;
		
		rc = checkFile(covariate_data, "check_design");
		
		**loop through covariates for each model variable;
		do var_j = 1 to num_variables;
		
			variable = model_variables[var_j];
			
			use (covariate_data);
			
				read point(var_j) var {num_covariates} into num_covariates_j;
				read point(var_j) var {intercept} into intercept_j;
				
				if num_covariates_j > 0 then do;
				
					read point(var_j) var (cat("covariate", 1:num_covariates_j)) into covariates_j;
				end;
			close (covariate_data); 
			
			num_coefficients_j = num_covariates_j + intercept_j;
			
			**Check for empty design matrix;
			if num_coefficients_j = 0 then do;
			
				run update_logger(3, "Design", cat("Empty design matrix for ", variable));
			end;
			
			**Check for linearly dependent columns;
			if num_coefficients_j > 1 then do;
			
				covariate_matrix_j = j(num_records, num_coefficients_j, .);
				
				if intercept_j = 1 then do;
				
					covariate_matrix_j[,1] = 1;
				end;
			
				use _input;
						
					read all var (covariates_j) into covariate_data_j;
				close _input;
				
				covariate_matrix_j[,intercept_j + (1:num_covariates_j)] = covariate_data_j;
/* BW CHANGE 2025-10-22 */	
				*cross = covariate_matrix_j` * covariate_matrix_j;
				cross = matMultAAt(covariate_matrix_j`, alg);
				
				*call eigen(eigvals, eigvecs, cross);
				tmp     = eigenDecomp(cross, alg, 1);  
				eigvals = tmp$"eigenvalues";
				free tmp;
				
				if any(eigvals < 10**-8) then do;
				
					run update_logger(3, "Design", cat("Linearly dependent columns in design matrix for ", variable));
				end;
			end;
		end;
		
		**never-consumer design matrix;
		if has_never_consumers = 1 then do;
		
			subject_id = unique(record_id)`;
			num_subjects = nrow(subject_id);
		
			rc = checkFile(cat(covariate_data, "nc"), "check_design");
			use (cat(covariate_data, "nc"));
			
				read all var {num_covariates} into num_never_consumer_covariates;
				read all var {intercept} into never_consumer_intercept;
				
				if num_never_consumer_covariates > 0 then do;
				
					read all var (cat("covariate", 1:num_never_consumer_covariates)) into never_consumer_covariates;
				end;
			close (cat(covariate_data, "nc"));
			
			num_never_consumer_coefficients = num_never_consumer_covariates + never_consumer_intercept;
			
			**Check for empty design matrix;
			if num_never_consumer_coefficients = 0 then do;
			
				run update_logger(3, "Design", "Empty never-consumer design matrix");
			end;
			
			**Check for linearly dependent columns;
			if num_never_consumer_coefficients > 1 then do;
			
				never_consumer_covariate_matrix = j(num_subjects, num_never_consumer_coefficients, .);
				
				if never_consumer_intercept = 1 then do;
				
					never_consumer_covariate_matrix[,1] = 1;
				end;
			
				use _input;
						
					read point(uniqueby(record_id)) var (never_consumer_covariates) into never_consumer_covariate_data;
				close _input;
				
				never_consumer_covariate_matrix[,never_consumer_intercept + (1:num_never_consumer_covariates)] = never_consumer_covariate_data;
/* BW CHANGE 2025-10-22 */
				*cross = never_consumer_covariate_matrix` * never_consumer_covariate_matrix;
				cross = matMultAAt(never_consumer_covariate_matrix`, alg);
				
				*call eigen(eigvals, eigvecs, cross);
				tmp     = eigenDecomp(cross, alg, 1);
				eigvals = tmp$"eigenvalues";
				free tmp;
				
				if any(eigvals < 10**-8) then do;
				
					run update_logger(3, "Design", "Linearly dependent columns in never-consumer design matrix");
				end;
			end;
		end;
		if symgetn("DEBUG") then print "end module: check_design";
		
	finish;
	
	start log_u_storage(save_mcmc_u,
											save_all_u,
											num_mcmc_iterations,
											num_burn,
											num_thin,
											num_post,
											pre_mcmc_data,
											id,
											variable_data);
											
		element_size = 8;
												
		**Find number of subjects;
		input_data = cat(pre_mcmc_data, "_mcmc_in");
		
		rc = checkFile(input_data, "log_u_storage");
		use (input_data);
		
			read all var id into record_id;
		close (input_data);
		
		subject_id = unique(record_id)`;
		num_subjects = nrow(subject_id);
		
		**Find number of variables;
		rc = checkFile(variable_data, "log_u_storage");
		use (variable_data);
		
			read all var {num_episodic} into num_episodic;
			read all var {num_daily} into num_daily;
		close (variable_data);
		
		num_variables = 2*num_episodic + num_daily;
		
		**calculate U matrix size;
		u_matrix_size = num_subjects * num_variables * element_size;
		
		**Main chain MCMC U matrices;
		if upcase(save_mcmc_u) = "Y" then do;
		
			if upcase(save_all_u) = "Y" then do;
			
				num_mcmc_u = num_mcmc_iterations;
			end;
			else do;
			
				num_mcmc_u = int((num_mcmc_iterations - num_burn)/num_thin);
			end;
			
			run update_logger(1, "U Storage", cat("Estimated memory for main chain MCMC U matrices: ", format_memsize(num_mcmc_u * u_matrix_size)));
		end;
		
		**Post-MCMC U matrices;
		if num_post > 0 then do;
		
			run update_logger(1, "U Storage", cat("Estimated memory for post-MCMC U matrices: ", format_memsize(num_post * u_matrix_size)));
		end;
	finish;
	
	reset storage=WORK.mcmc_modules;
	
	store module=_all_;
quit;