make_exp_dist <- function(rate = 1) {
  dist_spec(
    d    = function(x, p) p$rate * exp(-p$rate * x),
    p    = function(x, p) 1 - exp(-p$rate * x),
    q    = function(u, p) -log(1 - u) / p$rate,
    surv = function(x, p) exp(-p$rate * x),
    lower = 0, upper = Inf,
    param = list(rate = rate),
    name  = "Exp(1)"
  )
}

# -- dist_spec -----------------------------------------------------------------

test_that("dist_spec validates input", {
  expect_error(dist_spec(), "Provide at least one")
  d <- make_exp_dist()
  expect_s3_class(d, "dist_spec")
})

test_that("dist_spec works with only PDF", {
  d <- dist_spec(d = function(x, p) p$rate * exp(-p$rate * x),
                 lower = 0, upper = Inf, param = list(rate = 1))
  expect_s3_class(d, "dist_spec")
})

# -- Reproducibility -----------------------------------------------------------

test_that("reproducibility: same seed gives same sample", {
  d <- make_exp_dist()
  a <- r_complete(100, d, seed = 123)
  b <- r_complete(100, d, seed = 123)
  expect_equal(a$t_true, b$t_true)
  c <- r_type1_censor(100, d, cen_time = 2, seed = 456)
  e <- r_type1_censor(100, d, cen_time = 2, seed = 456)
  expect_equal(c$x, e$x)
  expect_equal(c$status, e$status)
})

test_that("reproducibility: middle censoring with same seed", {
  d <- make_exp_dist()
  a <- r_middle_censor(100, d, lower = 0.5, upper = 1.5, seed = 77)
  b <- r_middle_censor(100, d, lower = 0.5, upper = 1.5, seed = 77)
  expect_equal(a$t_true, b$t_true)
  expect_equal(a$x, b$x)
  expect_equal(a$status, b$status)
  expect_equal(a$time2, b$time2)
})

# -- Complete sample -----------------------------------------------------------

test_that("complete sample has all events", {
  d <- make_exp_dist()
  s <- r_complete(50, d, seed = 1)
  expect_length(s$x, 50)
  expect_true(all(s$status == 1L))
  expect_equal(s$x, s$t_true)
})

# -- Type-II -------------------------------------------------------------------

test_that("type2 censoring yields at least r_failures events", {
  d <- make_exp_dist()
  s <- r_type2_censor(100, d, r_failures = 30, seed = 99)
  expect_gte(sum(s$status == 1L), 30L)
  if (any(s$status == 0L)) {
    expect_true(all(s$x[s$status == 0L] == s$r_th_time))
  }
})

# -- Doubly Type-II ------------------------------------------------------------

test_that("doubly type2 censors correct counts", {
  d <- make_exp_dist()
  s <- r_doubly_type2_censor(50, d, r1 = 5, r2 = 5, seed = 11)
  expect_equal(sum(s$status == 2L), 5L)
  expect_equal(sum(s$status == 0L), 5L)
  expect_equal(sum(s$status == 1L), 40L)
  expect_error(r_doubly_type2_censor(10, d, r1 = 5, r2 = 5, seed = 1),
               "r1.*r2")
})

# -- Middle censoring ----------------------------------------------------------

test_that("middle censoring status codes are correct", {
  d <- make_exp_dist()
  s <- r_middle_censor(200, d, lower = 0.5, upper = 1.5, seed = 42)
  # All status codes must be 1 (event) or 4 (middle-censored)
  expect_true(all(s$status %in% c(1L, 4L)))
  # With n = 200, some should be middle-censored
  expect_gt(sum(s$status == 4L), 0L)
  # Middle-censored: x stores the lower bound
  mc <- s$status == 4L
  expect_true(all(s$x[mc] == 0.5))
  # Middle-censored: time2 stores the upper bound
  expect_true(all(s$time2[mc] == 1.5))
  # Uncensored: x == t_true
  obs <- s$status == 1L
  expect_equal(s$x[obs], s$t_true[obs])
  # Middle-censored: t_true is in [lower, upper]
  expect_true(all(s$t_true[mc] >= 0.5 & s$t_true[mc] <= 1.5))
  # Uncensored: t_true is outside [lower, upper]
  expect_true(all(s$t_true[obs] < 0.5 | s$t_true[obs] > 1.5))
})

test_that("middle censoring validates lower < upper", {
  d <- make_exp_dist()
  expect_error(r_middle_censor(30, d, lower = 2, upper = 1, seed = 1),
               "lower.*upper")
  expect_error(r_middle_censor(30, d, lower = 1, upper = 1, seed = 1),
               "lower.*upper")
})

# -- Right / Left censoring ---------------------------------------------------

test_that("right censoring caps at cen_time", {
  d <- make_exp_dist()
  s <- r_right_censor(100, d, cen_time = 1, seed = 5)
  expect_true(all(s$x <= 1))
  expect_true(all(s$x[s$status == 0L] == 1))
  expect_equal(s$x[s$status == 1L], s$t_true[s$status == 1L])
})

test_that("left censoring floors at threshold", {
  d <- make_exp_dist()
  s <- r_left_censor(100, d, threshold = 0.5, seed = 5)
  expect_true(all(s$x >= 0.5 | s$status == 1L))
  expect_true(all(s$x[s$status == 2L] == 0.5))
})

# -- Progressive Type-II ------------------------------------------------------

test_that("progressive type2 removes units correctly", {
  d <- make_exp_dist()
  s <- r_progressive_type2_censor(50, d, scheme = c(2, 2, 2), seed = 10)
  expect_length(s$x, 50)
  expect_true(sum(s$status == 0L) <= 6L)  # at most sum(scheme) removed
  # All censored x values should be <= some failure time
  if (any(s$status == 0L)) {
    expect_true(all(s$x[s$status == 0L] <= max(s$t_true)))
  }
})

# -- Hybrid censoring ---------------------------------------------------------

test_that("hybrid censoring stops at min(cen_time, r-th failure)", {
  d <- make_exp_dist()
  s <- r_hybrid_censor(100, d, cen_time = 5, r_failures = 50, seed = 3)
  expect_length(s$x, 100)
  expect_true(all(s$x[s$status == 0L] == s$stop_time))
})

# -- Truncation ----------------------------------------------------------------

test_that("right truncation generates values <= right", {
  d <- make_exp_dist()
  s <- r_right_truncation(100, d, right = 2, seed = 10)
  expect_true(all(s$x <= 2))
  expect_true(all(s$status == 1L))
})

test_that("left truncation generates values >= left", {
  d <- make_exp_dist()
  s <- r_left_truncation(100, d, left = 0.5, seed = 10)
  expect_true(all(s$x >= 0.5))
  expect_true(all(s$status == 1L))
})

# -- Plot functions ------------------------------------------------------------

test_that("plot functions run without error", {
  d <- make_exp_dist()
  s <- r_random_censor(80, d, seed = 7)
  grDevices::pdf(NULL)
  on.exit(grDevices::dev.off(), add = TRUE)
  expect_invisible(plot_censor_hist(s))
  expect_invisible(plot_censor_dot(s))
  expect_invisible(plot_censor_acf(s))
  expect_invisible(plot_censor_diagnostics(s))
})

test_that("plot functions work for middle censoring", {
  d <- make_exp_dist()
  s <- r_middle_censor(80, d, lower = 0.5, upper = 1.5, seed = 7)
  grDevices::pdf(NULL)
  on.exit(grDevices::dev.off(), add = TRUE)
  expect_invisible(plot_censor_hist(s))
  expect_invisible(plot_censor_dot(s))
  expect_invisible(plot_censor_acf(s))
})

# -- All scheme functions return censored_sample -------------------------------

test_that("all scheme functions return censored_sample", {
  d <- make_exp_dist()
  n <- 40L
  schemes <- list(
    r_complete(n, d, seed = 1),
    r_right_truncation(n, d, right = 3, seed = 1),
    r_left_truncation(n, d, left = 0.5, seed = 2),
    r_random_censor(n, d, seed = 3),
    r_right_censor(n, d, cen_time = 2, seed = 4),
    r_left_censor(n, d, threshold = 0.3, seed = 5),
    r_interval_censor(n, d, width = 0.5, seed = 6),
    r_middle_censor(n, d, lower = 0.5, upper = 1.5, seed = 6),
    r_block_random_censor(n, d, seed = 7),
    r_type1_censor(n, d, cen_time = 2, seed = 8),
    r_type2_censor(n, d, r_failures = 20, seed = 9),
    r_progressive_type2_censor(n, d, seed = 10),
    r_progressive_first_failure(n, d, seed = 11),
    r_hybrid_censor(n, d, cen_time = 2, r_failures = 15, seed = 12),
    r_hybrid_type1_censor(n, d, cen_time = 3,
                          prog_times = c(1, 2), scheme = c(2, 2), seed = 13),
    r_doubly_type2_censor(n, d, r1 = 5, r2 = 5, seed = 14),
    r_type1_hybrid_censor(n, d, cen_time = 2, r_failures = 15, seed = 15),
    r_hybrid_type2_censor(n, d, scheme = c(1, 1), r_failures = 10, seed = 16),
    r_progressive_hybrid_type2_censor(n, d, scheme = c(1, 1), r_failures = 10, cen_time = 3, seed = 17)
  )
  for (s in schemes) {
    expect_s3_class(s, "censored_sample")
    expect_equal(length(s$x), n)
    expect_equal(length(s$status), n)
    expect_equal(length(s$t_true), n)
  }
  j1 <- r_joint_type1_censor(20, d, cen_time = 2, k = 2, seed = 17)
  j2 <- r_joint_type2_censor(20, d, r_failures = 15, k = 2, seed = 18)
  bj <- r_bjpt2_censor(20, d, k = 2, seed = 19)
  expect_s3_class(j1, "censored_sample_list")
  expect_s3_class(j2, "censored_sample_list")
  expect_s3_class(bj, "censored_sample_list")
})

# -- Type-II Progressively Hybrid Censoring ------------------------------------

test_that("Type-II progressively hybrid censoring stops at min(m-th failure, T)", {
  d <- make_exp_dist()
  # Case I: m-th failure before T (T = 5, m = 10)
  s1 <- r_progressive_hybrid_type2_censor(50, d, scheme = rep(1L, 10), r_failures = 10, cen_time = 5, seed = 42)
  expect_s3_class(s1, "censored_sample")
  expect_length(s1$x, 50)
  # Check that stop_time is the 10th failure time
  expect_true(s1$stop_time < 5)
  expect_equal(sum(s1$status == 1L), 10L)
  
  # Case II: m-th failure after T (T = 0.5, m = 40)
  s2 <- r_progressive_hybrid_type2_censor(50, d, scheme = rep(1L, 40), r_failures = 40, cen_time = 0.5, seed = 42)
  expect_equal(s2$stop_time, 0.5)
  expect_true(sum(s2$status == 1L) < 40L)
  expect_true(all(s2$x[s2$status == 0L] <= 0.5))
})

# -- Coercion methods ----------------------------------------------------------

test_that("as.vector and as.data.frame work", {
  d <- make_exp_dist()
  s <- r_complete(20, d, seed = 1)
  v <- as.vector(s)
  expect_length(v, 20)
  df <- as.data.frame(s)
  expect_s3_class(df, "data.frame")
  expect_equal(nrow(df), 20)
  expect_true(all(c("x", "status", "t_true", "time2") %in% names(df)))
})

# -- Edge cases ----------------------------------------------------------------

test_that("n = 1 works for all basic schemes", {
  d <- make_exp_dist()
  expect_s3_class(r_complete(1, d, seed = 1), "censored_sample")
  expect_s3_class(r_right_censor(1, d, cen_time = 2, seed = 1), "censored_sample")
  expect_s3_class(r_type1_censor(1, d, cen_time = 2, seed = 1), "censored_sample")
  expect_s3_class(r_type2_censor(1, d, r_failures = 1, seed = 1), "censored_sample")
  expect_s3_class(r_middle_censor(1, d, lower = 0.5, upper = 1.5, seed = 1),
                  "censored_sample")
  expect_s3_class(r_progressive_hybrid_type2_censor(1, d, scheme = 0L, r_failures = 1, cen_time = 2, seed = 1),
                  "censored_sample")
})

test_that("invalid n is rejected", {
  d <- make_exp_dist()
  expect_error(r_complete(0, d), "positive")
  expect_error(r_complete(-5, d), "positive")
})
