#!/usr/bin/env Rscript

# Paper-grade single-panel figures supporting the conclusion that fixed-landmark
# horizon evaluation favors all_future over next_token.
#
# Outputs are written as separate panel files. This script intentionally does not
# combine panels with plot_grid().

suppressPackageStartupMessages({
  library(cowplot)
  library(dplyr)
  library(ggplot2)
  library(jsonlite)
  library(readr)
  library(stringr)
  library(tibble)
  library(tidyr)
})

root_dir <- "."
runs_dir <- file.path(root_dir, "runs")
out_dir <- file.path(root_dir, "figures_next_token_to_all_future_absolute_smoking_alcohol_bmi")
dir.create(out_dir, showWarnings = FALSE, recursive = TRUE)

required_time_mode <- "absolute"
required_extra_info_types <- c(11L, 66L, 67L)
required_extra_info_signature <- paste(sort(required_extra_info_types), collapse = ",")

theme_set(
  theme_cowplot(font_size = 9) +
    theme(
      plot.background = element_rect(fill = "white", color = NA),
      panel.background = element_rect(fill = "white", color = NA),
      legend.background = element_rect(fill = "white", color = NA),
      legend.key = element_rect(fill = "white", color = NA)
    )
)

target_cols <- c(
  "next_token" = "#B54A3A",
  "all_future" = "#2C7FB8"
)

dist_shapes <- c(
  "exponential" = 16,
  "weibull" = 17,
  "mixed" = 15
)

read_run_config <- function(run_path) {
  cfg_path <- file.path(run_path, "train_config.json")
  if (!file.exists(cfg_path)) return(NULL)
  cfg <- jsonlite::read_json(cfg_path, simplifyVector = TRUE)
  extra_info_types <- cfg$extra_info_types %||% integer(0)
  extra_info_signature <- paste(sort(as.integer(extra_info_types)), collapse = ",")
  tibble(
    run = basename(run_path),
    model_target_mode = as.character(cfg$model_target_mode %||% NA_character_),
    target_mode = as.character(cfg$target_mode %||% NA_character_),
    dist_mode = as.character(cfg$dist_mode %||% NA_character_),
    time_mode = as.character(cfg$time_mode %||% NA_character_),
    readout_name = as.character(cfg$readout_name %||% NA_character_),
    attn_mask_mode = as.character(cfg$attn_mask_mode %||% NA_character_),
    extra_info_signature = extra_info_signature
  )
}

`%||%` <- function(x, y) {
  if (is.null(x) || length(x) == 0) y else x
}

load_one_result <- function(run_path, file_name, eval_family) {
  cfg <- read_run_config(run_path)
  if (is.null(cfg)) return(NULL)
  fp <- file.path(run_path, file_name)
  if (!file.exists(fp)) return(NULL)

  df <- suppressMessages(readr::read_csv(fp, show_col_types = FALSE))
  if (!("auc" %in% names(df)) || nrow(df) == 0) return(NULL)

  out <- df %>%
    mutate(
      run = basename(run_path),
      eval_family = eval_family,
      auc = as.numeric(auc)
    ) %>%
    left_join(cfg, by = "run", suffix = c("", "_cfg"))

  coalesce_joined <- function(data, col) {
    cfg_col <- paste0(col, "_cfg")
    if (col %in% names(data) && cfg_col %in% names(data)) {
      dplyr::coalesce(data[[col]], data[[cfg_col]])
    } else if (col %in% names(data)) {
      data[[col]]
    } else if (cfg_col %in% names(data)) {
      data[[cfg_col]]
    } else {
      rep(NA_character_, nrow(data))
    }
  }

  for (col in c("model_target_mode", "target_mode", "dist_mode", "time_mode", "readout_name", "attn_mask_mode")) {
    out[[col]] <- coalesce_joined(out, col)
  }

  out %>%
    select(-any_of(c(
      "model_target_mode_cfg", "target_mode_cfg", "dist_mode_cfg",
      "time_mode_cfg", "readout_name_cfg", "attn_mask_mode_cfg"
    )))
}

run_paths <- list.dirs(runs_dir, recursive = FALSE, full.names = TRUE)

landmark_auc <- bind_rows(lapply(
  run_paths,
  load_one_result,
  file_name = "df_auc_landmark.csv",
  eval_family = "Fixed landmark + horizon"
)) %>%
  filter(time_mode == "absolute")

token_auc <- bind_rows(lapply(
  run_paths,
  load_one_result,
  file_name = "df_both.csv",
  eval_family = "Delphi2M-style token"
)) %>%
  filter(time_mode == "absolute")

if (nrow(landmark_auc) == 0) {
  stop("No landmark AUC files found under runs/*/df_auc_landmark.csv")
}
if (nrow(token_auc) == 0) {
  stop("No token AUC files found under runs/*/df_both.csv")
}

landmark_auc <- landmark_auc %>%
  filter(
    time_mode == required_time_mode,
    extra_info_signature == required_extra_info_signature
  )

token_auc <- token_auc %>%
  filter(
    time_mode == required_time_mode,
    extra_info_signature == required_extra_info_signature
  )

if (nrow(landmark_auc) == 0 || nrow(token_auc) == 0) {
  stop(
    "No AUC rows remain after filtering for time_mode='",
    required_time_mode,
    "' and extra_info_types='",
    required_extra_info_signature,
    "'."
  )
}

message(
  "Using runs with time_mode='", required_time_mode,
  "' and extra_info_types='", required_extra_info_signature, "':"
)
print(sort(unique(landmark_auc$run)))

classify_endpoint <- function(data) {
  data %>%
    mutate(
      endpoint_type = if_else(
        str_to_lower(as.character(label_code)) == "death",
        "Death",
        "Non-death disease"
      ),
      endpoint_type = factor(endpoint_type, levels = c("Non-death disease", "Death"))
    )
}

landmark_auc <- classify_endpoint(landmark_auc)
token_auc <- classify_endpoint(token_auc)

landmark_auc_disease <- landmark_auc %>% filter(endpoint_type == "Non-death disease")
token_auc_disease <- token_auc %>% filter(endpoint_type == "Non-death disease")
landmark_auc_death <- landmark_auc %>% filter(endpoint_type == "Death")
token_auc_death <- token_auc %>% filter(endpoint_type == "Death")

if (nrow(landmark_auc_death) == 0 || nrow(token_auc_death) == 0) {
  warning("Death rows were not found in one or both AUC tables.")
}

auc_all <- bind_rows(
  landmark_auc_disease %>% mutate(horizon = as.numeric(horizon), offset = NA_real_),
  token_auc_disease %>% mutate(horizon = NA_real_, offset = as.numeric(offset))
) %>%
  mutate(
    model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future")),
    eval_family = factor(eval_family, levels = c("Delphi2M-style token", "Fixed landmark + horizon")),
    dist_mode = factor(dist_mode, levels = c("exponential", "weibull", "mixed")),
    model_label = recode(
      as.character(model_target_mode),
      "next_token" = "next-token objective",
      "all_future" = "all-future objective"
    )
  )

mean_ci <- function(x) {
  x <- x[is.finite(x)]
  n <- length(x)
  m <- mean(x)
  se <- sd(x) / sqrt(n)
  tibble(mean = m, ymin = m - 1.96 * se, ymax = m + 1.96 * se, n = n)
}

save_panel <- function(plot, name, width = 3.6, height = 3.0) {
  pdf_path <- file.path(out_dir, paste0(name, ".pdf"))
  png_path <- file.path(out_dir, paste0(name, ".png"))
  cowplot::save_plot(pdf_path, plot, base_width = width, base_height = height, bg = "white")
  cowplot::save_plot(png_path, plot, base_width = width, base_height = height, dpi = 600, bg = "white")
  message("Wrote: ", pdf_path)
  message("Wrote: ", png_path)
}

# Panel 1: run-level mean AUC under the clinically aligned landmark/horizon task.
# Death is excluded here and plotted separately below.
landmark_run <- landmark_auc_disease %>%
  mutate(model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future"))) %>%
  group_by(run, model_target_mode, dist_mode, time_mode, target_mode) %>%
  summarise(mean_auc = mean(auc, na.rm = TRUE), median_auc = median(auc, na.rm = TRUE), .groups = "drop")

landmark_summary <- landmark_run %>%
  group_by(model_target_mode) %>%
  summarise(mean_ci(mean_auc), .groups = "drop")

p1 <- ggplot(landmark_run, aes(x = model_target_mode, y = mean_auc)) +
  geom_point(
    aes(color = model_target_mode, shape = dist_mode),
    position = position_jitter(width = 0.09, height = 0, seed = 1),
    size = 2.2,
    alpha = 0.88
  ) +
  geom_errorbar(
    data = landmark_summary,
    aes(x = model_target_mode, y = mean, ymin = ymin, ymax = ymax, color = model_target_mode),
    width = 0.12,
    linewidth = 0.55,
    inherit.aes = FALSE
  ) +
  geom_point(
    data = landmark_summary,
    aes(x = model_target_mode, y = mean, color = model_target_mode),
    size = 3.4,
    inherit.aes = FALSE
  ) +
  scale_color_manual(values = target_cols, guide = "none") +
  scale_shape_manual(values = dist_shapes, na.translate = FALSE) +
  scale_x_discrete(labels = c("next_token", "all_future")) +
  coord_cartesian(ylim = c(0.58, 0.78)) +
  labs(
    x = NULL,
    y = "Mean AUC per run",
    shape = "Risk head",
    title = "Non-death landmark AUC (absolute time)"
  ) +
  theme(
    plot.title = element_text(face = "bold", size = 10),
    axis.text.x = element_text(size = 9),
    legend.position = c(0.72, 0.20),
    legend.background = element_blank()
  )

save_panel(p1, "panel_01_landmark_overall")

# Panel 2: landmark AUC by prediction horizon.
landmark_horizon_run <- landmark_auc_disease %>%
  mutate(
    horizon = as.numeric(horizon),
    model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future"))
  ) %>%
  group_by(run, model_target_mode, horizon) %>%
  summarise(mean_auc = mean(auc, na.rm = TRUE), .groups = "drop")

landmark_horizon_summary <- landmark_horizon_run %>%
  group_by(model_target_mode, horizon) %>%
  summarise(mean_ci(mean_auc), .groups = "drop")

p2 <- ggplot(landmark_horizon_run, aes(x = horizon, y = mean_auc, color = model_target_mode)) +
  geom_line(aes(group = run), alpha = 0.18, linewidth = 0.35) +
  geom_point(alpha = 0.32, size = 1.1) +
  geom_ribbon(
    data = landmark_horizon_summary,
    aes(x = horizon, y = mean, ymin = ymin, ymax = ymax, fill = model_target_mode, group = model_target_mode),
    alpha = 0.13,
    color = NA,
    inherit.aes = FALSE
  ) +
  geom_line(data = landmark_horizon_summary, aes(y = mean), linewidth = 0.85) +
  geom_point(data = landmark_horizon_summary, aes(y = mean), size = 2.0) +
  scale_color_manual(
    values = target_cols,
    labels = c("next_token", "all_future"),
    name = NULL
  ) +
  scale_fill_manual(values = target_cols, guide = "none") +
  scale_x_continuous(breaks = c(1, 5, 10)) +
  coord_cartesian(ylim = c(0.58, 0.78)) +
  labs(
    x = "Prediction horizon, years",
    y = "Mean AUC per run",
    title = "Non-death landmark AUC across horizons"
  ) +
  theme(
    plot.title = element_text(face = "bold", size = 10),
    legend.position = c(0.31, 0.20),
    legend.background = element_blank()
  )

save_panel(p2, "panel_02_landmark_by_horizon", width = 3.8, height = 3.0)

# Panel 3: Delphi2M-style token AUC by offset. This documents why the old
# evaluation can make next_token look competitive, especially near the event.
token_offset_run <- token_auc_disease %>%
  mutate(
    offset = as.numeric(offset),
    model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future"))
  ) %>%
  group_by(run, model_target_mode, offset) %>%
  summarise(mean_auc = mean(auc, na.rm = TRUE), .groups = "drop")

token_offset_summary <- token_offset_run %>%
  group_by(model_target_mode, offset) %>%
  summarise(mean_ci(mean_auc), .groups = "drop")

p3 <- ggplot(token_offset_run, aes(x = offset, y = mean_auc, color = model_target_mode)) +
  geom_line(aes(group = run), alpha = 0.18, linewidth = 0.35) +
  geom_point(alpha = 0.32, size = 1.1) +
  geom_ribbon(
    data = token_offset_summary,
    aes(x = offset, y = mean, ymin = ymin, ymax = ymax, fill = model_target_mode, group = model_target_mode),
    alpha = 0.13,
    color = NA,
    inherit.aes = FALSE
  ) +
  geom_line(data = token_offset_summary, aes(y = mean), linewidth = 0.85) +
  geom_point(data = token_offset_summary, aes(y = mean), size = 2.0) +
  scale_color_manual(
    values = target_cols,
    labels = c("next_token", "all_future"),
    name = NULL
  ) +
  scale_fill_manual(values = target_cols, guide = "none") +
  scale_x_continuous(breaks = c(0.1, 1, 5, 10), trans = "log10") +
  coord_cartesian(ylim = c(0.55, 0.82)) +
  labs(
    x = "Minimum offset before event, years",
    y = "Mean AUC per run",
    title = "Non-death token AUC by offset"
  ) +
  theme(
    plot.title = element_text(face = "bold", size = 10),
    legend.position = c(0.31, 0.20),
    legend.background = element_blank()
  )

save_panel(p3, "panel_03_token_auc_by_offset", width = 3.8, height = 3.0)

# Panel 4: within-run contrast between old token evaluation and landmark
# evaluation. Each run contributes one point per evaluation family.
run_eval_contrast <- auc_all %>%
  group_by(run, model_target_mode, dist_mode, eval_family) %>%
  summarise(mean_auc = mean(auc, na.rm = TRUE), .groups = "drop")

p4 <- ggplot(run_eval_contrast, aes(x = eval_family, y = mean_auc, color = model_target_mode)) +
  geom_line(aes(group = run), alpha = 0.34, linewidth = 0.45) +
  geom_point(aes(shape = dist_mode), size = 2.0, alpha = 0.84) +
  stat_summary(
    aes(group = model_target_mode),
    fun = mean,
    geom = "point",
    size = 3.3,
    shape = 18,
    position = position_dodge(width = 0.16)
  ) +
  scale_color_manual(
    values = target_cols,
    labels = c("next_token", "all_future"),
    name = NULL
  ) +
  scale_shape_manual(values = dist_shapes, na.translate = FALSE, name = "Risk head") +
  coord_cartesian(ylim = c(0.58, 0.78)) +
  labs(
    x = NULL,
    y = "Mean AUC per run",
    title = "Evaluation choice changes the conclusion (absolute time)"
  ) +
  theme(
    plot.title = element_text(face = "bold", size = 10),
    axis.text.x = element_text(angle = 18, hjust = 1),
    legend.position = "right"
  )

save_panel(p4, "panel_04_evaluation_contrast", width = 4.3, height = 3.1)

# Panel 5: disease-level distribution for the landmark task, pooled over
# horizons and runs. This shows the shift without hiding heterogeneity.
landmark_density <- landmark_auc_disease %>%
  mutate(model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future"))) %>%
  filter(is.finite(auc))

p5 <- ggplot(landmark_density, aes(x = auc, fill = model_target_mode, color = model_target_mode)) +
  geom_density(alpha = 0.20, linewidth = 0.65, adjust = 1.1) +
  geom_vline(
    data = landmark_density %>%
      group_by(model_target_mode) %>%
      summarise(mean_auc = mean(auc), .groups = "drop"),
    aes(xintercept = mean_auc, color = model_target_mode),
    linewidth = 0.75,
    linetype = "22"
  ) +
  scale_color_manual(values = target_cols, labels = c("next_token", "all_future"), name = NULL) +
  scale_fill_manual(values = target_cols, labels = c("next_token", "all_future"), name = NULL) +
  coord_cartesian(xlim = c(0.35, 1.0)) +
  labs(
    x = "AUC",
    y = "Density",
    title = "Non-death landmark AUC distribution"
  ) +
  theme(
    plot.title = element_text(face = "bold", size = 10),
    legend.position = c(0.24, 0.82),
    legend.background = element_blank()
  )

save_panel(p5, "panel_05_landmark_auc_distribution", width = 3.8, height = 3.0)

# Panel 6: death-only fixed landmark + horizon AUC. Death has one endpoint token,
# so each line is a run trajectory across horizons.
death_landmark_run <- landmark_auc_death %>%
  mutate(
    horizon = as.numeric(horizon),
    model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future")),
    dist_mode = factor(dist_mode, levels = c("exponential", "weibull", "mixed"))
  ) %>%
  group_by(run, model_target_mode, dist_mode, horizon) %>%
  summarise(mean_auc = mean(auc, na.rm = TRUE), .groups = "drop")

death_landmark_summary <- death_landmark_run %>%
  group_by(model_target_mode, horizon) %>%
  summarise(mean_ci(mean_auc), .groups = "drop")

p6 <- ggplot(death_landmark_run, aes(x = horizon, y = mean_auc, color = model_target_mode)) +
  geom_line(aes(group = run), alpha = 0.42, linewidth = 0.45) +
  geom_point(aes(shape = dist_mode), alpha = 0.9, size = 2.0) +
  geom_line(data = death_landmark_summary, aes(y = mean, group = model_target_mode), linewidth = 0.9) +
  geom_point(data = death_landmark_summary, aes(y = mean), size = 2.2) +
  scale_color_manual(values = target_cols, labels = c("next_token", "all_future"), name = NULL) +
  scale_shape_manual(values = dist_shapes, na.translate = FALSE, name = "Risk head") +
  scale_x_continuous(breaks = c(1, 5, 10)) +
  coord_cartesian(ylim = c(0.58, 0.95)) +
  labs(
    x = "Prediction horizon, years",
    y = "AUC",
    title = "Death-only landmark AUC"
  ) +
  theme(
    plot.title = element_text(face = "bold", size = 10),
    legend.position = "right"
  )

save_panel(p6, "panel_06_death_landmark_by_horizon", width = 3.9, height = 3.0)

# Panel 7: death-only Delphi2M-style token AUC by offset.
death_token_run <- token_auc_death %>%
  mutate(
    offset = as.numeric(offset),
    model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future")),
    dist_mode = factor(dist_mode, levels = c("exponential", "weibull", "mixed"))
  ) %>%
  group_by(run, model_target_mode, dist_mode, offset) %>%
  summarise(mean_auc = mean(auc, na.rm = TRUE), .groups = "drop")

death_token_summary <- death_token_run %>%
  group_by(model_target_mode, offset) %>%
  summarise(mean_ci(mean_auc), .groups = "drop")

p7 <- ggplot(death_token_run, aes(x = offset, y = mean_auc, color = model_target_mode)) +
  geom_line(aes(group = run), alpha = 0.42, linewidth = 0.45) +
  geom_point(aes(shape = dist_mode), alpha = 0.9, size = 2.0) +
  geom_line(data = death_token_summary, aes(y = mean, group = model_target_mode), linewidth = 0.9) +
  geom_point(data = death_token_summary, aes(y = mean), size = 2.2) +
  scale_color_manual(values = target_cols, labels = c("next_token", "all_future"), name = NULL) +
  scale_shape_manual(values = dist_shapes, na.translate = FALSE, name = "Risk head") +
  scale_x_continuous(breaks = c(0.1, 1, 5, 10), trans = "log10") +
  coord_cartesian(ylim = c(0.58, 0.95)) +
  labs(
    x = "Minimum offset before event, years",
    y = "AUC",
    title = "Death-only token AUC"
  ) +
  theme(
    plot.title = element_text(face = "bold", size = 10),
    legend.position = "right"
  )

save_panel(p7, "panel_07_death_token_auc_by_offset", width = 3.9, height = 3.0)

# Panel 8: death-only contrast between the two evaluation families.
death_eval_contrast <- bind_rows(
  landmark_auc_death %>% mutate(horizon = as.numeric(horizon), offset = NA_real_),
  token_auc_death %>% mutate(horizon = NA_real_, offset = as.numeric(offset))
) %>%
  mutate(
    model_target_mode = factor(model_target_mode, levels = c("next_token", "all_future")),
    eval_family = factor(eval_family, levels = c("Delphi2M-style token", "Fixed landmark + horizon")),
    dist_mode = factor(dist_mode, levels = c("exponential", "weibull", "mixed"))
  ) %>%
  group_by(run, model_target_mode, dist_mode, eval_family) %>%
  summarise(mean_auc = mean(auc, na.rm = TRUE), .groups = "drop")

p8 <- ggplot(death_eval_contrast, aes(x = eval_family, y = mean_auc, color = model_target_mode)) +
  geom_line(aes(group = run), alpha = 0.38, linewidth = 0.5) +
  geom_point(aes(shape = dist_mode), size = 2.2, alpha = 0.9) +
  stat_summary(
    aes(group = model_target_mode),
    fun = mean,
    geom = "point",
    size = 3.4,
    shape = 18,
    position = position_dodge(width = 0.16)
  ) +
  scale_color_manual(values = target_cols, labels = c("next_token", "all_future"), name = NULL) +
  scale_shape_manual(values = dist_shapes, na.translate = FALSE, name = "Risk head") +
  coord_cartesian(ylim = c(0.58, 0.95)) +
  labs(
    x = NULL,
    y = "Mean AUC per run",
    title = "Death endpoint evaluated separately"
  ) +
  theme(
    plot.title = element_text(face = "bold", size = 10),
    axis.text.x = element_text(angle = 18, hjust = 1),
    legend.position = "right"
  )

save_panel(p8, "panel_08_death_evaluation_contrast", width = 4.3, height = 3.1)

# Export the exact run-level summaries used by the figures.
readr::write_csv(landmark_run, file.path(out_dir, "landmark_run_summary.csv"))
readr::write_csv(token_offset_run, file.path(out_dir, "token_offset_run_summary.csv"))
readr::write_csv(run_eval_contrast, file.path(out_dir, "run_evaluation_contrast.csv"))
readr::write_csv(death_landmark_run, file.path(out_dir, "death_landmark_run_summary.csv"))
readr::write_csv(death_token_run, file.path(out_dir, "death_token_offset_run_summary.csv"))
readr::write_csv(death_eval_contrast, file.path(out_dir, "death_evaluation_contrast.csv"))

message("Done. Panels are in: ", normalizePath(out_dir, winslash = "/"))