sf1 <- survfit2(Surv(time, status) ~ 1, data = df_lung)
sf2 <- survfit2(Surv(time, status) ~ sex, data = df_lung)
sf3 <- survfit2(Surv(time, status) ~ sex + ph.ecog, data = df_lung)

test_that("tidy_survfit() works with survfit2()", {
  expect_error(
    list(sf1, sf2, sf3) %>% lapply(tidy_survfit),
    NA
  )

  expect_error(
    list(sf1, sf2, sf3) %>% lapply(tidy_survfit, type = "risk"),
    NA
  )

  expect_error(
    list(sf1, sf2, sf3) %>% lapply(tidy_survfit, type = "cumhaz"),
    NA
  )

  expect_error(
    list(sf1, sf2, sf3) %>% lapply(tidy_survfit, type = function(x) 1 - x),
    NA
  )

  expect_error(
    list(sf1, sf2, sf3) %>% lapply(tidy_survfit, times = 12:14),
    NA
  )

  expect_error(
    list(sf1, sf2, sf3) %>% lapply(tidy_survfit, times = 1),
    NA
  )

  expect_error(
    list(sf1, sf2, sf3) %>% lapply(tidy_survfit, times = 0),
    NA
  )

  # checking calculation of n.risk at timepoints
  times <- 0:4 * 10
  expect_equal(
    sf1 %>% tidy_survfit(times = times) %>% dplyr::pull(n.risk),
    lapply(times, function(x) sum(df_lung$time >= x)) %>% unlist()
  )

  df_nrisk_check <- sf2 %>% tidy_survfit(times = times) %>% dplyr::select(strata, time, n.risk)
  expect_equal(
    df_nrisk_check %>% dplyr::filter(strata == "Female") %>% dplyr::pull(n.risk),
    lapply(times, function(x) sum(df_lung$time >= x & df_lung$sex == "Female")) %>% unlist()
  )
  expect_equal(
    df_nrisk_check %>% dplyr::filter(strata == "Male") %>% dplyr::pull(n.risk),
    lapply(times, function(x) sum(df_lung$time >= x & df_lung$sex == "Male")) %>% unlist()
  )

  # checking calculation of event counts at timepoints
  expect_equal(
    sf1 %>% tidy_survfit(times = times) %>% dplyr::pull(cum.event),
    lapply(times, function(x) sum(df_lung$time <= x & df_lung$status == 2)) %>% unlist()
  )

  df_event_check <- sf2 %>% tidy_survfit(times = times) %>% dplyr::select(strata, time, cum.event)
  expect_equal(
    df_event_check %>% dplyr::filter(strata == "Female") %>% dplyr::pull(cum.event),
    lapply(times, function(x) sum(df_lung$time <= x & df_lung$status == 2 & df_lung$sex == "Female")) %>% unlist()
  )
  expect_equal(
    df_event_check %>% dplyr::filter(strata == "Male") %>% dplyr::pull(cum.event),
    lapply(times, function(x) sum(df_lung$time <= x & df_lung$status == 2 & df_lung$sex == "Male")) %>% unlist()
  )

  expect_error(sf2 %>% tidy_survfit(type = mtcars))

  expect_true(
    sf1 %>%
      tidy_survfit(type = "risk") %>%
      dplyr::mutate(..flag.. = conf.low <= conf.high) %>%
      dplyr::pull("..flag..") %>%
      all()
  )
})

test_that("tidy_survfit() throws appropriate errors", {
  expect_error(tidy_survfit(mtcars))
  expect_error(tidy_survfit(sf1, type = "not_a_type"))
})

test_that("tidy_survfit() works with survfit()", {
  sf1 <- survfit(Surv(time, status) ~ 1, data = df_lung)
  sf2 <- survfit(Surv(time, status) ~ sex, data = df_lung)
  sf3 <- survfit(Surv(time, status) ~ sex + ph.ecog, data = df_lung)

  expect_error(
    list(sf1, sf2, sf3) %>% lapply(tidy_survfit),
    NA
  )

  expect_error(
    list(sf1, sf2, sf3) %>% lapply(tidy_survfit, type = "risk"),
    NA
  )

  expect_error(
    list(sf1, sf2, sf3) %>% lapply(tidy_survfit, type = "cumhaz"),
    NA
  )

  expect_error(
    list(sf1, sf2, sf3) %>% lapply(tidy_survfit, type = function(x) 1 - x),
    NA
  )

  expect_error(
    list(sf1, sf2, sf3) %>% lapply(tidy_survfit, times = 12:14),
    NA
  )

  expect_error(
    list(sf1, sf2, sf3) %>% lapply(tidy_survfit, times = 1),
    NA
  )

  expect_error(
    list(sf1, sf2, sf3) %>% lapply(tidy_survfit, times = 0),
    NA
  )
})


test_that("tidy_survfit() works with multi-state models", {
  sfms1 <- survfit2(Surv(ttdeath, death_cr) ~ 1, data = tidycmprsk::trial)
  sfms2 <- survfit2(Surv(ttdeath, death_cr) ~ trt, data = tidycmprsk::trial)
  sfms3 <- survfit2(Surv(ttdeath, death_cr) ~ trt + grade, data = tidycmprsk::trial)

  cuminc1 <- tidycmprsk::cuminc(Surv(ttdeath, death_cr) ~ 1, data = tidycmprsk::trial)
  cuminc2 <- tidycmprsk::cuminc(Surv(ttdeath, death_cr) ~ trt, data = tidycmprsk::trial)
  cuminc3 <- tidycmprsk::cuminc(Surv(ttdeath, death_cr) ~ trt + grade, data = tidycmprsk::trial)

  expect_equal(
    sfms1 %>%
      tidy_survfit() %>%
      dplyr::select(dplyr::any_of(c("time", "outcome", "strata", "estimate"))) %>%
      dplyr::mutate(dplyr::across(dplyr::any_of(c("strata", "outcome")), ~as.character(.) %>% trimws())) %>%
      dplyr::arrange(dplyr::across(dplyr::any_of(c("time", "outcome", "strata")))),
    cuminc1 %>%
      tidy_cuminc() %>%
      dplyr::select(dplyr::any_of(c("time", "outcome", "strata", "estimate"))) %>%
      dplyr::mutate(dplyr::across(dplyr::any_of(c("strata", "outcome")), ~as.character(.) %>% trimws())) %>%
      dplyr::arrange(dplyr::across(dplyr::any_of(c("time", "outcome", "strata"))))
  )


  df_cmprsk_survfit <-
    sfms2 %>%
    tidy_survfit() %>%
    dplyr::select(dplyr::any_of(c("time", "outcome", "strata", "estimate"))) %>%
    dplyr::mutate(dplyr::across(dplyr::any_of(c("strata", "outcome")), ~as.character(.) %>% trimws())) %>%
    dplyr::arrange(dplyr::across(dplyr::any_of(c("outcome", "strata", "time"))))
  df_cmprsk_cuminc <-
    cuminc2 %>%
    tidy_cuminc() %>%
    dplyr::select(dplyr::any_of(c("time", "outcome", "strata", "estimate"))) %>%
    dplyr::mutate(dplyr::across(dplyr::any_of(c("strata", "outcome")), ~as.character(.) %>% trimws())) %>%
    dplyr::arrange(dplyr::across(dplyr::any_of(c("outcome", "strata", "time"))))

  expect_equal(
    df_cmprsk_cuminc %>%
      dplyr::select(c("outcome", "strata", "time", "estimate")),
    # adding in potential rows that were unobserved
    df_cmprsk_cuminc %>%
      dplyr::select(c("outcome", "strata", "time")) %>%
      dplyr::full_join(df_cmprsk_survfit, by = c("outcome", "strata", "time")) %>%
      dplyr::arrange(dplyr::across(dplyr::any_of(c("outcome", "strata", "time")))) %>%
      dplyr::group_by(dplyr::across(dplyr::any_of(c("outcome", "strata")))) %>%
      tidyr::fill("estimate", .direction = "down") %>%
      dplyr::ungroup()
  )

  df_sfms3 <-
    sfms3 %>%
    tidy_survfit() %>%
    dplyr::select(dplyr::any_of(c("time", "outcome", "strata", "estimate"))) %>%
    dplyr::mutate(dplyr::across(dplyr::any_of(c("strata", "outcome")), ~as.character(.) %>% trimws())) %>%
    dplyr::arrange(dplyr::across(dplyr::any_of(c("time", "outcome", "strata"))))
  df_cuminc3 <-
    cuminc3 %>%
    tidy_cuminc() %>%
    dplyr::select(dplyr::any_of(c("time", "outcome", "strata", "estimate"))) %>%
    dplyr::mutate(dplyr::across(dplyr::any_of(c("strata", "outcome")), ~as.character(.) %>% trimws())) %>%
    dplyr::arrange(dplyr::across(dplyr::any_of(c("time", "outcome", "strata"))))

  expect_equal(
    df_cuminc3 %>%
      dplyr::arrange(dplyr::across(dplyr::any_of(c("outcome", "strata", "time")))) %>%
      dplyr::select(c("outcome", "strata", "time", "estimate")),
    # adding in potential rows that were unobserved
    df_cuminc3 %>%
      dplyr::select(c("outcome", "strata", "time")) %>%
      dplyr::full_join(df_sfms3, by = c("outcome", "strata", "time")) %>%
      dplyr::group_by(dplyr::across(dplyr::any_of(c("outcome", "strata")))) %>%
      tidyr::fill("estimate", .direction = "down") %>%
      dplyr::ungroup() %>%
      dplyr::arrange(dplyr::across(dplyr::any_of(c("outcome", "strata", "time"))))
  )
})

test_that("tidy_survfit() messaging", {
  # expect a message about the start time when there are negative times and no start.time
  expect_snapshot(
    survfit(Surv(time - 500, status) ~ 1, df_lung) %>%
      tidy_survfit() %>%
      invisible()
  )
  # no longer see message when start.time specified
  expect_invisible(
    survfit(Surv(time - 500, status) ~ 1, df_lung, start.time = -500) %>%
      tidy_survfit() %>%
      invisible()
  )
})



test_that("tidy_survfit() handles custom transformation functions correctly with CI ordering", {

  # Test data: regular survival model
  sf_regular <- survfit2(Surv(time, status) ~ sex, data = df_lung)

  # Test data: multi-state model
  sfms <- survfit2(Surv(ttdeath, death_cr) ~ 1, data = tidycmprsk::trial)

  # Custom function that MAINTAINS monotonicity (decreasing -> decreasing)
  # Example: square the survival probability (still decreasing)
  maintain_monotonicity <- function(x) x^2

  # Custom function that FLIPS monotonicity (decreasing -> increasing)
  # Example: convert survival to risk (1-x)
  flip_monotonicity <- function(x) 1 - x

  # Test 1: Regular model with custom function that maintains monotonicity
  # Should NOT swap CIs (since monotonicity stays decreasing)
  result_maintain <- tidy_survfit(sf_regular, type = maintain_monotonicity)

  expect_true(
    result_maintain %>%
      dplyr::filter(!is.na(conf.low) & !is.na(conf.high)) %>%
      dplyr::mutate(ci_correct = conf.low <= conf.high) %>%
      dplyr::pull(ci_correct) %>%
      all()
  )

  # Test 2: Regular model with custom function that flips monotonicity
  # Should swap CIs (since monotonicity changes from decreasing to increasing)
  result_flip <- tidy_survfit(sf_regular, type = flip_monotonicity)

  expect_true(
    result_flip %>%
      dplyr::filter(!is.na(conf.low) & !is.na(conf.high)) %>%
      dplyr::mutate(ci_correct = conf.low <= conf.high) %>%
      dplyr::pull(ci_correct) %>%
      all()
  )

  # Test 3: Multi-state model with custom functions

  result_ms_maintain <- tidy_survfit(sfms, type = maintain_monotonicity)
  result_ms_flip <- tidy_survfit(sfms, type = flip_monotonicity)

  expect_true(
    result_ms_maintain %>%
      dplyr::filter(!is.na(conf.low) & !is.na(conf.high)) %>%
      dplyr::mutate(ci_correct = conf.low <= conf.high) %>%
      dplyr::pull(ci_correct) %>%
      all()
  )

  expect_true(
    result_ms_flip %>%
      dplyr::filter(!is.na(conf.low) & !is.na(conf.high)) %>%
      dplyr::mutate(ci_correct = conf.low <= conf.high) %>%
      dplyr::pull(ci_correct) %>%
      all()
  )

  # Test 4: Verify the transformations actually worked
  # Check that estimates were transformed as expected
  original_survival <- tidy_survfit(sf_regular, type = "survival")

  # Compare transformed estimates
  expect_equal(
    result_maintain$estimate,
    original_survival$estimate^2
  )

  expect_equal(
    result_flip$estimate,
    1 - original_survival$estimate
  )
})