test_that("bqr_count returns correct class and structure", {
  set.seed(42)
  n <- 60
  X <- matrix(rnorm(n * 2), n, 2)
  y <- rpois(n, exp(0.5 + 0.3 * X[, 1]))

  fit <- bqr_count(y, X, tau = 0.5, method = "lasso",
                   N_iter = 300, burn_in = 100, thin = 2, n_chains = 1)


  expect_s3_class(fit, "bqr_count")
  expect_true(is.array(fit$beta_samples))
  expect_equal(length(dim(fit$beta_samples)), 3)
  expect_equal(dim(fit$beta_samples)[2], 3)  # intercept + 2 covariates
  expect_equal(dim(fit$beta_samples)[3], 1)  # 1 chain
  expect_true(is.matrix(fit$sigma_samples))
  expect_true(is.matrix(fit$lambda_samples))
  expect_true(is.data.frame(fit$summary))
  expect_true(is.numeric(fit$coefficients))
  expect_equal(length(fit$coefficients), 3)
  expect_true(is.numeric(fit$fitted_quantiles))
  expect_equal(length(fit$fitted_quantiles), n)
  expect_true(is.logical(fit$variable_selection))
  expect_equal(fit$tau, 0.5)
  expect_equal(fit$method, "lasso")
})

test_that("bqr_count input validation works", {
  set.seed(42)
  n <- 30
  X <- matrix(rnorm(n * 2), n, 2)
  y <- rpois(n, 2)

  expect_error(bqr_count(y = c(-1, 2, 3), X = X[1:3, ], tau = 0.5),
               "non-negative")
  expect_error(bqr_count(y = c(1.5, 2, 3), X = X[1:3, ], tau = 0.5),
               "integer")
  expect_error(bqr_count(y, X, tau = 1.5), "between 0 and 1")
  expect_error(bqr_count(y, X, tau = 0.5, N_iter = 50, burn_in = 100),
               "greater than")
})

test_that("no NaN/NA/Inf in outputs", {
  set.seed(42)
  n <- 60
  X <- matrix(rnorm(n * 2), n, 2)
  y <- rpois(n, exp(0.5 + 0.3 * X[, 1]))

  fit <- bqr_count(y, X, tau = 0.5, method = "random_bridge",
                   N_iter = 300, burn_in = 100, thin = 2, n_chains = 1)

  expect_true(all(is.finite(fit$coefficients)))
  expect_true(all(is.finite(fit$fitted_quantiles)))
  expect_true(all(is.finite(as.vector(fit$beta_samples))))
  expect_true(all(is.finite(fit$sigma_samples)))
  expect_true(all(is.finite(fit$lambda_samples)))
  expect_true(all(is.finite(fit$xi_samples)))
  expect_true(all(is.finite(fit$summary$Mean)))
  expect_true(all(is.finite(fit$summary$SD)))
})

test_that("all three methods produce valid results", {
  set.seed(42)
  n <- 60
  X <- matrix(rnorm(n * 2), n, 2)
  y <- rpois(n, exp(0.5 + 0.3 * X[, 1]))

  for (meth in c("random_bridge", "fixed_bridge", "lasso")) {
    fit <- bqr_count(y, X, tau = 0.5, method = meth,
                     N_iter = 300, burn_in = 100, thin = 2, n_chains = 1)
    expect_s3_class(fit, "bqr_count")
    expect_equal(fit$method, meth)
    expect_true(all(is.finite(fit$coefficients)))
  }
})

test_that("predict returns correct dimensions", {
  set.seed(42)
  n <- 60
  X <- matrix(rnorm(n * 2), n, 2)
  y <- rpois(n, exp(0.5 + 0.3 * X[, 1]))

  fit <- bqr_count(y, X, tau = 0.5, method = "lasso",
                   N_iter = 300, burn_in = 100, thin = 2, n_chains = 1)

  preds <- predict(fit)
  expect_equal(length(preds), n)
  expect_true(all(preds >= 0))

  new_X <- matrix(rnorm(10 * 2), 10, 2)
  preds_new <- predict(fit, newdata = new_X)
  expect_equal(length(preds_new), 10)
  expect_true(all(preds_new >= 0))

  preds_ci <- predict(fit, interval = "credible")
  expect_true(is.data.frame(preds_ci))
  expect_equal(nrow(preds_ci), n)
  expect_true(all(c("fit", "lower", "upper") %in% names(preds_ci)))
})

test_that("S3 methods work", {
  set.seed(42)
  n <- 60
  X <- matrix(rnorm(n * 2), n, 2)
  y <- rpois(n, exp(0.5 + 0.3 * X[, 1]))

  fit <- bqr_count(y, X, tau = 0.5, method = "lasso",
                   N_iter = 300, burn_in = 100, thin = 2, n_chains = 1)

  expect_output(print(fit), "Bayesian Quantile Regression")
  expect_s3_class(summary(fit), "summary.bqr_count")
  expect_equal(length(coef(fit)), 3)
  expect_equal(length(fitted(fit)), n)
  expect_equal(length(residuals(fit)), n)
})

test_that("variable selection works", {
  set.seed(42)
  n <- 60
  X <- matrix(rnorm(n * 2), n, 2)
  y <- rpois(n, exp(0.5 + 0.8 * X[, 1]))

  fit <- bqr_count(y, X, tau = 0.5, method = "lasso",
                   N_iter = 300, burn_in = 100, thin = 2, n_chains = 1)

  vs <- variable_select(fit)
  expect_true(is.list(vs))
  expect_true("selected" %in% names(vs))
  expect_true("n_selected" %in% names(vs))
  expect_true(vs$n_selected >= 0 && vs$n_selected <= 2)

  sm <- selection_metrics(fit, true_beta = c(0.8, 0))
  expect_true(is.data.frame(sm))
  expect_true(all(c("F1", "Precision", "Recall", "MSE", "MAE", "Bias") %in% names(sm)))
  expect_true(all(is.finite(as.numeric(sm[1, ]))))
})

test_that("compare_models works", {
  set.seed(42)
  n <- 60
  X <- matrix(rnorm(n * 2), n, 2)
  y <- rpois(n, exp(0.5 + 0.3 * X[, 1]))

  fit1 <- bqr_count(y, X, tau = 0.5, method = "lasso",
                    N_iter = 300, burn_in = 100, thin = 2, n_chains = 1)
  fit2 <- bqr_count(y, X, tau = 0.5, method = "random_bridge",
                    N_iter = 300, burn_in = 100, thin = 2, n_chains = 1)

  comp <- compare_models(Lasso = fit1, RB = fit2)
  expect_true(is.data.frame(comp))
  expect_equal(nrow(comp), 2)
  expect_true(all(c("RMSE", "MAE") %in% names(comp)))
})

test_that("gelman_rubin works with multiple chains", {
  set.seed(42)
  n <- 60
  X <- matrix(rnorm(n * 2), n, 2)
  y <- rpois(n, exp(0.5 + 0.3 * X[, 1]))

  fit <- bqr_count(y, X, tau = 0.5, method = "lasso",
                   N_iter = 300, burn_in = 100, thin = 2, n_chains = 2)

  gr <- gelman_rubin(fit)
  expect_true(is.data.frame(gr))
  expect_true("Rhat" %in% names(gr))
  expect_true(all(is.finite(gr$Rhat)))
  expect_true(all(gr$Rhat > 0))
})

test_that("boundary case: y = 0 works", {
  set.seed(42)
  n <- 60
  X <- matrix(rnorm(n * 2), n, 2)
  y <- rep(0L, n)

  fit <- bqr_count(y, X, tau = 0.5, method = "lasso",
                   N_iter = 300, burn_in = 100, thin = 2, n_chains = 1)

  expect_s3_class(fit, "bqr_count")
  expect_true(all(is.finite(fit$coefficients)))
  expect_true(all(fit$fitted_quantiles >= 0))
})