# Tests for the plotting helpers (R/plotting.R). Every plot function consumes
# a plain summary data frame (not a live brms fit), so these are brms-free: we
# build synthetic inputs whose shape matches the documented producer output,
# then assert via ggplot2::ggplot_build() that the *plotted geometry encodes
# the input data correctly* and that each option flips the expected scale /
# layer. This is stronger than "returns a ggplot": it checks the mapping.

# --- helpers ---------------------------------------------------------------
built <- function(p) ggplot2::ggplot_build(p)

# Does any scale for `aes` use a log-10 transform?
has_log_scale <- function(p, aes) {
  any(vapply(p$scales$scales, function(s) {
    aes_ok <- any(s$aesthetics %in% aes)
    tr <- s$trans %||% s$transform
    aes_ok && !is.null(tr) && identical(tr$name, "log-10")
  }, logical(1)))
}
`%||%` <- function(a, b) if (is.null(a)) b else a

panel_axis_ranges <- function(p, axis = c("x", "y")) {
  axis <- match.arg(axis)
  bd <- built(p)
  rng_name <- paste0(axis, ".range")
  rng <- do.call(rbind, lapply(bd$layout$panel_params, `[[`, rng_name))
  out <- cbind(bd$layout$layout, as.data.frame(rng))
  names(out)[(ncol(out) - 1):ncol(out)] <- c("axis_min", "axis_max")
  out
}

# ============================ theme_tdt ====================================

test_that("theme_tdt returns a theme honouring base_size and bold axis titles", {
  th <- theme_tdt(base_size = 20)
  expect_s3_class(th, "theme")
  expect_equal(th$text$size, 20)
  expect_equal(th$axis.title$face, "bold")
  # Attaches cleanly to a plot and builds.
  p <- ggplot2::ggplot(mtcars, ggplot2::aes(mpg, hp)) +
    ggplot2::geom_point() + theme_tdt()
  expect_silent(built(p))
})

# ===================== plot_temperature_scenarios ==========================

test_that("plot_temperature_scenarios maps each scenario to its own facet and traces the temps", {
  scens <- make_temperature_scenarios(baseline = 20, spike_temp = 30,
                                      n_hours = 24, spike_times_single = 12,
                                      spike_times_multi = c(6, 12, 18))
  p  <- plot_temperature_scenarios(scens)
  bd <- built(p)
  # One facet panel per scenario.
  expect_equal(length(unique(bd$data[[1]]$PANEL)), length(scens))
  # The plotted y-values are exactly the concatenated scenario temperatures.
  plotted_y <- sort(bd$data[[1]]$y)
  raw_y     <- sort(unlist(lapply(scens, function(s) s$temp), use.names = FALSE))
  expect_equal(plotted_y, raw_y, tolerance = 1e-9)
})

test_that("plot_temperature_scenarios adds a T_c reference line only when asked", {
  scens <- make_temperature_scenarios(n_hours = 12)
  expect_equal(length(plot_temperature_scenarios(scens)$layers), 1L)
  p_tc <- plot_temperature_scenarios(scens, T_c = 24)
  expect_equal(length(p_tc$layers), 2L)            # line + hline
  # The hline sits at T_c.
  hline_layer <- p_tc$layers[[2]]
  expect_equal(hline_layer$data$yintercept, 24)
})

# ========================= plot_repair_rate ================================

test_that("plot_repair_rate plots exactly the Schoolfield repair rate", {
  rp <- list(TA = 14065, TAL = 50000, TAH = 120000,
             TL = 10.5 + 273.15, TH = 22.5 + 273.15,
             TREF = 17 + 273.15, r_ref = 0.01)
  grid  <- seq(5, 30, length.out = 50)
  p     <- plot_repair_rate(grid, rp)
  line_y <- built(p)$data[[1]]$y
  truth  <- repair_rate_schoolfield(grid, rp$TA, rp$TAL, rp$TAH,
                                    rp$TL, rp$TH, rp$TREF, rp$r_ref)
  expect_equal(line_y, truth, tolerance = 1e-9)    # the line IS the rate
})

# ====================== plot_temperature_density ===========================

# Synthetic derive_temperature_for_duration()-style temperature posterior.
fake_temp_post <- function(center = 35, sd = 1.2, n = 4000) {
  set.seed(1)
  draws <- tibble::tibble(.draw = seq_len(n), temp = stats::rnorm(n, center, sd))
  q <- stats::quantile(draws$temp, c(0.025, 0.5, 0.975), names = FALSE)
  list(draws = draws,
       summary = tibble::tibble(temp_lower = q[1], temp_median = q[2],
                                temp_upper = q[3]))
}

test_that("plot_temperature_density draws the density and the CI segment at the summary bounds", {
  tp <- fake_temp_post()
  p  <- plot_temperature_density(tp)
  # Layers: density, CI segment, median point.
  expect_gte(length(p$layers), 3L)
  seg_layer <- p$layers[[2]]
  # The credible-interval segment spans temp_lower -> temp_upper.
  expect_equal(seg_layer$data$temp_lower, tp$summary$temp_lower)
  expect_equal(seg_layer$data$temp_upper, tp$summary$temp_upper)
  expect_silent(built(p))
})

test_that("plot_temperature_density adds a truth line only when supplied", {
  tp <- fake_temp_post()
  expect_equal(length(plot_temperature_density(tp)$layers), 3L)
  p_truth <- plot_temperature_density(tp, truth = 36)
  expect_equal(length(p_truth$layers), 4L)
  expect_equal(p_truth$layers[[4]]$data$xintercept %||%
               p_truth$layers[[4]]$geom_params$xintercept, 36)
})

# ========================= plot_heat_injury ================================

fake_hi <- function() {
  t <- 0:48
  hi <- pmin(150, 3 * pmax(0, t - 6))
  list(summary = tibble::tibble(
    time = t,
    hi_median = hi, hi_lower = hi * 0.8, hi_upper = hi * 1.2,
    surv_median = pmax(0, 1 - hi / 200),
    surv_lower  = pmax(0, 1 - hi / 160),
    surv_upper  = pmin(1, 1 - hi / 260)))
}

test_that("plot_heat_injury returns two stacked panels with the HI threshold line", {
  hi <- fake_hi()
  p  <- plot_heat_injury(hi, lt50_threshold = 100)
  expect_s3_class(p, "patchwork")
  # Two sub-plots: HI (top) and survival (bottom).
  expect_equal(length(p$patches$plots) + 1L, 2L)   # patchwork holds n-1 + base
  # Rebuild the HI sub-plot directly to confirm the threshold + data.
  p_hi <- plot_heat_injury(hi)[[1]]
  expect_equal(built(p_hi)$data[[2]]$y, hi$summary$hi_median, tolerance = 1e-9)
})

# ====================== plot_survival_curves ===============================

fake_pred <- function() {
  grid <- expand.grid(temp = c(30, 34, 38), duration = c(0.5, 1, 2, 4))
  grid$survival_median <- plogis(2 - 0.5 * grid$duration - 0.1 * (grid$temp - 34))
  grid$survival_lower  <- pmax(0, grid$survival_median - 0.1)
  grid$survival_upper  <- pmin(1, grid$survival_median + 0.1)
  list(summary = tibble::as_tibble(grid))
}

fake_grouped_pred <- function() {
  grid <- expand.grid(oxygen = c("normoxia", "hyperoxia"),
                      temp = c(26, 38, 39),
                      duration = c(10, 30, 60))
  grid$survival_median <- plogis(5 - 0.04 * grid$duration -
                                   0.6 * (grid$temp - 35) +
                                   ifelse(grid$oxygen == "hyperoxia", 1, 0))
  grid$survival_lower  <- pmax(0, grid$survival_median - 0.1)
  grid$survival_upper  <- pmin(1, grid$survival_median + 0.1)
  list(summary = tibble::as_tibble(grid))
}

test_that("plot_survival_curves encodes the survival summary and one line per temperature", {
  pr <- fake_pred()
  p  <- plot_survival_curves(pr)
  bd <- built(p)
  # Ribbon (layer 1) ymin/ymax track the credible band.
  expect_equal(sort(bd$data[[1]]$ymin), sort(pr$summary$survival_lower),
               tolerance = 1e-9)
  # One colour group per temperature.
  expect_equal(length(unique(bd$data[[2]]$colour)),
               length(unique(pr$summary$temp)))
})

test_that("plot_survival_curves log_time flips the duration axis to log10", {
  pr <- fake_pred()
  expect_false(has_log_scale(plot_survival_curves(pr, log_time = FALSE), "x"))
  expect_true(has_log_scale(plot_survival_curves(pr, log_time = TRUE), "x"))
})

test_that("plot_survival_curves overlays observed points only when provided", {
  pr  <- fake_pred()
  obs <- tibble::tibble(temp = c(30, 34), duration = c(1, 2),
                        survival = c(0.8, 0.4))
  expect_equal(length(plot_survival_curves(pr)$layers), 2L)        # ribbon+line
  expect_equal(length(plot_survival_curves(pr, observed = obs)$layers), 3L)
})

test_that("clip_to_observed caps the UPPER observed duration but keeps shorter predictions", {
  pr  <- fake_pred()  # temps 30/34/38 x durations 0.5/1/2/4
  obs <- tibble::tibble(temp = c(30, 30, 38, 38),
                        duration = c(0.5, 1, 2, 4),
                        survival = c(0.9, 0.8, 0.5, 0.2))
  # side = "upper": each temp is capped at its longest observed duration, but
  # shorter grid durations are retained. temp 30 obs max 1 -> {0.5, 1} (2, 4
  # dropped). temp 38 obs [2, 4] -> {0.5, 1, 2, 4}: 4 is the cap, and 0.5/1
  # (below the observed minimum) are KEPT. temp 34 has no observed cell -> dropped.
  cd <- plot_survival_curves(pr, observed = obs)$data
  expect_setequal(cd$duration[cd$temp == 30], c(0.5, 1))
  expect_setequal(cd$duration[cd$temp == 38], c(0.5, 1, 2, 4))
  expect_false(any(cd$temp == 34))
  # No observed supplied -> no clip (no error), full grid.
  no_obs <- plot_survival_curves(pr)$data
  expect_setequal(no_obs$duration[no_obs$temp == 30], c(0.5, 1, 2, 4))
  # Explicit FALSE keeps the full grid even with observed.
  full <- plot_survival_curves(pr, observed = obs, clip_to_observed = FALSE)$data
  expect_setequal(full$duration[full$temp == 30], c(0.5, 1, 2, 4))
})

test_that("plot_survival_curves drops unsupported grouped temperature cells", {
  pr <- fake_grouped_pred()
  obs <- tibble::tibble(
    oxygen = c("normoxia", "normoxia", "hyperoxia", "hyperoxia"),
    temp = c(26, 26, 38, 38),
    duration = c(10, 30, 30, 60),
    survival = c(1, 1, 0.9, 0.5)
  )
  p <- plot_survival_curves(pr, observed = obs)
  cd <- p$data
  cells <- cd[, c("oxygen", "temp")]
  expect_equal(nrow(unique(cells)), 2L)
  expect_true(any(cd$oxygen == "normoxia" & cd$temp == 26))
  expect_true(any(cd$oxygen == "hyperoxia" & cd$temp == 38))
  expect_false(any(cd$oxygen == "hyperoxia" & cd$temp == 26))
  expect_false(any(cd$oxygen == "normoxia" & cd$temp == 38))
  expect_false(any(cd$temp == 39))
  # Duration is capped at each cell's longest observed exposure but keeps shorter
  # grid values (side = "upper"): normoxia@26 obs max 30 -> {10, 30} (60 dropped);
  # hyperoxia@38 obs [30, 60] -> {10, 30, 60} (10 is below the observed min but kept).
  expect_setequal(cd$duration[cd$oxygen == "normoxia" & cd$temp == 26],
                  c(10, 30))
  expect_setequal(cd$duration[cd$oxygen == "hyperoxia" & cd$temp == 38],
                  c(10, 30, 60))

  full <- plot_survival_curves(pr, observed = obs,
                               clip_to_observed = FALSE)$data
  expect_equal(nrow(unique(full[, c("oxygen", "temp")])), 6L)
})

test_that("plot_survival_curves facet_scales frees the duration axis per group (auto)", {
  pr <- fake_grouped_pred()
  obs <- tibble::tibble(
    oxygen = c("normoxia", "normoxia", "hyperoxia", "hyperoxia"),
    temp = c(26, 26, 38, 38),
    duration = c(10, 30, 30, 60),
    survival = c(1, 1, 0.9, 0.5)
  )
  # auto (default): with clipping + observed, free the duration (x) axis per
  # group while survival (y) stays shared.
  p <- plot_survival_curves(pr, observed = obs)
  expect_true(p$facet$params$free$x)
  expect_false(p$facet$params$free$y)
  # Explicit "fixed" frees neither axis.
  p_fixed <- plot_survival_curves(pr, observed = obs, facet_scales = "fixed")
  expect_false(p_fixed$facet$params$free$x)
  expect_false(p_fixed$facet$params$free$y)
})

# ========================= plot_tdt_landscape ==============================

fake_landscape <- function() {
  # Log-spaced duration grid (the derive_tdt_landscape() default) on a linear
  # axis: this is the real default the plot must render. geom_tile handles the
  # irregular spacing; geom_raster would mis-place every cell. Survival spans
  # the full 0-1 range so the 0.25/0.5/0.75 contours all exist.
  grid <- expand.grid(temp = seq(30, 38, length.out = 12),
                      duration = 10^seq(log10(0.5), log10(30), length.out = 12))
  grid$survival_median <- plogis(8 - 0.4 * grid$duration - 0.5 * (grid$temp - 30))
  list(summary = tibble::as_tibble(grid))
}

test_that("plot_tdt_landscape builds a tile heatmap of survival with contours", {
  lsp <- fake_landscape()
  p   <- plot_tdt_landscape(lsp)
  # geom_tile (not geom_raster) so the log-spaced default duration grid renders
  # without the "uneven raster intervals" shift.
  expect_s3_class(p$layers[[1]]$geom, "GeomTile")
  expect_s3_class(p$layers[[2]]$geom, "GeomContour")
  bd <- built(p)
  # One tile cell per grid row.
  expect_equal(nrow(bd$data[[1]]), nrow(lsp$summary))
  # Contours are drawn at exactly the requested survival levels (and they exist
  # because the synthetic surface spans 0-1).
  expect_setequal(unique(bd$data[[2]]$level), c(0.25, 0.5, 0.75))
  # Survival fill scale is clamped to the [0, 1] probability range.
  fill_scale <- p$scales$scales[[which(vapply(p$scales$scales,
                  function(s) any(s$aesthetics == "fill"), logical(1)))[1]]]
  expect_equal(fill_scale$limits, c(0, 1))
})

test_that("plot_tdt_landscape log_time flips the duration axis and overlays observed", {
  lsp <- fake_landscape()
  expect_false(has_log_scale(plot_tdt_landscape(lsp), "y"))
  expect_true(has_log_scale(plot_tdt_landscape(lsp, log_time = TRUE), "y"))
  obs <- tibble::tibble(temp = 34, duration = 1, survival = 0.5)
  expect_gt(length(plot_tdt_landscape(lsp, observed = obs)$layers),
            length(plot_tdt_landscape(lsp)$layers))
})

test_that("plot_tdt_landscape facets into one heatmap panel per group for a grouped fit", {
  lsp <- fake_landscape()
  # Single-condition summary -> exactly one panel (no moderator column).
  bd1 <- built(plot_tdt_landscape(lsp))
  expect_equal(length(unique(bd1$data[[1]]$PANEL)), 1L)

  # A grouped derive_tdt_landscape() prepends the moderator column and stacks one
  # block per level. The plot must split them into one heatmap panel each rather
  # than overplotting both surfaces onto a single (temp, duration) tile grid.
  g <- lsp$summary
  lsp_g <- list(summary = tibble::as_tibble(rbind(
    cbind(oxygen = "normoxia",  g),
    cbind(oxygen = "hyperoxia", g))))
  p_g <- plot_tdt_landscape(lsp_g)
  expect_s3_class(p_g$facet, "FacetWrap")
  bd2 <- built(p_g)
  expect_equal(length(unique(bd2$data[[1]]$PANEL)), 2L)
  # Each panel carries exactly one grid's worth of tiles (no overplot).
  expect_equal(nrow(bd2$data[[1]]), 2L * nrow(g))
})

test_that("plot_tdt_landscape clips to the observed temp x duration box (default)", {
  lsp <- fake_landscape()
  obs <- tibble::tibble(temp = c(32, 36), duration = c(1, 10), survival = c(0.6, 0.3))
  cd  <- plot_tdt_landscape(lsp, observed = obs)$data     # default clip = TRUE
  expect_true(all(cd$temp >= 32 & cd$temp <= 36))         # temperature clipped both ends
  expect_true(all(cd$duration <= 10))                     # duration capped at the upper end
  expect_true(any(cd$duration < 1))                       # ...shorter durations retained (side = "upper")
  # No observed / explicit FALSE -> full grid.
  expect_equal(nrow(plot_tdt_landscape(lsp)$data), nrow(lsp$summary))
  expect_equal(nrow(plot_tdt_landscape(lsp, observed = obs,
                                       clip_to_observed = FALSE)$data),
               nrow(lsp$summary))
})

test_that("plot_tdt_landscape clips grouped surfaces to each group's observed range", {
  lsp <- fake_landscape()
  g <- lsp$summary
  lsp_g <- list(summary = tibble::as_tibble(rbind(
    cbind(oxygen = "normoxia",  g),
    cbind(oxygen = "hyperoxia", g))))
  obs <- tibble::tibble(
    oxygen = c("normoxia", "normoxia", "hyperoxia", "hyperoxia"),
    temp = c(30, 38, 36, 38),
    duration = c(0.5, 30, 0.5, 30),
    survival = c(1, 0, 1, 0)
  )
  p <- plot_tdt_landscape(lsp_g, observed = obs)
  cd <- p$data
  expect_equal(range(cd$temp[cd$oxygen == "normoxia"]), c(30, 38))
  expect_true(min(cd$temp[cd$oxygen == "hyperoxia"]) >= 36)
  expect_equal(max(cd$temp[cd$oxygen == "hyperoxia"]), 38)
  expect_true(p$facet$params$free$x)
  expect_true(p$facet$params$free$y)
  xr <- panel_axis_ranges(p, "x")
  hyper_x <- unlist(xr[xr$oxygen == "hyperoxia", c("axis_min", "axis_max")])
  norm_x <- unlist(xr[xr$oxygen == "normoxia", c("axis_min", "axis_max")])
  expect_lt(diff(hyper_x), diff(norm_x))
  expect_gt(xr$axis_min[xr$oxygen == "hyperoxia"],
            xr$axis_min[xr$oxygen == "normoxia"])

  p_fixed <- plot_tdt_landscape(lsp_g, observed = obs, facet_scales = "fixed")
  expect_false(p_fixed$facet$params$free$x)
  expect_false(p_fixed$facet$params$free$y)
  xr_fixed <- panel_axis_ranges(p_fixed, "x")
  expect_equal(unique(xr_fixed$axis_min), xr_fixed$axis_min[1])
  expect_equal(unique(xr_fixed$axis_max), xr_fixed$axis_max[1])
})

# ========================== plot_tdt_curve =================================

fake_ltx <- function(target = "p=0.500", unit = "min", big = FALSE,
                     grouped = FALSE) {
  temp <- seq(30, 38, by = 1)
  med  <- 10^(3 - 0.15 * (temp - 30))            # falls with temperature
  upr  <- med * if (big) 50 else 1.5             # `big` pushes past the 48 h cap
  summary <- data.frame(target_surv = target, temp = temp,
                        duration_lower = med * 0.7,
                        duration_median = med,
                        duration_upper = upr,
                        check.names = FALSE)
  if (isTRUE(grouped)) {
    summary <- rbind(
      data.frame(grp = "A", summary, check.names = FALSE),
      data.frame(grp = "B",
                 transform(summary,
                           duration_lower = duration_lower * 1.25,
                           duration_median = duration_median * 1.25,
                           duration_upper = duration_upper * 1.25),
                 check.names = FALSE)
    )
  }
  list(summary = tibble::as_tibble(summary),
       output_time_unit = unit)
}

test_that("plot_tdt_curve returns two panels by default and single panels on request", {
  ltx <- fake_ltx()
  expect_s3_class(plot_tdt_curve(ltx), "patchwork")              # both
  expect_s3_class(plot_tdt_curve(ltx, panels = "linear"), "ggplot")
  p_log <- plot_tdt_curve(ltx, panels = "log")
  expect_s3_class(p_log, "ggplot")
  expect_true(has_log_scale(p_log, "y"))                        # log panel
})

test_that("plot_tdt_curve renders the target-survival label for each label style", {
  expect_match(plot_tdt_curve(fake_ltx("p=0.500"), panels = "linear")$labels$y,
               "50% absolute survival")
  expect_match(plot_tdt_curve(fake_ltx("(low+up)/2"), panels = "linear")$labels$y,
               "relative survival")
  expect_match(plot_tdt_curve(fake_ltx("0.1"), panels = "linear")$labels$y,
               "10% absolute survival")
})

test_that("plot_tdt_curve clamps the y-axis to 48 h only when the band exceeds it", {
  # Minutes: 48 h cap = 2880. `big = TRUE` pushes duration_upper well past it.
  capped <- plot_tdt_curve(fake_ltx(unit = "min", big = TRUE), panels = "linear")
  expect_equal(capped$coordinates$limits$y, c(0, 48 * 60))
  # Band stays under the cap -> no coord clamp.
  free <- plot_tdt_curve(fake_ltx(unit = "min", big = FALSE), panels = "linear")
  expect_null(free$coordinates$limits$y)
})

test_that("plot_tdt_curve carries the output time unit into the axis label", {
  expect_match(plot_tdt_curve(fake_ltx(unit = "hours"), panels = "linear")$labels$y,
               "hours")
})

test_that("plot_tdt_curve clips its temperature axis to observed (default) when observed supplied", {
  ltx <- fake_ltx()  # temp seq(30, 38, by = 1)
  obs <- tibble::tibble(temp = c(33, 36), duration = c(10, 10), survival = c(0.5, 0.5))
  cd  <- plot_tdt_curve(ltx, panels = "linear", observed = obs)$data
  expect_true(all(cd$temp >= 33 & cd$temp <= 36))
  # No observed -> full temperature range retained.
  expect_equal(range(plot_tdt_curve(ltx, panels = "linear")$data$temp), c(30, 38))
})

test_that("plot_tdt_curve facets grouped LT curves instead of pooling them", {
  p <- plot_tdt_curve(fake_ltx(grouped = TRUE), panels = "linear")
  expect_s3_class(p$facet, "FacetWrap")
  bd <- built(p)
  expect_equal(length(unique(bd$data[[1]]$PANEL)), 2L)
})