test_that("run_llm_rem (roleplay_sender) returns correct structure", {
  res <- run_llm_rem(
    test_edgelist, test_formula,
    query_fn     = make_query_roleplay_sender(mock_cfg),
    n_nonevents  = NULL, base_riskset = test_riskset, seed = 1L
  )
  expect_named(res, c("sequence", "responses", "rem"))
  expect_s3_class(res$rem, "coxph")
})

test_that("run_llm_rem sequence has the correct number of rows", {
  res <- run_llm_rem(
    test_edgelist, test_formula,
    query_fn     = make_query_roleplay_sender(mock_cfg),
    n_nonevents  = NULL, base_riskset = test_riskset, seed = 1L
  )
  expect_equal(nrow(res$sequence), nrow(test_edgelist))
})

test_that("roleplay_sender preserves the empirical sender sequence", {
  res <- run_llm_rem(
    test_edgelist, test_formula,
    query_fn     = make_query_roleplay_sender(mock_cfg),
    n_nonevents  = NULL, base_riskset = test_riskset, seed = 1L
  )
  expect_identical(res$sequence$sender, test_edgelist$sender)
})

test_that("generated sequence contains no self-loops", {
  res <- run_llm_rem(
    test_edgelist, test_formula,
    query_fn     = make_query_roleplay_sender(mock_cfg),
    n_nonevents  = NULL, base_riskset = test_riskset, seed = 1L
  )
  expect_true(all(res$sequence$sender != res$sequence$receiver))
})

test_that("all generated receivers are valid node IDs", {
  res <- run_llm_rem(
    test_edgelist, test_formula,
    query_fn     = make_query_roleplay_sender(mock_cfg),
    n_nonevents  = NULL, base_riskset = test_riskset, seed = 1L
  )
  expect_true(all(res$sequence$receiver %in% test_nodes))
})

test_that("run_llm_rem (predict) produces no self-loops and valid IDs", {
  res <- run_llm_rem(
    test_edgelist, test_formula,
    query_fn     = make_query_predict(mock_cfg),
    n_nonevents  = NULL, base_riskset = test_riskset, seed = 2L
  )
  expect_true(all(res$sequence$sender   %in% test_nodes))
  expect_true(all(res$sequence$receiver %in% test_nodes))
  expect_true(all(res$sequence$sender   != res$sequence$receiver))
})

test_that("run_llm_rem (roleplay_random) senders deviate from empirical", {
  res <- run_llm_rem(
    test_edgelist, test_formula,
    query_fn     = make_query_roleplay_random(mock_cfg),
    n_nonevents  = NULL, base_riskset = test_riskset, seed = 3L
  )
  # Exclude the seeded first event; with 6 nodes at least one should differ
  n_diff <- sum(res$sequence$sender[-1L] != test_edgelist$sender[-1L])
  expect_gte(n_diff, 1L)
})

# fit_rem_sep -----------------------------------------------------------------

test_that("fit_rem_sep returns a rem_sep object with correct structure", {
  fit <- fit_rem_sep(
    test_edgelist,
    formula_init = survival::Surv(time1, occurred) ~
      RRecInitiation + NOTPSnd + survival::strata(time2),
    formula_cont = survival::Surv(time1, occurred) ~
      Repetition + RecRecencyLog + survival::strata(time2),
    n_nonevents  = NULL,
    base_riskset = test_riskset
  )
  expect_s3_class(fit, "rem_sep")
  expect_named(fit, c("init","cont","loglik","n_init","n_cont",
                       "init_fraction","method"), ignore.order = TRUE)
  expect_true(fit$n_init + fit$n_cont == nrow(test_edgelist) - 1L)
  expect_true(fit$init_fraction >= 0 && fit$init_fraction <= 1)
})

test_that("fit_rem_sep loglik is sum of sub-model logliks", {
  fit <- fit_rem_sep(
    test_edgelist,
    formula_init = survival::Surv(time1, occurred) ~
      RRecInitiation + survival::strata(time2),
    formula_cont = survival::Surv(time1, occurred) ~
      Repetition + survival::strata(time2),
    n_nonevents  = NULL,
    base_riskset = test_riskset
  )
  expect_equal(unname(fit$loglik["sep"]),
               unname(fit$loglik["init"]) + unname(fit$loglik["cont"]),
               tolerance = 1e-6)
})

test_that("fit_rem_sep initiation events are genuinely first contacts", {
  fit <- fit_rem_sep(
    test_edgelist,
    formula_init = survival::Surv(time1, occurred) ~
      RRecInitiation + survival::strata(time2),
    formula_cont = survival::Surv(time1, occurred) ~
      Repetition + survival::strata(time2),
    n_nonevents  = NULL,
    base_riskset = test_riskset
  )
  # n_init + n_cont == nrow(test_edgelist) - 1 (event 1 has no risk set)
  expect_equal(fit$n_init + fit$n_cont, nrow(test_edgelist) - 1L)
})

test_that("fit_rem_sep clogit method returns rem_clogit sub-models", {
  fit <- fit_rem_sep(
    test_edgelist,
    formula_init = survival::Surv(time1, occurred) ~
      RRecInitiation + survival::strata(time2),
    formula_cont = survival::Surv(time1, occurred) ~
      Repetition + survival::strata(time2),
    n_nonevents  = NULL,
    base_riskset = test_riskset,
    method       = "clogit"
  )
  expect_s3_class(fit, "rem_sep")
  if (!is.null(fit$init)) expect_s3_class(fit$init, "rem_clogit")
  if (!is.null(fit$cont)) expect_s3_class(fit$cont, "rem_clogit")
})
