## -----------------------------------------------------------------------------
#| label: setup
#| echo: false

library(NBvarsel)
library(ggplot2)
library(dplyr)
library(tidyr)
library(stringr)
library(patchwork)
library(rms)
library(gt)
library(glmnet)
library(pROC)
library(ggsci)


## -----------------------------------------------------------------------------
#| label: simple-data

set.seed(42)
n <- 1000

X1 <- rnorm(n, mean = 0, sd = 1)
X2 <- rbinom(n, size = 1, prob = 0.7)
X3 <- rnorm(n, mean = 0, sd = 1)
X4 <- rbinom(n, size = 1, prob = 0.5)

log_odds <- -3 + 3 * X1 + 2 * X2 + 0.1 * X3 - 0.05 * X4
prob <- plogis(log_odds)
Y <- rbinom(n, size = 1, prob = prob)

df <- data.frame(X1 = X1, X2 = X2, X3 = X3, X4 = X4, Y = Y)

harms <- c(X1 = 0.05, X2 = 0.025, X3 = 0.05, X4 = 0.00005)


## -----------------------------------------------------------------------------
#| label: simple-bw

model_full <- lrm(Y ~ X1 + X2 + X3 + X4, data = df)
bw_result <- fastbw(model_full)
sprintf("Backward elimination retains: %s", paste(bw_result$names.kept, collapse = ", "))


## -----------------------------------------------------------------------------
#| label: simple-exhaustive

exhaustive <- nb_varsel(
  data = df,
  outcome_var = "Y",
  include_interactions = FALSE,
  costs = harms,
  cv_folds = 5,
  mode = "exhaustive",
  thresholds = 0.5,
  allow_parallel = FALSE,
  permutation = TRUE,
  splines = FALSE,
  n_knots = 3,
  verbose = FALSE
)


## -----------------------------------------------------------------------------
#| label: tbl-train-subsets
#| tbl-cap: "Results for the simple illustration. Models are ranked by cost-adjusted Net Benefit."

tbl_train <- exhaustive$all_models

tbl_train |>
  select(Model, AUC, Brier, Total_Cost, Avg_Adj_Net_Benefit, Avg_Net_Benefit) |>
  arrange(-Avg_Adj_Net_Benefit) |>
  gt(id = "tb-train") |>
  fmt_number(
    columns = c("AUC", "Brier", "Avg_Adj_Net_Benefit", "Avg_Net_Benefit"),
    decimals = 3
  ) |>
  cols_label(
    Model = "Included predictors",
    AUC = "AUC",
    Brier = "Brier Score",
    Avg_Net_Benefit = "Avg. Net Benefit",
    Total_Cost = "Total Cost",
    Avg_Adj_Net_Benefit = "Avg. Adj. Net Benefit"
  )


## -----------------------------------------------------------------------------
#| label: fig-simple-illustration
#| fig-cap:
#|   - "Net benefit based variable importance for the simple illustration."
#|   - "All subset plot for the simple illustration."

VIF_plot(tbl_train)$plot +
  theme(axis.text.x = element_text(angle = 0, hjust = 1))

all_subset_plot(
  tbl_train,
  filter = 100,
  metric = "Avg_Adj_Net_Benefit",
  y_axis = "Adjusted Net Benefit"
)


## -----------------------------------------------------------------------------
#| label: case-data

# Clean up simple illustration objects to avoid namespace conflicts
rm(list = setdiff(ls(), lsf.str()))

set.seed(8156)
n_train <- 3000
n_test <- 1500
n_total <- n_train + n_test

# Clinical history (free)
patient_age <- round(rnorm(n_total, mean = 68, sd = 12))
prior_admission <- rbinom(n_total, 1, 0.30)
specialist_centre <- rbinom(n_total, 1, 0.45)
symptom_severity <- rbinom(n_total, 1, 0.40)

# Cardiac imaging (moderate cost)
ejection_fraction <- pmin(75, pmax(10, rnorm(n_total, 42, 14)))
wall_thickness <- pmin(1, pmax(0.05, rbeta(n_total, 2, 4)))
valve_abnormalities <- rpois(n_total, lambda = 0.8)
pericardial_effusion <- rbinom(n_total, 1, 0.12)
chamber_dilation <- rbinom(n_total, 1, 0.22)

# Blood biomarker (higher cost)
nt_probnp <- round(pmax(20, rlnorm(n_total, meanlog = 6.5, sdlog = 1.4)))

# True model — strong signal for key predictors
lp <- -16.0 +
  0.07 * patient_age +
  3.0 * prior_admission +
  -0.12 * ejection_fraction +
  10.0 * wall_thickness +
  1.5 * valve_abnormalities +
  2.5 * pericardial_effusion +
  1.8 * chamber_dilation +
  0.9 * log2(nt_probnp) +
  # Weak / null predictors:
  0.03 * specialist_centre +
  0.02 * symptom_severity

readmitted <- rbinom(n_total, 1, plogis(lp))

clinical_df <- data.frame(
  readmitted, patient_age, nt_probnp, prior_admission,
  specialist_centre, ejection_fraction, valve_abnormalities,
  pericardial_effusion, chamber_dilation, symptom_severity,
  wall_thickness
)

training_data <- clinical_df[1:n_train, ]
test_data <- clinical_df[(n_train + 1):n_total, ]

sprintf("Training prevalence: %.1f%%", 100 * mean(training_data$readmitted))
sprintf("Test prevalence: %.1f%%", 100 * mean(test_data$readmitted))


## -----------------------------------------------------------------------------
#| label: case-costs

prevalence <- mean(test_data$readmitted)

grouped_costs <- list(
  history = list(
    cost = 0,
    vars = c("patient_age", "prior_admission", "specialist_centre",
             "symptom_severity")
  ),
  imaging = list(
    cost = prevalence * 0.03,
    vars = c(
      "ejection_fraction", "wall_thickness",
      "valve_abnormalities", "pericardial_effusion", "chamber_dilation"
    )
  ),
  biomarker = list(
    cost = prevalence * 0.08,
    vars = c("nt_probnp")
  )
)


## -----------------------------------------------------------------------------
#| label: case-exhaustive

vars <- names(training_data)

start_time <- Sys.time()
exhaustive_results <- nb_varsel(
  data = training_data,
  outcome_var = "readmitted",
  costs = grouped_costs,
  thresholds = seq(0.01, 0.2, by = 0.01),
  include_interactions = FALSE,
  cv_folds = 5,
  mode = "exhaustive",
  allow_parallel = TRUE,
  permutation = TRUE,
  splines = FALSE,
  verbose = FALSE
)
end_time <- Sys.time()

sprintf("Exhaustive search took: %s", format(end_time - start_time))


## -----------------------------------------------------------------------------
#| label: tbl-case-best
#| tbl-cap: "Best model from exhaustive search"

exhaustive_results$best_model_stats |>
  select(Model, AUC, Brier, Total_Cost, Avg_Adj_Net_Benefit, Avg_Net_Benefit) |>
  gt() |>
  fmt_number(
    columns = c("AUC", "Brier", "Total_Cost", "Avg_Adj_Net_Benefit", "Avg_Net_Benefit"),
    decimals = 3
  ) |>
  cols_label(
    Model = "Included predictors",
    AUC = "AUC",
    Brier = "Brier Score",
    Avg_Net_Benefit = "Avg. Net Benefit",
    Total_Cost = "Total Cost",
    Avg_Adj_Net_Benefit = "Avg. Adj. Net Benefit"
  )


## -----------------------------------------------------------------------------
#| label: tbl-case-all
#| tbl-cap: "All models from exhaustive search"

all_models <- exhaustive_results$all_models
all_models$Rank_NB <- rank(-all_models$Avg_Net_Benefit, ties.method = "min")
all_models$Rank_adj_NB <- rank(-all_models$Avg_Adj_Net_Benefit, ties.method = "min")

tbl_data <- all_models |>
  select(
    Rank_NB, Rank_adj_NB, Model, AUC, Brier,
    Total_Cost, Avg_Adj_Net_Benefit, Avg_Net_Benefit
  )

tbl <- tbl_data |>
  gt() |>
  fmt_number(
    columns = c("AUC", "Brier", "Total_Cost", "Avg_Adj_Net_Benefit", "Avg_Net_Benefit"),
    decimals = 3
  ) |>
  cols_label(
    Rank_NB = "Rank (NB)",
    Rank_adj_NB = "Rank (Adj. NB)",
    Model = "Included predictors",
    AUC = "AUC",
    Brier = "Brier Score",
    Avg_Net_Benefit = "Avg. NB",
    Total_Cost = "Total Cost",
    Avg_Adj_Net_Benefit = "Avg. Adj. NB"
  )

if (knitr::is_html_output()) {
  tbl <- tbl |>
    opt_interactive(
      use_pagination = TRUE,
      page_size_default = 10,
      use_filters = TRUE,
      use_compact_mode = TRUE
    )
}
tbl


## -----------------------------------------------------------------------------
#| label: fig-case-vif
#| fig-width: 5.5
#| fig-height: 3.2
#| fig-cap: "Variable importance plot (exhaustive search)"

VIF_plot(
  exhaustive_results$all_models,
  filter = round(nrow(exhaustive_results$all_models) * 0.1)
)$plot


## -----------------------------------------------------------------------------
#| label: fig-case-subset
#| fig-width: 6.4
#| fig-height: 8
#| fig-cap:
#|   - "All subset plot (Net Benefit)"
#|   - "All subset plot (Adjusted Net Benefit)"

all_subset_plot(
  exhaustive_results$all_models,
  filter = 7,
  size_dot = 1
)

all_subset_plot(
  exhaustive_results$all_models,
  filter = 7,
  size_dot = 1,
  metric = "Avg_Adj_Net_Benefit",
  y_axis = "Adjusted Net Benefit"
)


## -----------------------------------------------------------------------------
#| label: case-groupwise

start_time <- Sys.time()
groupwise <- nb_varsel(
  data = training_data,
  outcome_var = "readmitted",
  costs = grouped_costs,
  thresholds = seq(0.01, 0.2, by = 0.01),
  include_interactions = FALSE,
  cv_folds = 5,
  mode = "groupwise",
  allow_parallel = TRUE,
  permutation = TRUE,
  splines = FALSE,
  group_size = 1,
  verbose = TRUE
)
end_time <- Sys.time()
sprintf("Groupwise search took: %s", format(end_time - start_time))


## -----------------------------------------------------------------------------
#| label: tbl-case-groupwise
#| tbl-cap: "Groupwise selection: best model"

groupwise$best_model_stats |>
  select(Model, AUC, Brier, Total_Cost, Avg_Adj_Net_Benefit, Avg_Net_Benefit) |>
  gt() |>
  fmt_number(
    columns = c("AUC", "Brier", "Total_Cost", "Avg_Adj_Net_Benefit", "Avg_Net_Benefit"),
    decimals = 3
  ) |>
  cols_label(
    Model = "Included predictors",
    AUC = "AUC",
    Brier = "Brier Score",
    Avg_Net_Benefit = "Avg. NB",
    Total_Cost = "Total Cost",
    Avg_Adj_Net_Benefit = "Avg. Adj. NB"
  )


## -----------------------------------------------------------------------------
#| label: fig-case-group-vif
#| fig-width: 5.5
#| fig-height: 3.2
#| fig-cap: "Variable importance (groupwise selection)"

VIF_plot(groupwise$all_models)$plot


## -----------------------------------------------------------------------------
#| label: comparison-models

preds <- setdiff(vars, "readmitted")

# --- Backward elimination ---
fmla_full <- reformulate(termlabels = preds, response = "readmitted")
model_full <- lrm(fmla_full, data = training_data)
bw <- fastbw(model_full)
model_bw <- lrm(
  reformulate(termlabels = bw$names.kept, response = "readmitted"),
  data = training_data
)
test_data$bw_pred <- plogis(predict(model_bw, newdata = test_data))

# --- LASSO ---
x_train <- as.matrix(training_data[, preds])
y_train <- training_data$readmitted
x_test <- as.matrix(test_data[, preds])

set.seed(4738)
cv_lasso <- cv.glmnet(x_train, y_train, alpha = 1, family = "binomial")
lasso_coefs <- coef(cv_lasso, s = "lambda.1se")
lasso_kept <- rownames(lasso_coefs)[as.vector(lasso_coefs != 0)]
lasso_kept <- setdiff(lasso_kept, "(Intercept)")
test_data$lasso_pred <- as.vector(
  plogis(predict(cv_lasso, newx = x_test, s = "lambda.1se"))
)

# --- NB models ---
preds_best_nb <- str_trim(
  str_split(
    all_models |> arrange(-Avg_Net_Benefit) |> slice(1) |> pull(Model),
    ","
  )[[1]]
)
preds_best_adj <- str_trim(
  str_split(exhaustive_results$best_model_stats$Model, ",")[[1]]
)

model_nb <- lrm(
  reformulate(preds_best_nb, "readmitted"),
  data = training_data
)
model_adj_nb <- lrm(
  reformulate(preds_best_adj, "readmitted"),
  data = training_data
)
test_data$nb_pred <- plogis(predict(model_nb, newdata = test_data))
test_data$adj_nb_pred <- plogis(predict(model_adj_nb, newdata = test_data))


## -----------------------------------------------------------------------------
#| label: selected-vars

sprintf("NB model retains: %s", paste(preds_best_nb, collapse = ", "))
sprintf("Adjusted NB model retains: %s", paste(preds_best_adj, collapse = ", "))
sprintf("Backward elimination retains: %s", paste(bw$names.kept, collapse = ", "))
sprintf("LASSO retains: %s", paste(lasso_kept, collapse = ", "))


## -----------------------------------------------------------------------------
#| label: tbl-discrimination
#| tbl-cap: "AUC on test data for each variable selection method"

auc_results <- data.frame(
  Method = c("NB Model", "Adjusted NB Model", "Backward elimination", "LASSO"),
  AUC = c(
    as.numeric(pROC::auc(test_data$readmitted, test_data$nb_pred)),
    as.numeric(pROC::auc(test_data$readmitted, test_data$adj_nb_pred)),
    as.numeric(pROC::auc(test_data$readmitted, test_data$bw_pred)),
    as.numeric(pROC::auc(test_data$readmitted, test_data$lasso_pred))
  ),
  n_predictors = c(
    length(preds_best_nb),
    length(preds_best_adj),
    length(bw$names.kept),
    length(lasso_kept)
  )
)

auc_results |>
  gt() |>
  fmt_number(columns = "AUC", decimals = 3) |>
  cols_label(n_predictors = "N predictors")


## -----------------------------------------------------------------------------
#| label: fig-dca
#| fig-width: 6
#| fig-height: 4
#| fig-cap:
#|   - "Decision curve analysis comparing net benefit across methods."
#|   - "Decision curve analysis comparing harm-adjusted net benefit."

thresholds <- seq(0.01, 0.30, by = 0.005)
n_test_obs <- nrow(test_data)
y_test <- test_data$readmitted
prev <- mean(y_test)

calc_nb <- function(probs, y, thresholds) {
  n <- length(y)
  vapply(thresholds, function(pt) {
    pred_pos <- probs >= pt
    tp <- sum(y == 1 & pred_pos)
    fp <- sum(y == 0 & pred_pos)
    (tp / n) - (fp / n) * (pt / (1 - pt))
  }, numeric(1))
}

nb_all <- prev - (1 - prev) * (thresholds / (1 - thresholds))

dca_df <- bind_rows(
  data.frame(
    threshold = thresholds,
    net_benefit = calc_nb(test_data$nb_pred, y_test, thresholds),
    label = "NB Model"
  ),
  data.frame(
    threshold = thresholds,
    net_benefit = calc_nb(test_data$adj_nb_pred, y_test, thresholds),
    label = "Adjusted NB Model"
  ),
  data.frame(
    threshold = thresholds,
    net_benefit = calc_nb(test_data$bw_pred, y_test, thresholds),
    label = "Backward elimination"
  ),
  data.frame(
    threshold = thresholds,
    net_benefit = calc_nb(test_data$lasso_pred, y_test, thresholds),
    label = "LASSO"
  ),
  data.frame(
    threshold = thresholds,
    net_benefit = nb_all,
    label = "Treat all"
  ),
  data.frame(
    threshold = thresholds,
    net_benefit = 0,
    label = "Treat none"
  )
)

# Compute costs per method
cost_nb <- NBvarsel:::calculate_model_cost(preds_best_nb, grouped_costs)
cost_adj <- NBvarsel:::calculate_model_cost(preds_best_adj, grouped_costs)
cost_bw <- NBvarsel:::calculate_model_cost(bw$names.kept, grouped_costs)
cost_lasso <- NBvarsel:::calculate_model_cost(lasso_kept, grouped_costs)

dca_df <- dca_df |>
  mutate(
    cost = case_when(
      label == "NB Model" ~ cost_nb,
      label == "Adjusted NB Model" ~ cost_adj,
      label == "Backward elimination" ~ cost_bw,
      label == "LASSO" ~ cost_lasso,
      TRUE ~ 0
    ),
    adj_net_benefit = net_benefit - cost
  )

model_labels <- c(
  "NB Model", "Adjusted NB Model", "Backward elimination",
  "LASSO", "Treat all", "Treat none"
)
dca_df$label <- factor(dca_df$label, levels = model_labels)

# Raw NB
ggplot(dca_df, aes(x = threshold, y = net_benefit, color = label, linetype = label)) +
  geom_line(linewidth = 0.7) +
  theme_classic(base_size = 12) +
  labs(x = "Decision Threshold", y = "Net Benefit", color = "Strategy",
       linetype = "Strategy") +
  coord_cartesian(xlim = c(0, 0.30), ylim = c(-0.01, max(dca_df$net_benefit) * 1.05)) +
  scale_color_manual(values = c(
    "NB Model" = "#0072B5", "Adjusted NB Model" = "#BC3C29",
    "Backward elimination" = "#E18727", "LASSO" = "#20854E",
    "Treat all" = "black", "Treat none" = "black"
  )) +
  scale_linetype_manual(values = c(
    "NB Model" = "solid", "Adjusted NB Model" = "solid",
    "Backward elimination" = "solid", "LASSO" = "solid",
    "Treat all" = "dashed", "Treat none" = "dotted"
  ))

# Adjusted NB
ggplot(dca_df, aes(x = threshold, y = adj_net_benefit, color = label, linetype = label)) +
  geom_line(linewidth = 0.7) +
  theme_classic(base_size = 12) +
  labs(
    x = "Decision Threshold",
    y = "Harm-adjusted Net Benefit",
    color = "Strategy",
    linetype = "Strategy"
  ) +
  coord_cartesian(xlim = c(0, 0.30), ylim = c(-0.01, max(dca_df$adj_net_benefit) * 1.05)) +
  scale_color_manual(values = c(
    "NB Model" = "#0072B5", "Adjusted NB Model" = "#BC3C29",
    "Backward elimination" = "#E18727", "LASSO" = "#20854E",
    "Treat all" = "black", "Treat none" = "black"
  )) +
  scale_linetype_manual(values = c(
    "NB Model" = "solid", "Adjusted NB Model" = "solid",
    "Backward elimination" = "solid", "LASSO" = "solid",
    "Treat all" = "dashed", "Treat none" = "dotted"
  ))


## -----------------------------------------------------------------------------
#| label: adnex-costs
#| eval: false

# vars <- c(
#   "malignant", "age", "ca125", "family_history", "locules_gt_10",
#   "oncology_center", "max_diam_lesion", "papillary_count",
#   "acoustic_shadows", "ascites", "ireg_walls", "bilateral",
#   "color_score", "pain", "max_diam_solid", "papillary_presence",
#   "prop_solid"
# )
# 
# prevalence <- mean(test_data$malignant)
# 
# grouped_costs <- list(
#   history = list(
#     cost = 0,
#     vars = c("age", "family_history", "oncology_center", "pain")
#   ),
#   US = list(
#     cost = prevalence * 0.02,
#     vars = c(
#       "max_diam_lesion", "prop_solid", "locules_gt_10",
#       "papillary_count", "acoustic_shadows", "ascites",
#       "bilateral", "ireg_walls", "papillary_presence",
#       "color_score", "max_diam_solid"
#     )
#   ),
#   blood = list(
#     cost = prevalence * 0.05,
#     vars = c("ca125")
#   )
# )


## -----------------------------------------------------------------------------
#| label: adnex-exhaustive-call
#| eval: false

# exhaustive_results <- nb_varsel(
#   data = training_data[, vars],
#   outcome_var = "malignant",
#   costs = grouped_costs,
#   thresholds = seq(0.01, 0.2, by = 0.01),
#   include_interactions = FALSE,
#   cv_folds = 20,
#   mode = "exhaustive",
#   allow_parallel = TRUE,
#   permutation = TRUE,
#   splines = TRUE
# )


## -----------------------------------------------------------------------------
#| label: adnex-load-data

data(adnex_results)


## -----------------------------------------------------------------------------
#| label: tbl-adnex-best
#| tbl-cap: "Best model from the ADNEX exhaustive search (ranked by cost-adjusted Net Benefit)"

best <- attr(adnex_results, "best_model_stats")

best |>
  select(Model, n_Preds, AUC, Brier, Total_Cost,
         Avg_Adj_Net_Benefit, Avg_Net_Benefit) |>
  gt() |>
  fmt_number(
    columns = c("AUC", "Brier", "Total_Cost",
                "Avg_Adj_Net_Benefit", "Avg_Net_Benefit"),
    decimals = 3
  ) |>
  cols_label(
    Model = "Included predictors",
    n_Preds = "N",
    AUC = "AUC",
    Brier = "Brier Score",
    Avg_Net_Benefit = "Avg. NB",
    Total_Cost = "Total Cost",
    Avg_Adj_Net_Benefit = "Avg. Adj. NB"
  )


## -----------------------------------------------------------------------------
#| label: tbl-adnex-all
#| tbl-cap: "Top models per predictor count from the ADNEX exhaustive search"

adnex_results$Rank_adj_NB <- rank(
  -adnex_results$Avg_Adj_Net_Benefit, ties.method = "min"
)

tbl_data <- adnex_results |>
  select(Rank_adj_NB, Model, n_Preds, AUC, Brier,
         Total_Cost, Avg_Adj_Net_Benefit, Avg_Net_Benefit)

tbl <- tbl_data |>
  gt() |>
  fmt_number(
    columns = c("AUC", "Brier", "Total_Cost",
                "Avg_Adj_Net_Benefit", "Avg_Net_Benefit"),
    decimals = 3
  ) |>
  cols_label(
    Rank_adj_NB = "Rank (Adj. NB)",
    Model = "Included predictors",
    n_Preds = "N",
    AUC = "AUC",
    Brier = "Brier Score",
    Avg_Net_Benefit = "Avg. NB",
    Total_Cost = "Total Cost",
    Avg_Adj_Net_Benefit = "Avg. Adj. NB"
  )

if (knitr::is_html_output()) {
  tbl <- tbl |>
    opt_interactive(
      use_pagination = TRUE,
      page_size_default = 10,
      use_filters = TRUE,
      use_compact_mode = TRUE
    )
}
tbl


## -----------------------------------------------------------------------------
#| label: fig-adnex-vif
#| fig-width: 5.5
#| fig-height: 3.2
#| fig-cap: "Variable importance for the ADNEX case study. Bars show the average drop in Net Benefit when each predictor is permuted."

VIF_plot(adnex_results)$plot


## -----------------------------------------------------------------------------
#| label: fig-adnex-subset
#| fig-width: 6.4
#| fig-height: 8
#| fig-cap:
#|   - "All subset plot (Net Benefit) for the ADNEX case study."
#|   - "All subset plot (cost-adjusted Net Benefit) for the ADNEX case study."

all_subset_plot(
  adnex_results,
  filter = 7,
  size_dot = 1
)

all_subset_plot(
  adnex_results,
  filter = 7,
  size_dot = 1,
  metric = "Avg_Adj_Net_Benefit",
  y_axis = "Adjusted Net Benefit"
)