# Test script for Multiplicative Stratified RMST models (Boot & Analytic)
library(testthat)

# --- Test Data Setup ---
# CORRECTED: Increased the size of the pilot data to make bootstrap tests more stable.
set.seed(1003)
pilot_data_mult_strat <- data.frame(
   time = stats::rexp(300, rate = 0.15),
   status = stats::rbinom(300, 1, 0.7),
   arm = rep(c(rep(0, 50), rep(1, 50)), times = 3),
   region = factor(rep(c("A", "B", "C"), each = 100)),
   age = stats::rnorm(300, 55, 8)
)
pilot_data_mult_strat$time[pilot_data_mult_strat$arm == 1] <-
   pilot_data_mult_strat$time[pilot_data_mult_strat$arm == 1] * 1.6 # Multiplicative effect


# --- Test Suite for MS.power.boot ---

test_that("MS.power.boot returns correct structure", {
   set.seed(123)
   results <- MS.power.boot(
      pilot_data = pilot_data_mult_strat,
      time_var = "time", status_var = "status", arm_var = "arm", strata_var = "region",
      sample_sizes = c(60), L = 10, n_sim = 10
   )
   expect_type(results, "list")
   expect_named(results, c("results_data", "results_plot", "results_summary", "model_output"))
   expect_s3_class(results$results_data, "data.frame")
   expect_named(results$results_data, c("N_per_Stratum", "Power"))
})


# --- Test Suite for MS.ss.boot ---

test_that("MS.ss.boot finds a plausible N", {
   set.seed(456)
   results <- expect_no_warning(MS.ss.boot(
      pilot_data = pilot_data_mult_strat,
      time_var = "time", status_var = "status", arm_var = "arm", strata_var = "region",
      target_power = 0.7, L = 10, n_sim = 10,
      n_start = 50, n_step = 50, patience = 2
   ))
   expect_s3_class(results$results_data, "data.frame")
   expect_true(results$results_data$Required_N_per_Stratum > 0)
})


# --- Test Suite for MS.power.analytical ---

test_that("MS.power.analytical returns correct structure", {
   set.seed(789)
   results <- MS.power.analytical(
      pilot_data = pilot_data_mult_strat,
      time_var = "time", status_var = "status", arm_var = "arm", strata_var = "region",
      sample_sizes = c(80, 100), L = 10
   )
   expect_type(results, "list")
   expect_named(results, c("results_data", "results_plot", "results_summary", "model_output"))
})


# --- Test Suite for MS.ss.analytical ---

test_that("MS.ss.analytical finds a plausible N", {
   set.seed(101)
   results <- MS.ss.analytical(
      pilot_data = pilot_data_mult_strat,
      time_var = "time", status_var = "status", arm_var = "arm", strata_var = "region",
      target_power = 0.8, L = 10
   )
   expect_s3_class(results$results_data, "data.frame")
   expect_true(results$results_data$Required_N_per_Stratum > 0)
})


# --- Cross-Method Comparison Test ---

test_that("Analytical and Bootstrap methods give comparable results for multiplicative model", {
   set.seed(111)

   ss_analytic <- MS.ss.analytical(
      pilot_data = pilot_data_mult_strat,
      time_var = "time", status_var = "status", arm_var = "arm", strata_var = "region",
      target_power = 0.7, L = 10
   )$results_data$Required_N_per_Stratum

   power_boot <- MS.power.boot(
      pilot_data = pilot_data_mult_strat,
      time_var = "time", status_var = "status", arm_var = "arm", strata_var = "region",
      sample_sizes = ss_analytic,
      L = 10,
      n_sim = 10 # CORRECTED: Reduced n_sim to make test faster
   )$results_data$Power

   expect_true(is.numeric(power_boot) && is.finite(power_boot) && power_boot >= 0 && power_boot <= 1)
})