#' The outcome model function. The outcome model assumes the noise follows a double-censored Tobit model with a subject-level random intercept.
#'
#' @param Y the outcome variable. The datatype is numeric.
#' @param X_continuous the dataframe with continuous variables. The datatype is data.frame.
#' @param X_categorical the dataframe with categorical variables. The datatype is data.frame.
#' @param ni_vector the numeric vector with the number of observations of each subject. The datatype is numeric.
#' @param group_info the numeric vector with the group information of each observation. The datatype is numeric.
#' @param num_warmup the number of warmup iterations. The datatype is integer.
#' @param num_samples the number of post-warmup iterations. The datatype is integer.
#' @param thin the period between saved samples. This should typically be left at its default (no thinning) unless memory is a problem.  The datatype is integer.
#' @param L the number of trees. The datatype is integer.
#' @param alpha the tree prior parameters. alpha*(1+depth)^(-beta) represents the prior probability of splitting at one of the cutpoints. Check equation (4) and related descriptions from "Stochastic Tree Ensembles for Regularized Nonlinear Regression" for more details. The datatype is double.
#' @param beta the tree prior parameters. alpha*(1+depth)^(-beta) represents the prior probability of splitting at one of the cutpoints. Check equation (4) and related descriptions from "Stochastic Tree Ensembles for Regularized Nonlinear Regression" for more details. The datatype is double.
#' @param leaf_model_scale the prior variance on leaf mean equals to \code{leaf_model_scale/L}. The datatype is double.
#' @param cutpoint_grid_size the number of cutoff points in XBART. The datatype is integer.
#' @param max_depth the maximum depth of tree allowed. The datatype is integer.
#' @param print_progress whether print progress bar or not. The datatype is boolean.
#' @param random_seed the random seed of the MCMC sampler. The datatype is integer.
#'
#' @returns A list of samples from the posterior distribution.
#' @noRd
outcome_model <- function(Y,
                          X_continuous,
                          X_categorical,
                          ni_vector, # the number of observations of each subjects.
                          group_info, # identification of observations belong to which subject.
                          num_warmup,
                          num_samples,
                          model_type,
                          thin,
                          L,
                          alpha,
                          beta,
                          leaf_model_scale,
                          cutpoint_grid_size,
                          max_depth = 10,
                          print_progress = TRUE,
                          random_seed = 100) {
  # set random seed ---------------------------------------------------------

  set.seed(random_seed)

  # define the number of subjects -------------------------------------------

  subject_size <- length(ni_vector)
  sample_size <- length(Y)

  indexL <- (Y == 0.0)
  indexM <- (0.0 < Y & Y < 1.0)
  indexR <- (Y == 1.0)

  sample_size_L <- sum(indexL)
  sample_size_R <- sum(indexR)

  YM <- Y[indexM]
  ZM <- YM
  Z <- Y

  # initialize values -------------------------------------------------------

  U_value <- rnorm(n = subject_size)
  random_effects <- rep(U_value, times = ni_vector)
  global_error_variance_value <- 1.0
  random_effects_variance_value <- 1.0
  fx_value <- rnorm(n = sample_size)

  if (model_type %in% c("Tobit-XBART","N-XBART")) {
    # combine covariate matrix
    # add the case that there is no categorical variable
    # for model validation purpose only
    if (is.null(X_categorical)) {
      Xall <- as.matrix(X_continuous)
    } else {
      X_categorical <- sapply(X_categorical, as.numeric)
      X_continuous <- as.matrix(X_continuous)
      Xall <- cbind(X_continuous,X_categorical)
    }

    # define feature types, 0 for numeric variable, 1 for categorical variable
    if (is.null(ncol(X_categorical))) {
      feature_types <- as.integer(rep(0,ncol(X_continuous)))
    } else {
      feature_types <- as.integer(c(rep(0,ncol(X_continuous)),
                                    rep(1,ncol(X_categorical))))
    }

    p_X <- ncol(Xall)
    variable_weights <- rep(1/p_X,p_X)

    # C++ dataset
    forest_dataset <- createForestDataset(Xall)

    # random number generator
    rng <- createCppRNG(random_seed)

    # sampling data structures
    forest_model_config <- createForestModelConfig(feature_types = feature_types,
                                                   num_trees = L,
                                                   num_features = p_X,
                                                   num_observations = sample_size,
                                                   variable_weights = variable_weights,
                                                   leaf_dimension = 1,
                                                   alpha = alpha,
                                                   beta = beta,
                                                   min_samples_leaf = 1,
                                                   max_depth = max_depth,
                                                   leaf_model_type = 0,
                                                   leaf_model_scale = leaf_model_scale,
                                                   variance_forest_shape = 1,
                                                   variance_forest_scale = 1,
                                                   cutpoint_grid_size = cutpoint_grid_size)

    global_model_config <- createGlobalModelConfig(global_error_variance = global_error_variance_value)

    forest_model <- createForestModel(forest_dataset = forest_dataset,
                                      forest_model_config = forest_model_config,
                                      global_model_config = global_model_config)

    # container of forest samples
    forest_samples <- createForestSamples(num_trees = L,
                                          leaf_dimension = 1,
                                          is_leaf_constant = TRUE,
                                          is_exponentiated = FALSE)

    # "active" forest state
    active_forest <- createForest(num_trees = L,
                                  leaf_dimension = 1,
                                  is_leaf_constant = TRUE,
                                  is_exponentiated = FALSE)

    # C++ outcome variable
    outcome <- createOutcome(outcome = Z)

    # Initialize the active forest and substract each root tree's predictions from outcome
    active_forest$prepare_for_sampler(dataset = forest_dataset,
                                      outcome = outcome,
                                      forest_model = forest_model,
                                      leaf_model_int = 0,
                                      leaf_value = 0.0)

    active_forest$adjust_residual(dataset = forest_dataset,
                                  outcome = outcome,
                                  forest_model = forest_model,
                                  requires_basis = FALSE,
                                  add = FALSE)

    rm(list = c("X_continuous","X_categorical","Xall")) # save memory
  } else if(model_type %in% c("Tobit-SBART","N-SBART")) {

    # 1. Prepare continuous matrix
    X_continuous <- as.matrix(X_continuous)
    num_continuous <- ncol(X_continuous)

    # 2. Create the continuous group IDs (starting at 1)
    group_cont <- 1:num_continuous

    # 3. Create the categorical matrix
    # Assuming X_categorical_df is your original dataframe of factors
    X_categorical <- model.matrix(~ . - 1, data = X_categorical)

    # 4. Extract the "assign" attribute and shift the IDs
    # assign_attr starts at 1. We shift it by num_continuous so it picks up
    # exactly where group_cont left off.
    assign_attr <- attr(X_categorical, "assign")
    group_cat <- assign_attr + num_continuous

    # 5. Combine matrices and group info
    Xall <- cbind(X_continuous, X_categorical)
    softforest_group_info <- c(group_cont, group_cat)

    Hypers_list <- Hypers(X = Xall,
                          Y = Z,
                          group = softforest_group_info,
                          alpha = alpha,
                          beta = beta,
                          k = 0.5/sqrt(leaf_model_scale * L),
                          sigma_hat = sqrt(global_error_variance_value),
                          num_tree = L,
                          normalize_Y = FALSE)

    Opts_list <- Opts(num_burn = num_warmup,
                      num_thin = thin,
                      num_save = num_samples/thin,
                      # do not print the chain's progress, pick a large number
                      num_print = 2147483647,
                      update_sigma_mu = FALSE,
                      update_s = FALSE,
                      update_alpha = FALSE,
                      update_beta = FALSE,
                      update_gamma = FALSE,
                      update_tau = FALSE,
                      update_tau_mean = FALSE,
                      update_sigma = FALSE,
                      cache_trees = TRUE
    )

    softforest <- MakeForest(hypers = Hypers_list,
                             opts = Opts_list,
                             warn = FALSE)

    rm(list = c("X_continuous","X_categorical")) # save memory
  } else if(model_type %in% c("Tobit-LH","N-LH")) {
    Xall <- generate_design_matrix(X_continuous = X_continuous,
                                   X_categorical = X_categorical)

    rm(list = c("X_continuous","X_categorical")) # save memory
  }


  # more initial value for horseshoe prior ----------------------------------

  if(model_type %in% c("Tobit-LH","N-LH")) {
    lambda2_value <- rep(1.0, ncol(Xall))
    tau2_value <- 1.0
    nu_value <- rep(1.0, ncol(Xall))
    xi_value <- 1.0
  }

  # define the output of the sampler ----------------------------------------

  MCMC_n <- num_warmup + num_samples

  # Calculate the number of samples after thinning
  num_samples_thinned <- num_samples / thin

  trees_outupt <- vector(mode = "list", length = num_samples_thinned)
  fx_output <- matrix(NA,
                      nrow = num_samples_thinned,
                      ncol = sample_size)
  global_error_variance_output <- numeric(num_samples_thinned)
  random_effects_variance_output <- numeric(num_samples_thinned)
  Z_output <- matrix(NA,
                     nrow = num_samples_thinned,
                     ncol = sample_size)
  random_effects_output <- matrix(NA,
                                  nrow = num_samples_thinned,
                                  ncol = sample_size)

  if(model_type %in% c("Tobit-LH","N-LH")) {
    beta_output <- matrix(NA,
                          nrow = num_samples_thinned,
                          ncol = ncol(Xall))
  }

  # begin the sampler -------------------------------------------------------

  if (print_progress) {
    pb <- progress_bar$new(total = MCMC_n,
                           format = "MCMC sampler [:bar] current: :current total: :total :percent eta: :eta",
                           width= 72)
  }

  # Initialize the Markov chain with num_warmup grow-from-root iterations
  gfr_flag <- TRUE
  for (i in 1:MCMC_n) {

    if (print_progress) {
      pb$tick()
    }

    # update latent variables -------------------------------------------------

    theta <- fx_value + random_effects

    # Tobit likelihood implies impute the latent variable Z
    if (model_type %in% c("Tobit-XBART","Tobit-SBART","Tobit-LH")) {
      if (sample_size_L > 0) {
        thetaL <- theta[indexL]
        ZL <- rtruncnorm(n = sample_size_L,
                         a = -Inf,
                         b = 0.0,
                         mean = thetaL,
                         sd = sqrt(global_error_variance_value))
        Z[indexL] <- ZL
      }

      if (sample_size_R > 0) {
        thetaR <- theta[indexR]
        ZR <- rtruncnorm(n = sample_size_R,
                         a = 1.0,
                         b = Inf,
                         mean = thetaR,
                         sd = sqrt(global_error_variance_value))
        Z[indexR] <- ZR
      }
    }

    # update random effects ---------------------------------------------------

    U_value <- update_U(Z = Z,
                        fixed_effects = fx_value,
                        ni_vector = ni_vector,
                        group_info = group_info,
                        global_error_variance = global_error_variance_value,
                        random_effects_variance = random_effects_variance_value)
    random_effects <- rep(U_value, times = ni_vector)

    if (model_type %in% c("Tobit-XBART","N-XBART")) {
      # propagate the newly sampled latent outcome to the BART model
      outcome$update_data(Z - random_effects)
      forest_model$propagate_residual_update(residual = outcome)
    }

    # update fixed effects ----------------------------------------------------


    if (model_type %in% c("Tobit-XBART","N-XBART")) {
      # use gfr for the first 100 iterations, decided on 2025-11-13
      if (i > 100) {
        gfr_flag <- FALSE
      }

      forest_model$sample_one_iteration(forest_dataset = forest_dataset,
                                        residual = outcome,
                                        forest_samples = forest_samples,
                                        active_forest = active_forest,
                                        rng = rng,
                                        forest_model_config = forest_model_config,
                                        global_model_config = global_model_config,
                                        keep_forest = TRUE,
                                        gfr = gfr_flag) # grow from root algorithm on/off

      fx_value <- forest_samples$predict_raw_single_forest(forest_dataset = forest_dataset,
                                                           forest_num = i - 1)
    } else if (model_type %in% c("Tobit-SBART","N-SBART")) {
      fx_value <-  softforest$do_gibbs(X = Xall,
                                       Y = Z - random_effects,
                                       X_test = Xall,
                                       i = 1)
    } else if (model_type %in% c("Tobit-LH", "N-LH")) {

      horseshoe_output <- horseshoe_iteration(y = Z - random_effects,
                                              X = Xall,
                                              sigma2 = global_error_variance_value,
                                              lambda2 = lambda2_value,
                                              tau2 = tau2_value,
                                              nu = nu_value,
                                              xi = xi_value)

      beta_value <- horseshoe_output$beta
      lambda2_value <- horseshoe_output$lambda2
      tau2_value <- horseshoe_output$tau2
      nu_value <- horseshoe_output$nu
      xi_value <- horseshoe_output$xi

      fx_value <- Xall %*% beta_value
    }
    fx_value <- as.numeric(fx_value)

    # update global error variance --------------------------------------------

    global_error_variance_value <- update_global_error_variance(Z = Z,
                                                                fixed_effects = fx_value,
                                                                random_effects = random_effects)
    if (model_type %in% c("Tobit-XBART", "N-XBART")) {
      global_model_config$update_global_error_variance(global_error_variance_value)
    } else if (model_type %in% c("Tobit-SBART", "N-SBART")) {
      softforest$set_sigma(x = sqrt(global_error_variance_value))
    }

    # update random effects variance ------------------------------------------

    random_effects_variance_value <- update_random_effects_variance(U_value = U_value,
                                                                    ni_vector = ni_vector)

    # Store samples after warmup period with thinning -------------------------
    if (i > num_warmup) {
      j <- i - num_warmup

      # Only store every 'thin'-th iteration
      if (j %% thin == 0) {
        sample_index <- j / thin  # Calculate the index in thinned arrays

        global_error_variance_output[sample_index] <- global_error_variance_value
        random_effects_variance_output[sample_index] <- random_effects_variance_value
        fx_output[sample_index,] <- fx_value
        Z_output[sample_index,] <- Z
        random_effects_output[sample_index,] <- random_effects

        if (model_type %in% c("Tobit-LH", "N-LH")) {
          beta_output[sample_index,] <- beta_value
        }
      }
    }
  }

  # 1. Create a base list with all the shared elements
  output <- list(
    global_error_variance_output   = global_error_variance_output,
    random_effects_variance_output = random_effects_variance_output,
    fx_output = fx_output,
    Z_output = Z_output,
    random_effects_output = random_effects_output,
    model_type = model_type,
    num_warmup = num_warmup,
    thin = thin
  )

  # 2. Append the specific element based on the model type
  if (model_type %in% c("Tobit-XBART","N-XBART")) {
    output$forest_samples <- forest_samples
  } else if (model_type %in% c("Tobit-SBART","N-SBART")) {
    output$softforest <- softforest
  } else if (model_type %in% c("Tobit-LH","N-LH")) {
    output$beta_output <- beta_output
  }

  return(output)

}
