library(testthat)
library(EpidigiR)

test_that("epi_analyze returns expected output for valid input", {
  df <- data.frame(region = c("A", "B", "C"),
                   cases = c(10, 20, 30),
                   population = c(100, 200, 300))
  result <- epi_analyze(df, outcome = "cases",
                        population = "population",
                        type = "summary")
  expect_s3_class(result, "data.frame")
  expected_prevalence <- c(0.1, 0.1, 0.1)
  expect_equal(result$prevalence,
               expected_prevalence)
})

test_that("epi_analyze throws error for missing columns", {
  df <- data.frame(cases = c(5, 6, 7))
  expect_error(epi_analyze(df, outcome = "cases",
                           population = "population",
                           type = "summary"))
})

test_that("epi_analyze SIR model returns expected structure and population conservation", {
  result <- epi_analyze(data = NULL,
                        outcome = NULL,
                        population = NULL,
                        type = "sir",
                        N = 1000,
                        beta = 0.3,
                        gamma = 0.1,
                        days = 20)
  expect_s3_class(result, "data.frame")
  expect_true(all(c("time",
                    "Susceptible",
                    "Infected",
                    "Recovered") %in% names(result)))
  expect_equal(nrow(result), 20)
  total_pop <- result$Susceptible + result$Infected + result$Recovered
  expect_true(all(abs(total_pop - 1000) < 1e-6))
})