library(testthat)
library(EpidigiR)

test_that("epi_visualize executes without errors for various plot types", {
# Scatter
  df_scatter <- data.frame(region_num = 1:3,
                           prevalence = c(0.1, 0.2, 0.15))
  expect_error(epi_visualize(df_scatter,
                             x = "region_num",
                             y = "prevalence",
                             type = "scatter"), NA)

# Curve
  df_curve <- data.frame(day = 1:10, cases = c(1,2,3,5,8,13,21,34,55,89))
  expect_error(epi_visualize(df_curve,
                             x = "day",
                             y = "cases",
                             type = "curve"), NA)

# Boxplot
  df_box <- data.frame(group = rep(c("A","B"), each = 5), value = 1:10)
  expect_error(epi_visualize(df_box,
                             x = "group",
                             y = "value",
                             type = "boxplot"), NA)
})


# Survey


test_that("Survey data prevalence and visualization works", {

# Simulated survey data
  survey_data <- data.frame(
    respondent_id = 1:100,
    age_group = rep(c("18-29", "30-44", "45-59", "60+"),
                    each = 25),
    vaccinated = sample(c(0,1), 100, replace = TRUE)
  )


# Prevalence calculation
  prevalence_df <- aggregate(
    vaccinated ~ age_group,
    data = survey_data,
    FUN = function(x) mean(x, na.rm = TRUE)
  )

  expect_s3_class(prevalence_df, "data.frame")
  expect_true(all(c("age_group", "vaccinated")
                  %in% colnames(prevalence_df)))
  expect_type(prevalence_df$vaccinated, "double")
  expect_true(all(prevalence_df$vaccinated >= 0 & prevalence_df$vaccinated <= 1))

# Logistic regression
  res_log <- suppressWarnings(epi_model(
    data = survey_data,
    formula = vaccinated ~ age_group,
    type = "logistic"
  ))

  expect_true("coefficients" %in% names(res_log))
  expect_true("predictions" %in% names(res_log))
  expect_equal(length(res_log$predictions), nrow(survey_data))

# Visualization test — check it runs without error
  expect_error(
    epi_visualize(
      prevalence_df,
      x = "age_group",
      y = "vaccinated",
      type = "boxplot",
      main = "Vaccination Prevalence by Age Group"
    ),
    NA  # NA here means "expect no error"
  )

})