%include "&sysincludefiledir./internal/mcmc_trace_plot_functions.sas";

/*
MCMC Trace Plots

Description:

	Generate trace plots from an NCI multivar MCMC model.
	
Parameter:

	- multivar_mcmc_lib: The SAS library that contains the output datasets of an MCMC model. (default = WORK)
	- multivar_mcmc_model: Prefix for the output datasets of the MCMC model (max 23 characters).
												 See nci_multivar_mcmc macro for more details.
	- do_beta: Flag to specify whether trace plots of fixed effect (beta) coefficients should be produced. (Y/N, default = Y)
	- do_sigma_e: Flag to specify whether trace plots of residual error covariance matrix (Sigma-e) elements should be produced. (Y/N, default = Y)
	- do_sigma_u: Flag to specify whether trace plots of random effect covariance matrix (Sigma-u) elements should be produced. (Y/N, default = Y)
	- do_sigma_overlay: Flag to specify whether overlays of Sigma-e and Sigma-u element trace plots should be produced. (Y/N, default = N)
	- do_alpha1: Flag to specify whether trace plots of never-consumer model (alpha1) coefficients should be produced. 
							 No effect if no never-consumer variable was modeled. (Y/N, default = Y)
	- do_consumer_probability: Flag to specify whether trace plots of mean consumer probability should be produced. 
														 No effect if no never-consumer variable was modeled. (Y/N, default = Y)
	- plot_type: Space-delimited list specifying what types of plots to generate.
							 One or both of the following can be specified:
							 	 - VALUE: Plots of a parameter's value at every MCMC iteration. This is the standard type of trace plot. (default)
							 	 - MEAN: Plots of the rolling mean of a parameter at thinned MCMC iterations after the burn-in period.
							 	 
Output:

	This macro does not produce any output datasets.
	
	The trace plots that are created are output using ODS.
	
Details:

	Trace plots are a method of assessing the convergence of MCMC chain.
	The value of an MCMC parameter (such as a coefficient or a covariance
	matrix element) is plotted against the number of iterations to see how the
	value fluctuates over time. Plots of the parameter values and rolling
	means can be produced. A parameter value that moves around a fixed point with
	a uniform variance over time is said to have converged or mixed well.
*/

%macro trace_plots(multivar_mcmc_lib=WORK,
									 multivar_mcmc_model=,
									 do_beta=Y,
									 do_sigma_e=Y,
									 do_sigma_u=Y,
									 do_sigma_overlay=N,
									 do_alpha1=Y,
									 do_consumer_probability=Y,
									 plot_type=VALUE);
									 
	%local variables num_episodic num_daily i
				 has_never_consumers
				 num_mcmc_iterations num_burn num_thin;
									 
	**types of plots to generate;
	%let plot_type = %upcase(&plot_type.);
	
	**number of iterations and burn-in;
	data _NULL_;
		set &multivar_mcmc_lib..&multivar_mcmc_model._iters;
		
		call symputx("num_mcmc_iterations", num_mcmc_iterations, "L");
		call symputx("num_burn", num_burn, "L");
		call symputx("num_thin", num_thin, "L");
	run;
	
	**get variables to make trace plots for;
	data _NULL_;
		set &multivar_mcmc_lib..&multivar_mcmc_model._vars;
		
		call symputx("num_episodic", num_episodic, "L");
		call symputx("num_daily", num_daily, "L");
		call symputx("num_variables", num_variables, "L");
	run;
		
	data _model_variables (keep = variable1-variable&num_variables.);
		set &multivar_mcmc_lib..&multivar_mcmc_model._vars;
		
		length variable1-variable&num_variables. $32.;
		
		array variables{&num_variables.} variable1-variable&num_variables.;
			
		%if &num_episodic. > 0 %then %do;
		
			array episodic_indicators{&num_episodic.} episodic_ind1-episodic_ind&num_episodic.;
			array episodic_amounts{&num_episodic.} episodic_amt1-episodic_amt&num_episodic.;
		%end;
		
		%if &num_daily. > 0 %then %do;
		
			array daily_amounts{&num_daily.} daily_amt1-daily_amt&num_daily.;
		%end;
		
		%if &num_episodic. > 0 %then %do;
		
			do episodic_j = 1 to num_episodic;
			
				variables{2*episodic_j - 1} = episodic_indicators{episodic_j};
				variables{2*episodic_j} = episodic_amounts{episodic_j};
			end;
		%end;
		
		%if &num_daily. > 0 %then %do;
		
			do daily_j = 1 to num_daily;
			
				variables{2*num_episodic + daily_j} = daily_amounts{daily_j};
			end;
		%end;
	run;
	
	data _NULL_;
		set &multivar_mcmc_lib..&multivar_mcmc_model._covars end=last;
		
		retain max_num_covariates max_num_coefficients;
		
		if num_covariates > max_num_covariates then max_num_covariates = num_covariates;
		if num_coefficients > max_num_coefficients then max_num_coefficients = num_coefficients;
		
		if last = 1 then do;
		
			call symputx("max_num_covariates", max_num_covariates, "L");
			call symputx("max_num_coefficients", max_num_coefficients, "L");
		end;
	run;
	
	**beta trace plots;
	%if %upcase(&do_beta.) = Y %then %do;
	
		%beta_trace_plots(beta=&multivar_mcmc_lib..&multivar_mcmc_model._beta,
											variable_data=_model_variables,
											covariate_data=&multivar_mcmc_lib..&multivar_mcmc_model._covars,
											num_mcmc_iterations=&num_mcmc_iterations.,
											num_burn=&num_burn.,
											num_thin=&num_thin.,
											num_variables=&num_variables.,
											max_num_covariates=&max_num_covariates.,
											max_num_coefficients=&max_num_coefficients.,
											plot_type=&plot_type.);
	%end;
	
	**Sigma-e trace plots;
	%if %upcase(&do_sigma_e.) = Y %then %do;
	
		%sigma_trace_plots(sigma=&multivar_mcmc_lib..&multivar_mcmc_model._sigma_e,
											 component=e,
											 variable_data=_model_variables,
											 num_mcmc_iterations=&num_mcmc_iterations.,
											 num_burn=&num_burn.,
											 num_thin=&num_thin.,
											 num_variables=&num_variables.,
											 plot_type=&plot_type.);
	%end;
	
	**Sigma-u trace plots;
	%if %upcase(&do_sigma_u.) = Y %then %do;
	
		%sigma_trace_plots(sigma=&multivar_mcmc_lib..&multivar_mcmc_model._sigma_u,
											 component=u,
											 variable_data=_model_variables,
											 num_mcmc_iterations=&num_mcmc_iterations.,
											 num_burn=&num_burn.,
											 num_thin=&num_thin.,
											 num_variables=&num_variables.,
											 plot_type=&plot_type.);
	%end;
	
	**sigma overlay plots;
	%if %upcase(&do_sigma_overlay.) = Y %then %do;
	
		%sigma_overlay_plots(sigma_e=&multivar_mcmc_lib..&multivar_mcmc_model._sigma_e,
											 	 sigma_u=&multivar_mcmc_lib..&multivar_mcmc_model._sigma_u,
											 	 variable_data=_model_variables,
											 	 num_mcmc_iterations=&num_mcmc_iterations.,
											 	 num_burn=&num_burn.,
											 	 num_thin=&num_thin.,
											 	 num_variables=&num_variables.,
											 	 plot_type=&plot_type.);
	%end;
	
	**Never-consumer trace plots;
	data _NULL_;
		set &multivar_mcmc_lib..&multivar_mcmc_model._vars;
		
		call symputx("has_never_consumers", has_never_consumers, "L");
	run;
	
	%if &has_never_consumers. = 1 %then %do;
	
		data _NULL_;
			set &multivar_mcmc_lib..&multivar_mcmc_model._covarsnc;
			
			call symputx("num_never_consumer_covariates", num_covariates, "L");
			call symputx("num_never_consumer_coefficients", num_coefficients, "L");
		run;
	
		**alpha1 trace plots;
		%if %upcase(&do_alpha1.) = Y %then %do;
		
			%alpha1_trace_plots(alpha1=&multivar_mcmc_lib..&multivar_mcmc_model._alpha1,
													variable_data=_model_variables,
													never_consumer_covariate_data=&multivar_mcmc_lib..&multivar_mcmc_model._covarsnc,
													num_mcmc_iterations=&num_mcmc_iterations.,
													num_burn=&num_burn.,
													num_thin=&num_thin.,
													num_never_consumer_covariates=&num_never_consumer_covariates.,
													num_never_consumer_coefficients=&num_never_consumer_coefficients.,
													plot_type=&plot_type.);
		%end;
		
		**consumer probability trace plots;
		%if %upcase(&do_consumer_probability.) = Y %then %do;
		
			%consumer_probability_trace_plots(consumer_probabilities=&multivar_mcmc_lib..&multivar_mcmc_model._conprob1,
																				variable_data=_model_variables,
																				num_mcmc_iterations=&num_mcmc_iterations.,
																				num_burn=&num_burn.,
																				num_thin=&num_thin.,
																				plot_type=&plot_type.);
		%end;
	%end;
%mend trace_plots;