#' @title Bayesian Mixed Treatment-State Causal Model (MTSCM)
#'
#' @description
#' Fits a Bayesian Mixed Treatment-State Causal Model (MTSCM) tailored for longitudinal settings with irregular visits.
#' This model is specifically designed for bounded continuous outcomes with mass at both boundaries, such as the proportion of periodontal pockets exceeding 3 mm.
#'
#' @details
#' \strong{Outcome Model Description:}
#' The MTSCM utilizes a double-censored Tobit regression structure with subject-level random effects to capture within-subject correlation.
#' It uses tree-based ensemble priors (encompassing standard BART and Soft BART) to flexibly model complex, non-linear interactions without parametric restrictions.
#' Using a latent Gaussian variable \eqn{Z_{i,j+1}}, the data generating process for the outcome \eqn{Y_{i,j+1}} is formulated as:
#' \deqn{Z_{i,j+1} = f(A_{ij}, Y_{ij}, \mathbf{X}_{ij}) + U_i + \epsilon_{ij}, \quad \epsilon_{ij} \sim N(0, \sigma^2).}
#' \deqn{Y_{i,j+1} = \begin{cases} 0 & \text{if } Z_{i,j+1} \leq 0, \\ Z_{i,j+1} & \text{if } 0 < Z_{i,j+1} < 1, \\ 1 & \text{if } Z_{i,j+1} \geq 1. \end{cases}}
#'
#' \strong{Causal Estimand Description:}
#' The framework estimates causal effects using the G-computation formula to marginalize conditional expectations over the covariate distribution.
#'
#' \strong{Mixed Average Potential Outcome (MAPO):}
#' The MAPO represents the expected potential outcome across the population for a given treatment \eqn{a} and previous continuous disease state \eqn{y^{\star}}.
#' It is defined mathematically as:
#' \deqn{\theta(a,y^{\star}) = \frac{\sum_{i=1}^{N} \sum_{j=0}^{n_i - 1} \mathbb{E}(Y_{i,j+1} \mid A_{ij} = a, Y_{ij} = y^{\star}, \mathbf{X}_{ij} = \mathbf{x}_{ij})}{\sum_{i=1}^{N} n_{i}}.}
#'
#' \strong{Mixed Probability of Disease Resolution (MPDR):}
#' The MPDR estimates the population-level probability of achieving a zero disease burden under a specific treatment \eqn{a} and previous disease state \eqn{y^{\star}}.
#' It is computed as:
#' \deqn{\theta^{(0)}(a,y^{\star}) = \frac{\sum_{i=1}^{N} \sum_{j=0}^{n_i - 1} \mathbb{P}(Y_{i,j+1} = 0 \mid A_{ij} = a, Y_{ij} = y^{\star}, \mathbf{X}_{ij} = \mathbf{x}_{ij})}{\sum_{i=1}^{N} n_{i}}.}
#'
#' @param df the input dataframe. The datatype is \code{data.frame}.
#' @param y_name the name of the response variable. The datatype is \code{character}.
#' @param continuous_name the name of continuous predictors. The datatype is \code{character}.
#' @param categorical_name the name of categorical predictors. The datatype is \code{character}.
#' @param treatment_name the name of the treatment predictor. The datatype is \code{character}.
#' @param previous_status_name the name of variable that represents previous status of a subject. The datatype is \code{character}.
#' @param subjectID_name the name of variable that represents subjectID. The datatype is \code{character}.
#' @param num_warmup the number of warmup iterations. The datatype is \code{integer}.
#' @param num_samples the number of post-warmup iterations. The datatype is \code{integer}.
#' @param model_type the type of causal inference models. It must be one of \code{"Tobit-XBART"}, \code{"Tobit-SBART"}, \code{"Tobit-LH"}, \code{"N-XBART"}, \code{"N-SBART"}, or \code{"N-LH"}.
#' @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 \code{integer}.
#' @param L the number of trees. The datatype is \code{integer}.
#' @param alpha the tree prior parameters. \code{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 \code{double}.
#' @param beta the tree prior parameters. \code{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 \code{double}.
#' @param leaf_model_scale the prior variance on leaf mean equals to \code{leaf_model_scale/L}. The datatype is \code{double}.
#' @param cutpoint_grid_size the number of cutoff points in XBART. The datatype is \code{integer}.
#' @param max_depth the maximum depth of tree allowed. The datatype is \code{integer}.
#' @param credible_interval_level the nominal level of credible intervals of causal estimand. The datatype is \code{double}.
#' @param print_progress whether print progress bar or not. The datatype is \code{boolean}.
#' @param random_seed the random seed of the MCMC sampler. The datatype is \code{integer}.
#' @param calculate_causal_estimand calculate causal estimand or not. The datatype is \code{boolean}.
#' @param previous_status_grid_size the number of cutoff points of previous status for causal estimands. The datatype is \code{integer}.
#' @returns a list containing posterior inference of causal estimands.
#' @export
#'
#' @example inst/examples/main.r
causal_inference_model <- function(df,
                                   y_name,
                                   continuous_name,
                                   categorical_name,
                                   treatment_name,
                                   previous_status_name,
                                   subjectID_name,
                                   num_warmup,
                                   num_samples,
                                   model_type,
                                   thin = 1,
                                   L = 50,
                                   alpha = 0.95,
                                   beta = 1.25,
                                   leaf_model_scale = 0.3/50,
                                   cutpoint_grid_size = 100,
                                   max_depth = 10,
                                   credible_interval_level = 0.95,
                                   print_progress = TRUE,
                                   random_seed = 100,
                                   calculate_causal_estimand = FALSE,
                                   previous_status_grid_size = 100) {

  # Type checking for all parameters
  if (!is.data.frame(df)) {
    stop("df must be a data frame")
  }

  # avoid problem caused by tibble
  df <- as.data.frame(df)

  if (!is.character(y_name) || length(y_name) != 1) {
    stop("y_name must be a single character string")
  }

  if (!is.character(continuous_name)) {
    stop("continuous_name must be a character vector")
  }

  if (!is.character(categorical_name)) {
    stop("categorical_name must be a character vector")
  }

  if (!is.character(treatment_name) || length(treatment_name) != 1) {
    stop("treatment_name must be a single character string")
  }

  if (!is.character(previous_status_name) || length(previous_status_name) != 1) {
    stop("previous_status_name must be a single character string")
  }

  if (!is.character(subjectID_name) || length(subjectID_name) != 1) {
    stop("subjectID_name must be a single character string")
  }

  if (!is.numeric(num_warmup) || num_warmup <= 0 || num_warmup != as.integer(num_warmup)) {
    stop("num_warmup must be a positive integer")
  }

  if (!is.numeric(num_samples) || num_samples <= 0 || num_samples != as.integer(num_samples)) {
    stop("num_samples must be a positive integer")
  }

  if (!is.character(model_type) || ! model_type %in% c("Tobit-XBART",
                                                       "Tobit-SBART",
                                                       "Tobit-LH",
                                                       "N-XBART",
                                                       "N-SBART",
                                                       "N-LH")) {
    stop("model_type must be a string and one of Tobit-XBART, Tobit-SBART, Tobit-LH, N-XBART, N-SBART, N-LH")
  }

  if (!is.numeric(thin) || thin < 1 || thin != as.integer(thin)) {
    stop("thin must be an integer at least 1")
  }

  # Additional safe checks for thin and num_samples relationship
  if (num_samples <= thin) {
    stop("num_samples must be larger than thin")
  }

  if (num_samples %% thin != 0) {
    stop("num_samples must be divisible by thin (ratio must be an integer)")
  }

  if (!is.numeric(L) || L <= 0 || L != as.integer(L)) {
    stop("L must be a positive integer")
  }

  if (!is.numeric(alpha) || alpha <= 0) {
    stop("alpha must be a positive numeric value")
  }

  if (!is.numeric(beta) || beta <= 0) {
    stop("beta must be a positive numeric value")
  }

  if (!is.numeric(leaf_model_scale) || leaf_model_scale <= 0) {
    stop("leaf_model_scale must be a positive numeric value")
  }

  if (!is.numeric(cutpoint_grid_size) || cutpoint_grid_size <= 0 || cutpoint_grid_size != as.integer(cutpoint_grid_size)) {
    stop("cutpoint_grid_size must be a positive integer")
  }

  if (!is.numeric(max_depth) || max_depth <= 0 || max_depth != as.integer(max_depth)) {
    stop("max_depth must be a positive integer")
  }

  if (!is.numeric(credible_interval_level) || credible_interval_level <= 0 || credible_interval_level >= 1) {
    stop("credible_interval_level must be a numeric value between 0 and 1")
  }

  if (!is.logical(print_progress)) {
    stop("print_progress must be a logical value (TRUE or FALSE)")
  }

  if (!is.numeric(random_seed) || random_seed != as.integer(random_seed)) {
    stop("random_seed must be an integer")
  }

  if (!is.logical(calculate_causal_estimand)) {
    stop("calculate_causal_estimand must be a logical value (TRUE or FALSE)")
  }

  # Check if required columns exist in the dataframe
  required_columns <- c(y_name, continuous_name, categorical_name, treatment_name,
                        previous_status_name, subjectID_name)
  missing_columns <- setdiff(required_columns, names(df))

  if (length(missing_columns) > 0) {
    stop(paste("The following columns are missing from the dataframe:",
               paste(missing_columns, collapse = ", ")))
  }

  # Check if treatment variable is a factor
  if (!is.factor(df[[treatment_name]])) {
    stop("treatment_name must refer to a factor variable in the dataframe")
  }

  # Check each column to see if it is a factor
  is_factor_col <- sapply(df[, categorical_name, drop = FALSE], is.factor)
  if (!all(is_factor_col)) {
    # Find the names of the columns that failed the check
    non_factors <- names(df[, categorical_name, drop = FALSE])[!is_factor_col]

    # Stop execution and print a helpful error message
    stop(paste("Error: All categorical columns must be factors. The following columns are not:",
               paste(non_factors, collapse = ", ")))
  }

  Train_Data_Preparation_output <- Train_Data_Preparation(df = df,
                                                          y_name = y_name,
                                                          continuous_name = continuous_name,
                                                          categorical_name = categorical_name,
                                                          subjectID_name = subjectID_name)

  Y <- Train_Data_Preparation_output$Y
  X_continuous <- Train_Data_Preparation_output$X_continuous
  X_categorical <- Train_Data_Preparation_output$X_categorical

  ni_vector <- Train_Data_Preparation_output$ni_vector
  group_info <- Train_Data_Preparation_output$group_info

  outcome_model_results <- outcome_model(Y = Y,
                                         X_continuous = X_continuous,
                                         X_categorical = X_categorical,
                                         ni_vector = ni_vector,
                                         group_info = group_info,
                                         num_warmup = num_warmup,
                                         num_samples = num_samples,
                                         model_type = model_type,
                                         thin = thin,
                                         L = L,
                                         alpha = alpha,
                                         beta = beta,
                                         leaf_model_scale = leaf_model_scale,
                                         cutpoint_grid_size = cutpoint_grid_size,
                                         max_depth = max_depth,
                                         print_progress = print_progress,
                                         random_seed = random_seed)

  if (calculate_causal_estimand) {
    treatment_types <- levels(df[,treatment_name])
    causal_estimand_inference_list <- vector("list",length(treatment_types))
    names(causal_estimand_inference_list) <- treatment_types

    for (treatment_value in treatment_types) {
      if (print_progress) {
        message(paste("Calculate causal estimand for:", treatment_value))
      }
      causal_estimand_inference_list[[treatment_value]] <- causal_estimand_inference(outcome_model_results = outcome_model_results,
                                                                                     df = df,
                                                                                     continuous_name = continuous_name,
                                                                                     categorical_name = categorical_name,
                                                                                     treatment_name = treatment_name,
                                                                                     treatment_value = factor(treatment_value,
                                                                                                              levels = treatment_types),
                                                                                     previous_status_name = previous_status_name,
                                                                                     credible_interval_level = credible_interval_level,
                                                                                     num_of_grids = previous_status_grid_size)
    }

    output <- list(outcome_model_results = outcome_model_results,
                   causal_estimand = causal_estimand_inference_list,
                   df = df,
                   y_name = y_name,
                   continuous_name = continuous_name,
                   categorical_name = categorical_name,
                   treatment_name = treatment_name,
                   previous_status_name = previous_status_name,
                   subjectID_name = subjectID_name,
                   credible_interval_level = credible_interval_level)
  } else {

    output <- list(outcome_model_results = outcome_model_results,
                   df = df,
                   y_name = y_name,
                   continuous_name = continuous_name,
                   categorical_name = categorical_name,
                   treatment_name = treatment_name,
                   previous_status_name = previous_status_name,
                   subjectID_name = subjectID_name,
                   credible_interval_level = credible_interval_level)

  }

  return(output)

}
