# =============================================================================
# FAIRNESS ASSESSMENT IN CLINICAL PREDICTION MODELS
# Subgroup performance, parity diagnostics, and cutoff sensitivity
#
# Developed by: Planimeter Kft.
# Version:      1.0.0
# Date:         2026-09-24
# Copyright:    (c) 2026 Planimeter Kft.
# SPDX-License-Identifier: GPL-3.0-or-later
#
# AI use statement:
#   Generative AI tools were used to support code development, review, and
#   validation. The analytical design, final code, interpretation, and
#   responsibility for the released version remain with Planimeter Kft.
#
# Purpose:
#   Reproduce the Planimeter fairness case study from the released compact
#   Framingham dataset. The analysis is methodological and is not a validated
#   clinical prediction tool or a determination of discrimination.
#
# Expected project structure:
#   project/
#   |-- fairness_analysis_planimeter_v1.0.0.R
#   |-- data/framingham_analysis_dataset.csv
#   |-- figures/
#   `-- output/
# =============================================================================

required_packages <- c(
  "dplyr", "readr", "ggplot2", "scales", "yardstick", "ranger",
  "DALEX", "fairmodels", "purrr", "tidyr"
)
missing_packages <- required_packages[
  !vapply(required_packages, requireNamespace, logical(1), quietly = TRUE)
]
if (length(missing_packages) > 0) {
  stop(paste("Missing required packages:", paste(missing_packages, collapse = ", ")))
}

library(dplyr)
library(readr)
library(ggplot2)
library(scales)
library(yardstick)
library(ranger)
library(DALEX)
library(fairmodels)
library(purrr)
library(tidyr)

dir.create("output", recursive = TRUE, showWarnings = FALSE)
dir.create("figures", recursive = TRUE, showWarnings = FALSE)

df <- read_csv(file.path("data", "framingham_analysis_dataset.csv"), show_col_types = FALSE) %>%
  mutate(
    outcome = factor(outcome, levels = c(0, 1)),
    sex = factor(sex, levels = c("Male", "Female")),
    age_group = factor(age_group, levels = c("Younger", "Older")),
    cursmoke = factor(cursmoke),
    diabetes = factor(diabetes)
  )

train <- df %>%
  filter(analysis_split == "Train") %>%
  select(-analysis_split, -source_row_id)

test <- df %>%
  filter(analysis_split == "Test") %>%
  select(-analysis_split, -source_row_id)

x_test <- test %>% select(-outcome)
y_test <- test$outcome
positive_class <- "1"
y_test_num <- ifelse(y_test == positive_class, 1, 0)

fit_rf <- ranger(
  outcome ~ ., data = train, probability = TRUE,
  num.trees = 500, seed = 20260610
)
fit_glm <- glm(outcome ~ ., data = train, family = binomial())

pred_rf <- predict(fit_rf, data = x_test)$predictions[, positive_class]
pred_glm <- predict(fit_glm, newdata = x_test, type = "response")

perf_df <- x_test %>%
  mutate(
    outcome = factor(y_test, levels = c("0", "1")),
    pred_rf = pred_rf,
    pred_glm = pred_glm,
    rf_class = factor(ifelse(pred_rf >= 0.5, "1", "0"), levels = c("0", "1")),
    glm_class = factor(ifelse(pred_glm >= 0.5, "1", "0"), levels = c("0", "1"))
  )

write_csv(
  perf_df %>% mutate(source_row_id = df$source_row_id[df$analysis_split == "Test"], .before = 1),
  file.path("output", "test_predictions.csv")
)

explainer_rf <- DALEX::explain(
  model = fit_rf, data = x_test, y = y_test_num,
  predict_function = function(model, newdata) {
    predict(model, data = newdata)$predictions[, positive_class]
  },
  label = "Random forest", verbose = FALSE
)
explainer_glm <- DALEX::explain(
  model = fit_glm, data = x_test, y = y_test_num,
  predict_function = function(model, newdata) {
    predict(model, newdata = newdata, type = "response")
  },
  label = "Logistic regression", verbose = FALSE
)

fair_sex <- fairness_check(
  explainer_glm, explainer_rf,
  protected = x_test$sex,
  privileged = "Male",
  cutoff = 0.5,
  colorize = FALSE
)

fairness_parity <- fair_sex$fairness_check_data %>%
  mutate(privileged_group = "Male", protected_group = "Female", cutoff = 0.5)
write_csv(fairness_parity, file.path("output", "fairness_parity_metrics.csv"))

pop_df <- perf_df %>%
  group_by(sex) %>%
  summarise(N = n(), Events = sum(outcome == "1"), EventRate = mean(outcome == "1"), .groups = "drop")
write_csv(pop_df, file.path("output", "population_by_sex.csv"))

auc_df <- bind_rows(
  perf_df %>% group_by(sex) %>% roc_auc(truth = outcome, pred_rf, event_level = "second") %>%
    transmute(sex, model = "Random forest", metric = "AUC", estimate = .estimate),
  perf_df %>% group_by(sex) %>% roc_auc(truth = outcome, pred_glm, event_level = "second") %>%
    transmute(sex, model = "Logistic regression", metric = "AUC", estimate = .estimate)
)

metric_by_sex <- function(data, model_name, class_col) {
  bind_rows(
    data %>% group_by(sex) %>% accuracy(truth = outcome, estimate = {{ class_col }}) %>% mutate(metric = "Accuracy"),
    data %>% group_by(sex) %>% sens(truth = outcome, estimate = {{ class_col }}, event_level = "second") %>% mutate(metric = "Sensitivity"),
    data %>% group_by(sex) %>% spec(truth = outcome, estimate = {{ class_col }}, event_level = "second") %>% mutate(metric = "Specificity"),
    data %>% group_by(sex) %>% ppv(truth = outcome, estimate = {{ class_col }}, event_level = "second") %>% mutate(metric = "PPV"),
    data %>% group_by(sex) %>% npv(truth = outcome, estimate = {{ class_col }}, event_level = "second") %>% mutate(metric = "NPV")
  ) %>% transmute(sex, model = model_name, metric, estimate = .estimate)
}

summary_metrics <- bind_rows(
  auc_df,
  metric_by_sex(perf_df, "Random forest", rf_class),
  metric_by_sex(perf_df, "Logistic regression", glm_class)
)
write_csv(summary_metrics, file.path("output", "model_metrics_by_sex.csv"))

safe_ratio <- function(a, b) {
  if (length(a) == 0 || length(b) == 0 || is.na(a) || is.na(b) || a == 0 || b == 0) return(NA_real_)
  a / b
}

calc_cutoff_metrics <- function(cutoff_value) {
  tmp <- perf_df %>% mutate(pred_class = factor(ifelse(pred_rf >= cutoff_value, "1", "0"), levels = c("0", "1")))
  overall_accuracy <- accuracy_vec(truth = tmp$outcome, estimate = tmp$pred_class)
  by_sex <- tmp %>% group_by(sex) %>% summarise(
    accuracy = accuracy_vec(outcome, pred_class),
    sensitivity = sens_vec(outcome, pred_class, event_level = "second"),
    specificity = spec_vec(outcome, pred_class, event_level = "second"),
    ppv = ppv_vec(outcome, pred_class, event_level = "second"),
    pos_rate = mean(pred_class == "1"), .groups = "drop"
  )
  male <- by_sex %>% filter(sex == "Male")
  female <- by_sex %>% filter(sex == "Female")
  ratios <- c(
    accuracy_equality = safe_ratio(female$accuracy, male$accuracy),
    equal_opportunity = safe_ratio(female$sensitivity, male$sensitivity),
    predictive_equality = safe_ratio(1 - female$specificity, 1 - male$specificity),
    predictive_parity = safe_ratio(female$ppv, male$ppv),
    statistical_parity = safe_ratio(female$pos_rate, male$pos_rate)
  )
  apd <- if (any(is.na(ratios)) || any(is.infinite(ratios)) || any(ratios <= 0)) NA_real_ else sum(abs(log(ratios)))
  tibble(
    cutoff = cutoff_value, accuracy = overall_accuracy, APD = apd,
    accuracy_equality = ratios["accuracy_equality"],
    equal_opportunity = ratios["equal_opportunity"],
    predictive_equality = ratios["predictive_equality"],
    predictive_parity = ratios["predictive_parity"],
    statistical_parity = ratios["statistical_parity"]
  )
}

tradeoff_df <- map_dfr(seq(0.10, 0.90, by = 0.05), calc_cutoff_metrics)
write_csv(tradeoff_df, file.path("output", "cutoff_tradeoff.csv"))

# ---- Publication figures ---------------------------------------------------
bg <- "#07111f"; gridc <- "#24384F"; txt <- "white"; subtxt <- "#A9BED6"
cyan <- "#55DDE0"; gold <- "#FFD166"; red <- "#FF6B6B"; green <- "#80ED99"

theme_dark <- function(base_size = 16) {
  theme_minimal(base_size = base_size) + theme(
    plot.background = element_rect(fill = bg, color = NA), panel.background = element_rect(fill = bg, color = NA),
    panel.grid.major = element_line(color = gridc, linewidth = 0.35), panel.grid.minor = element_blank(),
    plot.title = element_text(color = txt, face = "bold", size = base_size + 8),
    plot.subtitle = element_text(color = subtxt), axis.title = element_text(color = subtxt),
    axis.text = element_text(color = subtxt), legend.text = element_text(color = subtxt),
    legend.title = element_text(color = txt), legend.background = element_rect(fill = bg, color = NA)
  )
}

p1 <- ggplot(pop_df, aes(sex, EventRate, fill = sex)) +
  geom_col(width = 0.6) + geom_text(aes(label = paste0(percent(EventRate, accuracy = 0.1), "\nN = ", N)), vjust = -0.25, color = txt) +
  scale_y_continuous(labels = percent_format(), limits = c(0, max(pop_df$EventRate)*1.35)) + scale_fill_manual(values = c(cyan, gold)) +
  labs(title = "The outcome is not equally distributed", subtitle = "Observed cardiovascular event rate differs by sex", x = NULL, y = "Observed event rate") +
  theme_dark() + theme(legend.position = "none")
ggsave(file.path("figures", "01_event_rate_by_sex.png"), p1, width = 14, height = 8, dpi = 300, bg = bg)

metric_order <- c("Accuracy", "Sensitivity", "Specificity", "PPV", "NPV")
p2df <- summary_metrics %>% filter(metric %in% metric_order) %>% mutate(metric = factor(metric, levels = metric_order))
p2 <- ggplot(p2df, aes(metric, estimate, fill = sex)) + geom_col(position = position_dodge(0.7), width = 0.6) +
  facet_wrap(~model) + scale_y_continuous(labels = percent_format(), limits = c(0,1.05)) + scale_fill_manual(values = c(cyan,gold)) +
  labs(title = "A single performance metric is not enough", subtitle = "Subgroup metrics reveal different error profiles", x = NULL, y = NULL, fill = "Sex") + theme_dark()
ggsave(file.path("figures", "02_multimetric_performance_by_sex.png"), p2, width = 15, height = 8, dpi = 300, bg = bg)

fair_plot_df <- fairness_parity %>% mutate(metric_short = case_when(
  grepl("Accuracy equality", metric) ~ "Accuracy equality",
  grepl("Predictive parity", metric) ~ "Predictive parity",
  grepl("Predictive equality", metric) ~ "Predictive equality",
  grepl("Equal opportunity", metric) ~ "Equal opportunity",
  grepl("Statistical parity", metric) ~ "Statistical parity",
  TRUE ~ metric
))
p3 <- ggplot(fair_plot_df, aes(score, metric_short, color = model)) + geom_vline(xintercept = 1, linetype = "dashed", color = subtxt) +
  geom_segment(aes(x = 1, xend = score, yend = metric_short), linewidth = 2.8, alpha = 0.75) + geom_point(size = 3.5) +
  scale_color_manual(values = c("Random forest" = cyan, "Logistic regression" = gold)) +
  labs(title = "Parity diagnostics depend on the metric", subtitle = "A ratio of 1 represents parity with the illustrative privileged group", x = "Female / Male parity ratio", y = NULL, color = "Model") + theme_dark()
ggsave(file.path("figures", "03_fairness_parity_metrics.png"), p3, width = 15, height = 9, dpi = 300, bg = bg)

p4 <- tradeoff_df %>% filter(!is.na(APD)) %>% ggplot(aes(APD, accuracy, color = cutoff)) + geom_path(color = subtxt, linewidth = 1) + geom_point(size = 4) +
  geom_point(data = tradeoff_df %>% filter(abs(cutoff-0.5)<1e-9), color = gold, size = 7) +
  scale_y_continuous(labels = percent_format(accuracy=1)) +
  labs(title = "Threshold choice changes the fairness profile", subtitle = "Random-forest accuracy and Aggregate Parity Deviation across cutoffs", x = "Aggregate Parity Deviation (lower = closer to parity across selected ratios)", y = "Overall accuracy", color = "Cutoff") + theme_dark()
ggsave(file.path("figures", "04_accuracy_apd_by_cutoff.png"), p4, width = 15, height = 8, dpi = 300, bg = bg)

cat("Fairness analysis completed.\n")