# data generation ---------------------------------------------------------

df1 <- data_generation(random_seed = 100,
                       N = 100,
                       sigma = 0.2,
                       sigma_u = 0.1)

# draw samples from the posterior distribution ----------------------------

inference_output <- causal_inference_model(df = df1,
                                           y_name = "current_value",
                                           continuous_name = c("previous_value",
                                                               "confounder"),
                                           categorical_name = c("treatment"),
                                           treatment_name = "treatment",
                                           previous_status_name = "previous_value",
                                           subjectID_name = "subjectID",
                                           num_warmup = 2,
                                           num_samples = 2,
                                           model_type = "Tobit-XBART",
                                           thin = 1,
                                           L = 5,
                                           alpha = 0.95,
                                           beta = 1.25,
                                           leaf_model_scale = 0.3/5,
                                           cutpoint_grid_size = 100,
                                           max_depth = 10,
                                           credible_interval_level = 0.95,
                                           random_seed = 100,
                                           calculate_causal_estimand = FALSE,
                                           previous_status_grid_size = 2)

# calculate the causal estimand over a grid of previous status values -------

outcome_model_results <- inference_output$outcome_model_results

inference_results <- causal_estimand_inference(outcome_model_results = outcome_model_results,
                 df = inference_output$df,
                 continuous_name = inference_output$continuous_name,
                 categorical_name = inference_output$categorical_name,
                 treatment_name = inference_output$treatment_name,
                 treatment_value = factor("TREATMENT A",
                                          levels = levels(df1$treatment)),
                 previous_status_name = inference_output$previous_status_name,
                 credible_interval_level = 0.95,
                 num_of_grids = 2) # Example uses a small 2-point grid

# View the newly calculated closed-form causal estimands ------------------

# 1. Print results for the first grid point
cat("--- Results for Grid Point 1 ---\n")
print(inference_results[[1]]$mapo_summary)
print(inference_results[[1]]$mpdr_summary)

# 2. Print results for the second grid point
cat("\n--- Results for Grid Point 2 ---\n")
print(inference_results[[2]]$mapo_summary)
print(inference_results[[2]]$mpdr_summary)
