/* 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.
		   Check the log argument in marginal_likelihood.
*/

/*
Program Purpose:
IML modules for estimating MCMC log-likelihood
*/

proc iml;

	**Performs LDLt decomposition of a matrix;
	start ldlt(l, d, x);
	
		u = root(x);
		diag_u = vecdiag(u);
		
		l = (u / diag_u)`;
		d = diag(diag_u ## 2);
	finish;
	
	**Performs rank 1 update to LDLt decomposition of a matrix;
	**From Gill, Golub, Murray, and Saunders (1974);
	start ldlt_update1(l_new, d_new, l, d, v, sigma);
	
		n = nrow(l);
		
		w = v;
		alpha = sigma;
		
		l_new = l;
		d_new = d;
		
		do j = 1 to n;
		
			p = w[j];
			
			d_new[j, j] = d[j, j] + alpha * p**2;
			
			beta = p*alpha/d_new[j, j];
			alpha = d[j, j]*alpha/d_new[j, j];
			
			do i = j+1 to n;
			
				w[i] = w[i] - p*l[i, j];
				
				l_new[i, j] = l[i, j] + beta*w[i];
			end;
		end;
	finish;
	
	**Approximates multivariate normal CDF using ME method;
	**From Bhat (2018);
	start cdf_mvnorm(x, mu=, sigma=);
	
		m = ncol(x);
		
		if isskipped(mu) then mu = j(1, m, 0);
		if isskipped(sigma) then sigma = i(m);
		
		**Normalize variables;
		z = (x - mu)/sqrt(vecdiag(sigma)`);
		
		mu_z = j(1, m, 0);
		sigma_z = cov2corr(sigma);
		
		**Special cases for univariate and bivariate CDFs;
		if m = 1 then do;
		
			p = cdf("Normal", z);
			
			return p;
		end;
		
		if m = 2 then do;
		
			p = probbnrm(z[1], z[2], sigma_z[1, 2]);
			
			return p;
		end;
		
		**Order variables to increase efficiency;
		call sortndx(z_order, z`);
		z = z[,z_order];
		sigma_z = sigma_z[z_order, z_order];
		
		**LDLt decomposition of covariance matrix;
		run ldlt(sigma_z_l, sigma_z_d, sigma_z);
		
		**Begin with univariate normal CDF for first variable;
		p = cdf("Normal", z[1]);
		
		**Iteratively approximate CDFs for variables conditional on joint CDF of previous variables;
		do h = 1 to m-1;
		
			**Calculate mean and variance of truncated normal distribution truncated at current variable;
			lambda = -pdf("Normal", z[h])/cdf("Normal", z[h]);
			mu_trunc = mu_z[1] + sigma_z[1, 1] * lambda;
			sigma_trunc = sigma_z[1, 1] * (1 + lambda * (z[h] - lambda));
			
			**Update mean and covariance matrix for remaining variables;
			mu_z = mu_z[,2:(m-h+1)] + (mu_trunc - mu_z[1]) # sigma_z_l[2:(m-h+1), 1]`;
			run ldlt_update1(sigma_z_l, sigma_z_d, sigma_z_l[2:(m-h+1), 2:(m-h+1)], sigma_z_d[2:(m-h+1), 2:(m-h+1)], sigma_z_l[2:(m-h+1), 1]`, sigma_trunc);
			
			**Calculate CDF of next variable;
			p = p * cdf("Normal", z[h+1], mu_z[1], sqrt(sigma_z_d[1, 1]));
		end;
		
		return p;
	finish;
	
	**Joint log-likelihood of observed parameters and random effects;
	**This is the likelihood from Appendix A.4 from Zhang, et al. (2011) with latent W matrix elements integrated out;
	**For never-consumer models, latent Ni variables are integrated out as in Bhadra, et al. (2020);
	start joint_likelihood(u_matrix_i)
				global(_w_matrix_i,
							 _episodic_indicator_i,
							 _xbeta_i,
							 _sigma_u,
							 _sigma_e,
							 _recall_availability_i,
							 _num_recalls,
							 _num_episodic,
							 _num_daily,
							 _num_variables,
							 _has_never_consumers,
							 _alpha1,
							 _g_i);
						
		alg = symgetn("v143code") + 1;
						
		pi = arcos(-1);
		
		ll = 0;
		
		if _has_never_consumers = 1 then do;
		
			has_no_consumption_days = all(_episodic_indicator_i[,1] = 0);
		
			g_alpha = _g_i * _alpha1;
			
			ll = log(max(cdf("Normal", g_alpha), 1e-15));
			if has_no_consumption_days = 1 then ll_nev = log(max(1 - cdf("Normal", g_alpha), 1e-15));
		end;
		
		do day_k = 1 to _num_recalls;
		
			if _recall_availability_i[day_k] = 1 then do;
			
				xbeta_u = _xbeta_i[day_k,] + u_matrix_i;
				w_minus_xbeta_u = _w_matrix_i[day_k,] - xbeta_u;
				
				**Observed variables;
				is_observed = j(_num_variables, 1, 0);
				
				if _num_episodic > 0 then is_observed[2 # (1:_num_episodic)] = _episodic_indicator_i[day_k,];
				
				if _num_daily > 0 then is_observed[2*_num_episodic + (1:_num_daily)] = 1;
				
				observed = loc(is_observed);
				num_observed = sum(is_observed);
				
				**Log-likelihood of probability of consumption given observed amounts;
				if _num_episodic > 0 then do;
				
					indicators = 2 # (1:_num_episodic) - 1;
				
					**Conditional distribution of probability of consumption given observed amounts;
					xbeta_u_prob = xbeta_u[,indicators];
					sigma_e_prob = _sigma_e[indicators, indicators];
					
					if num_observed > 0 then do;
/* BW CHANGE 2025-10-22 */		
						*xbeta_u_prob = xbeta_u_prob + w_minus_xbeta_u[,observed] * inv(_sigma_e[observed, observed]) * _sigma_e[observed, indicators];
						*sigma_e_prob = sigma_e_prob - _sigma_e[indicators, observed] * inv(_sigma_e[observed, observed]) * _sigma_e[observed, indicators];
						tmp          = inv(_sigma_e[observed, observed]);
						xbeta_u_prob = xbeta_u_prob + matMultABC(w_minus_xbeta_u[,observed],     tmp, _sigma_e[observed, indicators], alg);
						sigma_e_prob = sigma_e_prob - matMultABC(_sigma_e[indicators, observed], tmp, _sigma_e[observed, indicators], alg);
                        free tmp;
					end;
					
					prob = cdf_mvnorm((2 # _episodic_indicator_i[day_k,] - 1) # xbeta_u_prob, , sigma_e_prob);
					ll_prob = log(max(prob, 1e-15));
				end;
				else ll_prob = 0;
				
				**Log-likelihood of observed amounts;
				if num_observed > 0 then do;
				
					w_minus_xbeta_u_amt = w_minus_xbeta_u[,observed];
					sigma_e_amt = _sigma_e[observed, observed];
/* BW CHANGE 2025-10-22 */					
					*ll_amt = -0.5*num_observed*log(2*pi) - 0.5*log(det(sigma_e_amt)) - 0.5*(w_minus_xbeta_u_amt * inv(sigma_e_amt) * w_minus_xbeta_u_amt`);
					ll_amt = -0.5*num_observed*log(2*pi) - 0.5*log(det(sigma_e_amt)) - 0.5*matMultABAt(w_minus_xbeta_u_amt, inv(sigma_e_amt), alg);
				end;
				else ll_amt = 0;
				
				ll = ll + ll_prob + ll_amt;
				
				if _has_never_consumers = 1 then do;
				
					if has_no_consumption_days = 1 then ll_nev = ll_nev + ll_prob + ll_amt - log(max(1 - cdf("Normal", xbeta_u[1]), 1e-15));
				end;
			end;
		end;
		
		if _has_never_consumers = 1 then do;
		
			**Log-Sum-Exp trick to add consumer and never-consumer likelihoods;
			if has_no_consumption_days = 1 then do;
			
				ll_min = min(ll, ll_nev);
				ll = log(exp(ll - ll_min) + exp(ll_nev - ll_min)) + ll_min;
			end;
		end;
/* BW CHANGE 2025-10-22 */		
		**Log-likelihood of random effects;
		*ll = ll - 0.5*_num_variables*log(2*pi) - 0.5*log(det(_sigma_u)) - 0.5*(u_matrix_i * inv(_sigma_u) * u_matrix_i`);
		ll = ll - 0.5*_num_variables*log(2*pi) - 0.5*log(det(_sigma_u)) - 0.5*matMultABAt(u_matrix_i, inv(_sigma_u), alg);
		
		
		return ll;
	finish;
	
	start marginal_likelihood(w_matrix,
														episodic_indicator_matrices,
														xbeta,
														sigma_u,
														sigma_e,
														recall_availability,
														subject_weighting,
														num_subjects,
														num_recalls,
														num_episodic,
														num_daily,
														num_variables,
														has_never_consumers,
														alpha1,
														never_consumer_covariate_matrix)							
				global(_w_matrix_i,
							 _episodic_indicator_i,
							 _xbeta_i,
							 _sigma_u,
							 _sigma_e,
							 _recall_availability_i,
							 _num_recalls,
							 _num_episodic,
							 _num_daily,
							 _num_variables,
							 _has_never_consumers,
							 _alpha1,
							 _g_i);
		
		pi = arcos(-1);
		
		**Set decimal precision level for convergence tolerance;
		precision = 6;
		
		**Set option and control lists for optimizer call;
		optimizer_options = j(11, 1, .);
		optimizer_options[1] = 1; **Maximization;
		optimizer_options[4] = 3; **BFGS algorithm;
		optimizer_options[8] = 01; **Central difference for derivatives, interval calculated using only algorithm precision;
		
		optimizer_control = j(10, 1, .);
		optimizer_control[8] = precision; **Algorithm precision;
		
		**Set control list for derivative evaluation call;
		derivative_control = j(3, 1, .);
		derivative_control[2] = 1; **Central difference for derivatives, interval calculated using only algorithm precision;
		derivative_control[3] = precision; **Algorithm precision;
		
		**Find joint log-likelihood at posterior mode of U matrix;
		ll = j(num_subjects, 1, 0);
		log_det_hessian = j(num_subjects, 1, .);
		
		do i = 1 to num_subjects;
		
			**Define variables for log-likelihood function;
			_w_matrix_i = j(num_recalls, num_variables, .);
			
			if num_episodic > 0 then _episodic_indicator_i = j(num_recalls, num_episodic, .);
			else _episodic_indicator_i = {};
			
			_xbeta_i = j(num_recalls, num_variables, .);
			
			do day_k = 1 to num_recalls;
			
				w_matrix_k = w_matrix$day_k;
				_w_matrix_i[day_k,] = w_matrix_k[i,];
				
				if num_episodic > 0 then do;
				
					episodic_indicator_k = episodic_indicator_matrices$day_k;
					_episodic_indicator_i[day_k,] = episodic_indicator_k[i,];
				end;
				
				xbeta_k = xbeta$day_k;
				_xbeta_i[day_k,] = xbeta_k[i,];
			end;
			
			_sigma_u = sigma_u;
		
			_sigma_e = sigma_e;
			
			_recall_availability_i = recall_availability[i,];
			
			_num_recalls = num_recalls;
			_num_episodic = num_episodic;
			_num_daily = num_daily;
			_num_variables = num_variables;
			
			_has_never_consumers = has_never_consumers;
			if _has_never_consumers = 1 then do;
			
				_alpha1 = alpha1;
				_g_i = never_consumer_covariate_matrix[i,];
			end;
			else do;
			
				_alpha1 = {};
				_g_i = {};
			end;
			
			**Find value and Hessian of Ui at posterior mode;
			u0 = j(num_variables, 1, 0);
			call nlpqn(rc, ui_opt, "joint_likelihood", u0, optimizer_options, , , optimizer_control);
			call nlpfdd(ll_opt, ll_grad, ll_hess, "joint_likelihood", ui_opt, derivative_control);
			
			ll[i] = ll_opt;
/* CHANGE 2025-10-30 */
			tmp   = det(-ll_hess);
			if (tmp <= 0) then do;
				run update_logger(3, "marginal_likelihood module", "log argument <= 0");
				return .;
			end;
			log_det_hessian[i] = log(tmp);
			
		end;
		
		**Estimate marginal log-likelihood using Laplace approximation;
		ll = ll + 0.5*num_variables*log(2*pi) - 0.5 # log_det_hessian;
		
		**Return weighted sum of marginal log-likelihood;
		ll_total = sum(ll # subject_weighting);
		
		return ll_total;
	finish;
	
	reset storage=WORK.mcmc_modules;
	
	store module=_all_;
quit;