# app.R

# ---- Robust package setup ----

options(repos = c(CRAN = "https://packagemanager.posit.co/cran/latest"))

required_packages <- c(
  "shiny",
  "bslib",
  "survival",
  "ggplot2",
  "dplyr",
  "DT",
  "scales"
)

installed_packages <- rownames(installed.packages())
missing_packages <- setdiff(required_packages, installed_packages)

if (length(missing_packages) > 0) {
  install.packages(
    missing_packages,
    type = "binary",
    dependencies = TRUE
  )
}

library(shiny)
library(bslib)
library(survival)
library(ggplot2)
library(dplyr)
library(DT)
library(scales)

generate_survival_data <- function(n = 400, censor_rate = 0.3) {
  set.seed(789)

  data.frame(
    id = seq_len(n),
    treatment = sample(c("Treatment A", "Treatment B"), n, replace = TRUE),
    subgroup = sample(c("Subgroup 1", "Subgroup 2", "Subgroup 3"), n, replace = TRUE),
    age = round(rnorm(n, 60, 10)),
    sex = sample(c("Female", "Male"), n, replace = TRUE)
  ) |>
    mutate(
      # Adjusted rates to have more gradual decline and more events early
      # Treatment B has better survival (lower rate), Treatment A worse (higher rate)
      event_time = rweibull(
        n,
        shape = 1.2,  # Shape > 1 means more events early
        scale = ifelse(treatment == "Treatment B", 50, 40)
      ),
      censored = rbinom(n, 1, censor_rate),
      censor_time = ifelse(
        censored == 1,
        runif(n, 0.2, 0.95) * event_time,
        180
      ),
      time = pmin(event_time, censor_time, 150),
      time = round(time, 1),
      event = as.integer(event_time <= censor_time & event_time <= 150)
    ) |>
    select(id, treatment, subgroup, age, sex, time, event)
}

km_plot <- function(data) {
  fit <- survfit(Surv(time, event) ~ treatment, data = data)
  
  strata <- if (is.null(fit$strata)) {
    rep(unique(data$treatment)[1], length(fit$time))
  } else {
    sub("treatment=", "", rep(names(fit$strata), fit$strata))
  }

  km_data <- data.frame(
    time = fit$time,
    survival = fit$surv,
    upper = fit$upper,
    lower = fit$lower,
    strata = strata
  )

  p_value <- if (length(unique(data$treatment)) < 2) {
    NA_real_
  } else {
    test <- survdiff(Surv(time, event) ~ treatment, data = data)
    pchisq(test$chisq, length(test$n) - 1, lower.tail = FALSE)
  }

  ggplot(km_data, aes(x = time, colour = strata, fill = strata)) +
    geom_ribbon(aes(ymin = lower, ymax = upper), alpha = 0.2, colour = NA) +
    geom_step(aes(y = survival), linewidth = 1) +
    scale_colour_manual(
      values = c("Treatment A" = "#56B4E9", "Treatment B" = "#F0C419"),
      drop = FALSE
    ) +
    scale_fill_manual(
      values = c("Treatment A" = "#56B4E9", "Treatment B" = "#F0C419"),
      drop = FALSE
    ) +
    scale_x_continuous(breaks = seq(0, 150, 50), limits = c(0, 150)) +
    scale_y_continuous(limits = c(0, 1), labels = number_format(accuracy = 0.01)) +
    labs(
      x = "Time (months)",
      y = "Survival Probability",
      colour = "Treatment",
      fill = "Treatment"
    ) +
    annotate(
      "text",
      x = 5,
      y = 0.18,
      hjust = 0,
      label = ifelse(is.na(p_value), "", paste0("p = ", signif(p_value, 3)))
    ) +
    theme_minimal(base_size = 13) +
    theme(legend.position = "top")
}

risk_table_plot <- function(data) {
  fit <- survfit(Surv(time, event) ~ treatment, data = data)
  times <- seq(0, 150, 50)
  risk <- summary(fit, times = times, extend = TRUE)

  treatment <- if (is.null(risk$strata)) {
    rep(unique(data$treatment)[1], length(risk$time))
  } else {
    sub("treatment=", "", risk$strata)
  }

  risk_data <- data.frame(
    treatment = treatment,
    time = risk$time,
    n_risk = risk$n.risk
  ) |>
    mutate(treatment = factor(treatment, levels = c("Treatment B", "Treatment A")))

  label_data <- risk_data |> distinct(treatment)

  ggplot(risk_data, aes(time, treatment, label = n_risk)) +
    geom_text(size = 4, colour = "black") +
    geom_text(
      data = label_data,
      aes(x = -8, y = treatment, label = treatment, colour = treatment),
      inherit.aes = FALSE,
      hjust = 1,
      size = 3.2
    ) +
    scale_colour_manual(
      values = c("Treatment A" = "#56B4E9", "Treatment B" = "#F0C419"),
      drop = FALSE
    ) +
    scale_y_discrete(drop = TRUE) +
    scale_x_continuous(breaks = times, limits = c(-15, 150)) +
    coord_cartesian(clip = "off") +
    labs(x = "Time", y = NULL, title = "Number at risk") +
    theme_minimal(base_size = 12) +
    theme(
      legend.position = "none",
      axis.text.y = element_blank(),
      axis.ticks.y = element_blank(),
      plot.margin = margin(5.5, 5.5, 5.5, 70)
    )
}

hazard_ratio_summary <- function(data) {
  if (length(unique(data$treatment)) < 2) {
    return(tags$p("Select both treatment groups to estimate a hazard ratio.", style = "font-size: 0.9rem;"))
  }

  model <- coxph(Surv(time, event) ~ treatment, data = data)
  estimate <- summary(model)
  hr <- estimate$conf.int[1, "exp(coef)"]
  lower <- estimate$conf.int[1, "lower .95"]
  upper <- estimate$conf.int[1, "upper .95"]
  p_value <- estimate$coefficients[1, "Pr(>|z|)"]

  tagList(
    h2(paste0("HR: ", round(hr, 2)), class = "text-center", style = "margin-bottom: 0.05rem; color: #56B4E9;"),
    p(paste0("95% CI: [", round(lower, 2), ", ", round(upper, 2), "]"), class = "text-center", style = "font-size: 0.9rem; margin-bottom: 0.05rem;"),
    p(paste0("P-value: ", signif(p_value, 3)), class = "text-center", style = "font-size: 0.9rem; margin-bottom: 0.05rem;"),
    hr(style = "margin-top: 0.05rem; margin-bottom: 0.05rem;"),
    p(paste0("Total subjects: ", nrow(data)), class = "text-center text-muted", style = "font-size: 0.85rem; margin-bottom: 0.02rem;"),
    p(
      paste0("Events: ", sum(data$event), " (", round(mean(data$event) * 100, 1), "%)"),
      class = "text-center text-muted",
      style = "font-size: 0.85rem; margin-bottom: 0;"
    )
  )
}

ui <- page_navbar(
  title = "Clinical Trial Survival Analysis Dashboard",
  sidebar = sidebar(
    width = 300,
    selectInput(
      "treatment",
      "Treatment Group:",
      choices = c("All", "Treatment A", "Treatment B")
    ),
    selectInput(
      "subgroup",
      "Subgroup:",
      choices = c("All", "Subgroup 1", "Subgroup 2", "Subgroup 3")
    ),
    sliderInput(
      "censor_rate",
      "Censoring Rate:",
      min = 0,
      max = 0.5,
      value = 0.3,
      step = 0.05
    ),
    downloadButton("download_data", "Download Dataset", class = "btn-primary")
  ),
  nav_panel(
    title = "Analysis",
    card(
      card_header("Kaplan-Meier Survival Curve"),
      plotOutput("km_plot", height = "420px")
    ),
    layout_columns(
      col_widths = c(6, 6),
      card(
        card_header("Number at Risk Table"),
        plotOutput("risk_table", height = "260px")
      ),
      card(
        card_header("Hazard Ratio Summary"),
        uiOutput("hr_summary")
      )
    )
  ),
  nav_panel(
    title = "Dataset",
    card(
      card_header("Survival Data Table"),
      DTOutput("data_table")
    )
  )
)

server <- function(input, output, session) {
  survival_data <- reactive({
    generate_survival_data(censor_rate = input$censor_rate)
  })

  filtered_data <- reactive({
    data <- survival_data()

    if (input$treatment != "All") {
      data <- data |> filter(treatment == input$treatment)
    }

    if (input$subgroup != "All") {
      data <- data |> filter(subgroup == input$subgroup)
    }

    data
  })

  output$km_plot <- renderPlot({
    km_plot(filtered_data())
  })

  output$risk_table <- renderPlot({
    risk_table_plot(filtered_data())
  })

  output$hr_summary <- renderUI({
    hazard_ratio_summary(filtered_data())
  })

  output$data_table <- renderDT({
    datatable(
      filtered_data(),
      filter = "top",
      options = list(
        pageLength = 25,
        searching = TRUE,
        autoWidth = TRUE,
        lengthChange = FALSE,  # Removes the "Show entries" dropdown
        dom = 'frtip'  # Customizes table controls (f=filter, r=processing, t=table, i=info, p=pagination)
      ),
      rownames = FALSE
    )
  })

  output$download_data <- downloadHandler(
    filename = function() {
      paste0("survival_data_", Sys.Date(), ".csv")
    },
    content = function(file) {
      write.csv(filtered_data(), file, row.names = FALSE)
    }
  )
}

shinyApp(ui, server)