#' Random Effects State Management Routines
#' @name RandomEffectStateManagement
#' @description
#' A forest sampler features two types of state: ephemeral and persistent.
#' Persistent state includes objects like [ForestSamples] and [RandomEffectSamples]
#' which constitute part of the final sampled model.
#' Ephemeral state supports the sampling computations, but is not retained after the sampler finishes.
#'
#' The two primary random-effects-based bits of ephemeral state are the [RandomEffectsModel] and [RandomEffectsTracker] classes,
#' which represent the current state of a random effects model and its corresponding tracking data structures.
#'
#' In a linear sampling loop, this ephemeral state is updated with each iteration of the sampler
#' and any retained forests are copied to a [RandomEffectSamples] object. However, in multi-chain settings,
#' the state of a random effects model must typically be "reset" at the beginning of a new chain. These function enable
#' this process by synchronizing the state of a [RandomEffectsModel] and [RandomEffectsTracker] with a corresponding element
#' of a [RandomEffectSamples] object, or by resetting both to their default (root) state.
#'
#' `resetRandomEffectsModel` resets a `RandomEffectsModel` object based on the parameters indexed by `sample_num` in a `RandomEffectsSamples` object.
#' `resetRandomEffectsTracker` resets a `RandomEffectsTracker` object based on the parameters indexed by `sample_num` in a `RandomEffectsSamples` object.
#'
#' `rootResetRandomEffectsModel` resets a `RandomEffectsModel` object to its "default" state.
#' `rootResetRandomEffectsTracker` resets a `RandomEffectsTracker` object to its "default" state.
#'
#' These functions are intended for advanced use cases in which users require detailed control of sampling algorithms and data structures.
#' Minimal input validation and error checks are performed -- users are responsible for providing the correct inputs.
#' For tutorials on the "proper" usage of the stochtree's advanced workflow, we provide several vignettes at <https://stochtree.ai/>
#' @returns
#' All four functions have no return type and operate in-place on the relevant [RandomEffectsModel] or [RandomEffectsTracker] objects
#' @examples
#' n <- 100
#' p <- 10
#' rfx_group_ids <- sample(1:2, size = n, replace = TRUE)
#' rfx_basis <- matrix(rep(1.0, n), ncol=1)
#' rfx_dataset <- createRandomEffectsDataset(rfx_group_ids, rfx_basis)
#' y <- (-2*(rfx_group_ids==1)+2*(rfx_group_ids==2)) + rnorm(n)
#' y_std <- (y-mean(y))/sd(y)
#' outcome <- createOutcome(y_std)
#' rng <- createCppRNG(1234)
#' num_groups <- length(unique(rfx_group_ids))
#' num_components <- ncol(rfx_basis)
#' rfx_model <- createRandomEffectsModel(num_components, num_groups)
#' rfx_tracker <- createRandomEffectsTracker(rfx_group_ids)
#' rfx_samples <- createRandomEffectSamples(num_components, num_groups, rfx_tracker)
#' alpha_init <- rep(1,num_components)
#' xi_init <- matrix(rep(alpha_init, num_groups),num_components,num_groups)
#' sigma_alpha_init <- diag(1,num_components,num_components)
#' sigma_xi_init <- diag(1,num_components,num_components)
#' sigma_xi_shape <- 1
#' sigma_xi_scale <- 1
#' rfx_model$set_working_parameter(alpha_init)
#' rfx_model$set_group_parameters(xi_init)
#' rfx_model$set_working_parameter_cov(sigma_alpha_init)
#' rfx_model$set_group_parameter_cov(sigma_xi_init)
#' rfx_model$set_variance_prior_shape(sigma_xi_shape)
#' rfx_model$set_variance_prior_scale(sigma_xi_scale)
#' for (i in 1:3) {
#'     rfx_model$sample_random_effect(rfx_dataset=rfx_dataset, residual=outcome,
#'                                    rfx_tracker=rfx_tracker, rfx_samples=rfx_samples,
#'                                    keep_sample=TRUE, global_variance=1.0, rng=rng)
#' }
#' resetRandomEffectsModel(rfx_model, rfx_samples, 0, 1.0)
#' resetRandomEffectsTracker(rfx_tracker, rfx_model, rfx_dataset, outcome, rfx_samples)
#' rootResetRandomEffectsModel(rfx_model, alpha_init, xi_init, sigma_alpha_init,
#'                             sigma_xi_init, sigma_xi_shape, sigma_xi_scale)
#' rootResetRandomEffectsTracker(rfx_tracker, rfx_model, rfx_dataset, outcome)
#'
NULL
#> NULL

#' @title Random Effect Container C++ Wrapper
#' @description
#' Class that wraps the "persistent" aspects of a C++ random effects model, including
#' draws of the parameters and a map from the original label indices to the
#' 0-indexed label numbers used to place group samples in memory (i.e. the
#' first label is stored in column 0 of the sample matrix, the second label
#' is store in column 1 of the sample matrix, etc...)
#'
#' This class is intended for advanced use cases in which users require detailed control of sampling algorithms and data structures.
#' Minimal input validation and error checks are performed -- users are responsible for providing the correct inputs.
#' For tutorials on the "proper" usage of the stochtree's advanced workflow, we provide several vignettes at <https://stochtree.ai/>
#'
#' @description
#' Coordinates various C++ random effects classes and persists those
#' needed for prediction / serialization

RandomEffectSamples <- R6::R6Class(
  classname = "RandomEffectSamples",
  cloneable = FALSE,
  public = list(
    #' @field rfx_container_ptr External pointer to a C++ StochTree::RandomEffectsContainer class
    rfx_container_ptr = NULL,

    #' @field label_mapper_ptr External pointer to a C++ StochTree::LabelMapper class
    label_mapper_ptr = NULL,

    #' @field training_group_ids Unique vector of group IDs that were in the training dataset
    training_group_ids = NULL,

    #' @description
    #' Create a new RandomEffectSamples object.
    #' @return A new `RandomEffectSamples` object.
    initialize = function() {},

    #' @description
    #' Construct RandomEffectSamples object from other "in-session" R objects
    #' @param num_components Number of "components" or bases defining the random effects regression
    #' @param num_groups Number of random effects groups
    #' @param random_effects_tracker Object of type `RandomEffectsTracker`
    #' @return None
    load_in_session = function(
      num_components,
      num_groups,
      random_effects_tracker
    ) {
      # Initialize
      self$rfx_container_ptr <- rfx_container_cpp(
        num_components,
        num_groups
      )
      self$label_mapper_ptr <- rfx_label_mapper_cpp(
        random_effects_tracker$rfx_tracker_ptr
      )
      self$training_group_ids <- rfx_tracker_get_unique_group_ids_cpp(
        random_effects_tracker$rfx_tracker_ptr
      )
    },

    #' @description
    #' Construct RandomEffectSamples object from a json object
    #' @param json_object Object of class `CppJson`
    #' @param json_rfx_container_label Label referring to a particular rfx sample container (i.e. "random_effect_container_0") in the overall json hierarchy
    #' @param json_rfx_mapper_label Label referring to a particular rfx label mapper (i.e. "random_effect_label_mapper_0") in the overall json hierarchy
    #' @param json_rfx_groupids_label Label referring to a particular set of rfx group IDs (i.e. "random_effect_groupids_0") in the overall json hierarchy
    #' @return A new `RandomEffectSamples` object.
    load_from_json = function(
      json_object,
      json_rfx_container_label,
      json_rfx_mapper_label,
      json_rfx_groupids_label
    ) {
      self$rfx_container_ptr <- rfx_container_from_json_cpp(
        json_object$json_ptr,
        json_rfx_container_label
      )
      self$label_mapper_ptr <- rfx_label_mapper_from_json_cpp(
        json_object$json_ptr,
        json_rfx_mapper_label
      )
      self$training_group_ids <- rfx_group_ids_from_json_cpp(
        json_object$json_ptr,
        json_rfx_groupids_label
      )
    },

    #' @description
    #' Append random effect draws to `RandomEffectSamples` object from a json object
    #' @param json_object Object of class `CppJson`
    #' @param json_rfx_container_label Label referring to a particular rfx sample container (i.e. "random_effect_container_0") in the overall json hierarchy
    #' @param json_rfx_mapper_label Label referring to a particular rfx label mapper (i.e. "random_effect_label_mapper_0") in the overall json hierarchy
    #' @param json_rfx_groupids_label Label referring to a particular set of rfx group IDs (i.e. "random_effect_groupids_0") in the overall json hierarchy
    #' @return None
    append_from_json = function(
      json_object,
      json_rfx_container_label,
      json_rfx_mapper_label,
      json_rfx_groupids_label
    ) {
      rfx_container_append_from_json_cpp(
        self$rfx_container_ptr,
        json_object$json_ptr,
        json_rfx_container_label
      )
    },

    #' @description
    #' Construct RandomEffectSamples object from a json object
    #' @param json_string JSON string which parses into object of class `CppJson`
    #' @param json_rfx_container_label Label referring to a particular rfx sample container (i.e. "random_effect_container_0") in the overall json hierarchy
    #' @param json_rfx_mapper_label Label referring to a particular rfx label mapper (i.e. "random_effect_label_mapper_0") in the overall json hierarchy
    #' @param json_rfx_groupids_label Label referring to a particular set of rfx group IDs (i.e. "random_effect_groupids_0") in the overall json hierarchy
    #' @return A new `RandomEffectSamples` object.
    load_from_json_string = function(
      json_string,
      json_rfx_container_label,
      json_rfx_mapper_label,
      json_rfx_groupids_label
    ) {
      self$rfx_container_ptr <- rfx_container_from_json_string_cpp(
        json_string,
        json_rfx_container_label
      )
      self$label_mapper_ptr <- rfx_label_mapper_from_json_string_cpp(
        json_string,
        json_rfx_mapper_label
      )
      self$training_group_ids <- rfx_group_ids_from_json_string_cpp(
        json_string,
        json_rfx_groupids_label
      )
    },

    #' @description
    #' Append random effect draws to `RandomEffectSamples` object from a json object
    #' @param json_string JSON string which parses into object of class `CppJson`
    #' @param json_rfx_container_label Label referring to a particular rfx sample container (i.e. "random_effect_container_0") in the overall json hierarchy
    #' @param json_rfx_mapper_label Label referring to a particular rfx label mapper (i.e. "random_effect_label_mapper_0") in the overall json hierarchy
    #' @param json_rfx_groupids_label Label referring to a particular set of rfx group IDs (i.e. "random_effect_groupids_0") in the overall json hierarchy
    #' @return None
    append_from_json_string = function(
      json_string,
      json_rfx_container_label,
      json_rfx_mapper_label,
      json_rfx_groupids_label
    ) {
      # Append RFX objects
      rfx_container_append_from_json_string_cpp(
        self$rfx_container_ptr,
        json_string,
        json_rfx_container_label
      )
    },

    #' @description
    #' Predict random effects for each observation implied by `rfx_group_ids` and `rfx_basis`.
    #' If a random effects model is "intercept-only" the `rfx_basis` will be a vector of ones of size `length(rfx_group_ids)`.
    #' @param rfx_group_ids Indices of random effects groups in a prediction set
    #' @param rfx_basis (Optional) Basis used for random effects prediction
    #' @return Matrix with as many rows as observations provided and as many columns as samples drawn of the model.
    predict = function(rfx_group_ids, rfx_basis = NULL) {
      num_obs = length(rfx_group_ids)
      if (is.null(rfx_basis)) {
        rfx_basis <- matrix(rep(1, num_obs), ncol = 1)
      }
      num_samples = rfx_container_num_samples_cpp(self$rfx_container_ptr)
      num_components = rfx_container_num_components_cpp(
        self$rfx_container_ptr
      )
      num_groups = rfx_container_num_groups_cpp(self$rfx_container_ptr)
      rfx_group_ids_int <- as.integer(rfx_group_ids)
      stopifnot(sum(abs(rfx_group_ids_int - rfx_group_ids)) < 1e-6)
      stopifnot(sum(!(rfx_group_ids %in% self$training_group_ids)) == 0)
      stopifnot(ncol(rfx_basis) == num_components)
      rfx_dataset <- createRandomEffectsDataset(
        rfx_group_ids_int,
        rfx_basis
      )
      output <- rfx_container_predict_cpp(
        self$rfx_container_ptr,
        rfx_dataset$data_ptr,
        self$label_mapper_ptr
      )
      dim(output) <- c(num_obs, num_samples)
      return(output)
    },

    #' @description
    #' Extract the random effects parameters sampled. With the "redundant parameterization"
    #' of Gelman et al (2008), this includes four parameters: alpha (the "working parameter"
    #' shared across every group), xi (the "group parameter" sampled separately for each group),
    #' beta (the product of alpha and xi, which corresponds to the overall group-level random effects),
    #' and sigma (group-independent prior variance for each component of xi).
    #' @return List of arrays. The alpha array has dimension (`num_components`, `num_samples`) and is simply a vector if `num_components = 1`.
    #' The xi and beta arrays have dimension (`num_components`, `num_groups`, `num_samples`) and are simply matrices if `num_components = 1`.
    #' The sigma array has dimension (`num_components`, `num_samples`) and is simply a vector if `num_components = 1`.
    extract_parameter_samples = function() {
      num_samples = rfx_container_num_samples_cpp(self$rfx_container_ptr)
      num_components = rfx_container_num_components_cpp(
        self$rfx_container_ptr
      )
      num_groups = rfx_container_num_groups_cpp(self$rfx_container_ptr)
      beta_samples <- rfx_container_get_beta_cpp(self$rfx_container_ptr)
      xi_samples <- rfx_container_get_xi_cpp(self$rfx_container_ptr)
      alpha_samples <- rfx_container_get_alpha_cpp(self$rfx_container_ptr)
      sigma_samples <- rfx_container_get_sigma_cpp(self$rfx_container_ptr)
      if (num_components == 1) {
        dim(beta_samples) <- c(num_groups, num_samples)
        dim(xi_samples) <- c(num_groups, num_samples)
      } else if (num_components > 1) {
        dim(beta_samples) <- c(num_components, num_groups, num_samples)
        dim(xi_samples) <- c(num_components, num_groups, num_samples)
        dim(alpha_samples) <- c(num_components, num_samples)
        dim(sigma_samples) <- c(num_components, num_samples)
      } else {
        stop(
          "Invalid random effects sample container, num_components is less than 1"
        )
      }

      output = list(
        "beta_samples" = beta_samples,
        "xi_samples" = xi_samples,
        "alpha_samples" = alpha_samples,
        "sigma_samples" = sigma_samples
      )
      return(output)
    },

    #' @description
    #' Modify the `RandomEffectsSamples` object by removing the parameter samples index by `sample_num`.
    #' @param sample_num Index of the RFX sample to be removed
    delete_sample = function(sample_num) {
      rfx_container_delete_sample_cpp(self$rfx_container_ptr, sample_num)
    },

    #' @description
    #' Convert the mapping of group IDs to random effect components indices from C++ to R native format
    #' @return List mapping group ID to random effect components.
    extract_label_mapping = function() {
      keys_and_vals <- rfx_label_mapper_to_list_cpp(self$label_mapper_ptr)
      result <- as.list(keys_and_vals[[2]] + 1)
      setNames(result, keys_and_vals[[1]])
      return(result)
    },

    #' @description
    #' Query the number of samples in the `RandomEffectsSamples` object.
    #' @return Integer number of samples
    num_samples = function() {
      return(rfx_container_num_samples_cpp(self$rfx_container_ptr))
    },

    #' @description
    #' Query the number of components in the `RandomEffectsSamples` object.
    #' @return Integer number of components
    num_components = function() {
      return(rfx_container_num_components_cpp(self$rfx_container_ptr))
    },

    #' @description
    #' Query the number of groups in the `RandomEffectsSamples` object.
    #' @return Integer number of groups
    num_groups = function() {
      return(rfx_container_num_groups_cpp(self$rfx_container_ptr))
    }
  )
)

#' @title Random Effects Tracker C++ Wrapper
#'
#' @description
#' Class that defines a "tracker" for random effects models, most notably
#' storing the data indices available in each group for quicker posterior
#' computation and sampling of random effects terms.
#' The class stores a mapping from every observation to its group index, a mapping
#' from group indices to the training sample observations available in that
#' group, and predictions for each observation.
#'
#' This class is intended for advanced use cases in which users require detailed control of sampling algorithms and data structures.
#' Minimal input validation and error checks are performed -- users are responsible for providing the correct inputs.
#' For tutorials on the "proper" usage of the stochtree's advanced workflow, we provide several vignettes at <https://stochtree.ai/>

RandomEffectsTracker <- R6::R6Class(
  classname = "RandomEffectsTracker",
  cloneable = FALSE,
  public = list(
    #' @field rfx_tracker_ptr External pointer to a C++ StochTree::RandomEffectsTracker class
    rfx_tracker_ptr = NULL,

    #' @description
    #' Create a new RandomEffectsTracker object.
    #' @param rfx_group_indices Integer indices indicating groups used to define random effects
    #' @return A new `RandomEffectsTracker` object.
    initialize = function(rfx_group_indices) {
      # Initialize
      self$rfx_tracker_ptr <- rfx_tracker_cpp(rfx_group_indices)
    }
  )
)

#' @title Random Effects Model C++ Wrapper
#' @description
#' The core "model" class for sampling random effects.
#' Stores current model state, prior parameters, and procedures for
#' sampling from the conditional posterior of each parameter.
#'
#' This class is intended for advanced use cases in which users require detailed control of sampling algorithms and data structures.
#' Minimal input validation and error checks are performed -- users are responsible for providing the correct inputs.
#' For tutorials on the "proper" usage of the stochtree's advanced workflow, we provide several vignettes at <https://stochtree.ai/>

RandomEffectsModel <- R6::R6Class(
  classname = "RandomEffectsModel",
  cloneable = FALSE,
  public = list(
    #' @field rfx_model_ptr External pointer to a C++ StochTree::RandomEffectsModel class
    rfx_model_ptr = NULL,

    #' @field num_groups Number of groups in the random effects model
    num_groups = NULL,

    #' @field num_components Number of components (i.e. dimension of basis) in the random effects model
    num_components = NULL,

    #' @description
    #' Create a new RandomEffectsModel object.
    #' @param num_components Number of "components" or bases defining the random effects regression
    #' @param num_groups Number of random effects groups
    #' @return A new `RandomEffectsModel` object.
    initialize = function(num_components, num_groups) {
      # Initialize
      self$rfx_model_ptr <- rfx_model_cpp(num_components, num_groups)
      self$num_components <- num_components
      self$num_groups <- num_groups
    },

    #' @description
    #' Sample from random effects model.
    #' @param rfx_dataset Object of type `RandomEffectsDataset`
    #' @param residual Object of type `Outcome`
    #' @param rfx_tracker Object of type `RandomEffectsTracker`
    #' @param rfx_samples Object of type `RandomEffectSamples`
    #' @param keep_sample Whether sample should be retained in `rfx_samples`. If `FALSE`, the state of `rfx_tracker` will be updated, but the parameter values will not be added to the sample container. Samples are commonly discarded due to burn-in or thinning.
    #' @param global_variance Scalar global variance parameter
    #' @param rng Object of type `CppRNG`
    #' @return None
    sample_random_effect = function(
      rfx_dataset,
      residual,
      rfx_tracker,
      rfx_samples,
      keep_sample,
      global_variance,
      rng
    ) {
      rfx_model_sample_random_effects_cpp(
        self$rfx_model_ptr,
        rfx_dataset$data_ptr,
        residual$data_ptr,
        rfx_tracker$rfx_tracker_ptr,
        rfx_samples$rfx_container_ptr,
        keep_sample,
        global_variance,
        rng$rng_ptr
      )
    },

    #' @description
    #' Predict from (a single sample of a) random effects model.
    #' @param rfx_dataset Object of type `RandomEffectsDataset`
    #' @param rfx_tracker Object of type `RandomEffectsTracker`
    #' @return Vector of predictions with size matching number of observations in rfx_dataset
    predict = function(rfx_dataset, rfx_tracker) {
      pred <- rfx_model_predict_cpp(
        self$rfx_model_ptr,
        rfx_dataset$data_ptr,
        rfx_tracker$rfx_tracker_ptr
      )
      return(pred)
    },

    #' @description
    #' Set value for the "working parameter." This is typically
    #' used for initialization, but could also be used to interrupt
    #' or override the sampler.
    #' @param value Parameter input
    #' @return None
    set_working_parameter = function(value) {
      stopifnot(is.double(value))
      stopifnot(!is.matrix(value))
      stopifnot(length(value) == self$num_components)
      rfx_model_set_working_parameter_cpp(self$rfx_model_ptr, value)
    },

    #' @description
    #' Set value for the "group parameters." This is typically
    #' used for initialization, but could also be used to interrupt
    #' or override the sampler.
    #' @param value Parameter input
    #' @return None
    set_group_parameters = function(value) {
      stopifnot(is.double(value))
      stopifnot(is.matrix(value))
      stopifnot(nrow(value) == self$num_components)
      stopifnot(ncol(value) == self$num_groups)
      rfx_model_set_group_parameters_cpp(self$rfx_model_ptr, value)
    },

    #' @description
    #' Set value for the working parameter covariance. This is typically
    #' used for initialization, but could also be used to interrupt
    #' or override the sampler.
    #' @param value Parameter input
    #' @return None
    set_working_parameter_cov = function(value) {
      stopifnot(is.double(value))
      stopifnot(is.matrix(value))
      stopifnot(nrow(value) == self$num_components)
      stopifnot(ncol(value) == self$num_components)
      rfx_model_set_working_parameter_covariance_cpp(
        self$rfx_model_ptr,
        value
      )
    },

    #' @description
    #' Set value for the group parameter covariance. This is typically
    #' used for initialization, but could also be used to interrupt
    #' or override the sampler.
    #' @param value Parameter input
    #' @return None
    set_group_parameter_cov = function(value) {
      stopifnot(is.double(value))
      stopifnot(is.matrix(value))
      stopifnot(nrow(value) == self$num_components)
      stopifnot(ncol(value) == self$num_components)
      rfx_model_set_group_parameter_covariance_cpp(
        self$rfx_model_ptr,
        value
      )
    },

    #' @description
    #' Set shape parameter for the group parameter variance prior.
    #' @param value Parameter input
    #' @return None
    set_variance_prior_shape = function(value) {
      stopifnot(is.double(value))
      stopifnot(!is.matrix(value))
      stopifnot(length(value) == 1)
      rfx_model_set_variance_prior_shape_cpp(self$rfx_model_ptr, value)
    },

    #' @description
    #' Set shape parameter for the group parameter variance prior.
    #' @param value Parameter input
    #' @return None
    set_variance_prior_scale = function(value) {
      stopifnot(is.double(value))
      stopifnot(!is.matrix(value))
      stopifnot(length(value) == 1)
      rfx_model_set_variance_prior_scale_cpp(self$rfx_model_ptr, value)
    }
  )
)

#' @title Create RandomEffectSamples Object
#' @description
#' Create a `RandomEffectSamples` object
#'
#' This function is intended for advanced use cases in which users require detailed control of sampling algorithms and data structures.
#' Minimal input validation and error checks are performed -- users are responsible for providing the correct inputs.
#' For tutorials on the "proper" usage of the stochtree's advanced workflow, we provide several vignettes at <https://stochtree.ai/>
#'
#' @param num_components Number of "components" or bases defining the random effects regression
#' @param num_groups Number of random effects groups
#' @param random_effects_tracker Object of type `RandomEffectsTracker`
#' @return `RandomEffectSamples` object
#' @export
#'
#' @examples
#' n <- 100
#' rfx_group_ids <- sample(1:2, size = n, replace = TRUE)
#' rfx_basis <- matrix(rep(1.0, n), ncol=1)
#' num_groups <- length(unique(rfx_group_ids))
#' num_components <- ncol(rfx_basis)
#' rfx_tracker <- createRandomEffectsTracker(rfx_group_ids)
#' rfx_samples <- createRandomEffectSamples(num_components, num_groups, rfx_tracker)
createRandomEffectSamples <- function(
  num_components,
  num_groups,
  random_effects_tracker
) {
  invisible(output <- RandomEffectSamples$new())
  output$load_in_session(num_components, num_groups, random_effects_tracker)
  return(output)
}

#' @title Create RandomEffectsTracker Object
#' @description
#' Create a `RandomEffectsTracker` object
#'
#' This function is intended for advanced use cases in which users require detailed control of sampling algorithms and data structures.
#' Minimal input validation and error checks are performed -- users are responsible for providing the correct inputs.
#' For tutorials on the "proper" usage of the stochtree's advanced workflow, we provide several vignettes at <https://stochtree.ai/>
#'
#' @param rfx_group_indices Integer indices indicating groups used to define random effects
#' @return `RandomEffectsTracker` object
#' @export
#'
#' @examples
#' n <- 100
#' rfx_group_ids <- sample(1:2, size = n, replace = TRUE)
#' rfx_basis <- matrix(rep(1.0, n), ncol=1)
#' num_groups <- length(unique(rfx_group_ids))
#' num_components <- ncol(rfx_basis)
#' rfx_tracker <- createRandomEffectsTracker(rfx_group_ids)
createRandomEffectsTracker <- function(rfx_group_indices) {
  return(invisible((RandomEffectsTracker$new(rfx_group_indices))))
}

#' @title Create RandomEffectsModel Object
#' @description
#' Create a `RandomEffectsModel` object
#'
#' This function is intended for advanced use cases in which users require detailed control of sampling algorithms and data structures.
#' Minimal input validation and error checks are performed -- users are responsible for providing the correct inputs.
#' For tutorials on the "proper" usage of the stochtree's advanced workflow, we provide several vignettes at <https://stochtree.ai/>
#'
#' @param num_components Number of "components" or bases defining the random effects regression
#' @param num_groups Number of random effects groups
#' @return `RandomEffectsModel` object
#' @export
#'
#' @examples
#' n <- 100
#' rfx_group_ids <- sample(1:2, size = n, replace = TRUE)
#' rfx_basis <- matrix(rep(1.0, n), ncol=1)
#' num_groups <- length(unique(rfx_group_ids))
#' num_components <- ncol(rfx_basis)
#' rfx_model <- createRandomEffectsModel(num_components, num_groups)
createRandomEffectsModel <- function(num_components, num_groups) {
  return(invisible((RandomEffectsModel$new(num_components, num_groups))))
}

#' @title Reset RandomEffectsModel Object
#' @rdname RandomEffectStateManagement
#' @param rfx_model Object of type `RandomEffectsModel`.
#' @param rfx_samples Object of type `RandomEffectSamples`.
#' @param sample_num Index of sample stored in `rfx_samples` from which to reset the state of a random effects model. Zero-indexed, so resetting based on the first sample would require setting `sample_num = 0`.
#' @param sigma_alpha_init Initial value of the "working parameter" scale parameter.
#' @export
resetRandomEffectsModel <- function(
  rfx_model,
  rfx_samples,
  sample_num,
  sigma_alpha_init
) {
  if (!is.matrix(sigma_alpha_init)) {
    if (!is.double(sigma_alpha_init)) {
      stop("`sigma_alpha_init` must be a numeric scalar or matrix")
    }
    sigma_alpha_init <- as.matrix(sigma_alpha_init)
  }
  reset_rfx_model_cpp(
    rfx_model$rfx_model_ptr,
    rfx_samples$rfx_container_ptr,
    sample_num
  )
  rfx_model$set_working_parameter_cov(sigma_alpha_init)
}

#' @title Reset RandomEffectsTracker Object
#' @rdname RandomEffectStateManagement
#' @param rfx_tracker Object of type `RandomEffectsTracker`.
#' @param rfx_model Object of type `RandomEffectsModel`.
#' @param rfx_dataset Object of type `RandomEffectsDataset`.
#' @param residual Object of type `Outcome`.
#' @param rfx_samples Object of type `RandomEffectSamples`.
#' @export
resetRandomEffectsTracker <- function(
  rfx_tracker,
  rfx_model,
  rfx_dataset,
  residual,
  rfx_samples
) {
  reset_rfx_tracker_cpp(
    rfx_tracker$rfx_tracker_ptr,
    rfx_dataset$data_ptr,
    residual$data_ptr,
    rfx_model$rfx_model_ptr
  )
}

#' @title Reset RandomEffectsModel Object to Default State
#' @rdname RandomEffectStateManagement
#' @param rfx_model Object of type `RandomEffectsModel`.
#' @param alpha_init Initial value of the "working parameter".
#' @param xi_init Initial value of the "group parameters".
#' @param sigma_alpha_init Initial value of the "working parameter" scale parameter.
#' @param sigma_xi_init Initial value of the "group parameters" scale parameter.
#' @param sigma_xi_shape Shape parameter for the inverse gamma variance model on the group parameters.
#' @param sigma_xi_scale Scale parameter for the inverse gamma variance model on the group parameters.
#' @export
rootResetRandomEffectsModel <- function(
  rfx_model,
  alpha_init,
  xi_init,
  sigma_alpha_init,
  sigma_xi_init,
  sigma_xi_shape,
  sigma_xi_scale
) {
  rfx_model$set_working_parameter(alpha_init)
  rfx_model$set_group_parameters(xi_init)
  rfx_model$set_working_parameter_cov(sigma_alpha_init)
  rfx_model$set_group_parameter_cov(sigma_xi_init)
  rfx_model$set_variance_prior_shape(sigma_xi_shape)
  rfx_model$set_variance_prior_scale(sigma_xi_scale)
}

#' @title Reset RandomEffectsTracker Object to Default State
#' @rdname RandomEffectStateManagement
#' @param rfx_tracker Object of type `RandomEffectsTracker`.
#' @param rfx_model Object of type `RandomEffectsModel`.
#' @param rfx_dataset Object of type `RandomEffectsDataset`.
#' @param residual Object of type `Outcome`.
#' @export
rootResetRandomEffectsTracker <- function(
  rfx_tracker,
  rfx_model,
  rfx_dataset,
  residual
) {
  root_reset_rfx_tracker_cpp(
    rfx_tracker$rfx_tracker_ptr,
    rfx_dataset$data_ptr,
    residual$data_ptr,
    rfx_model$rfx_model_ptr
  )
}

#' @title Summarize a RandomEffectSamples Object
#' @description Prints a summary of the RandomEffectSamples object, including number of forests and the underlying model of each forest.
#' @param x RandomEffectSamples object
#' @param ... Additional arguments
#' @export
#' @return RandomEffectSamples object unchanged after printing summary
print.RandomEffectSamples <- function(x, ...) {
  # Construct summary message
  num_samples <- x$num_samples()
  num_components <- x$num_components()
  num_groups <- x$num_groups()
  summary_message <- paste0(
    "Random effects container with ",
    num_samples,
    " posterior draws of a model with ",
    num_components,
    " components and ",
    num_groups,
    " groups."
  )

  # Print the random effects container details
  cat(summary_message, "\n")

  # Return random effects container invisibly
  invisible(x)
}
