#' Test Script for Updated calculate_predicted_distributions() and bayesian_statistical_intervals()
#' with Covariate Support
#'
#' This script tests the new covariate functionality using a toy MCMC dataset

# Load required libraries
library(testthat)
library(data.table)
library(dplyr)

# ============================================================================
# SETUP: Create Toy Dataset
# ============================================================================

toy_data <- data.frame(
  EGTPTNUM = c(60, 60, 60, 60, 60, 60, 60, 60, 60, 60),
  Iteration = c(1001, 1301, 1601, 1901, 2201, 2501, 2801, 3101, 3401, 3701),
  s2 = c(74.1817, 58.917, 75.0402, 64.747, 67.8113, 64.5019, 69.6509, 70.4482, 55.461, 70.8997),
  b0 = c(45.8727, 46.3931, 63.4306, 37.1383, 37.4012, 14.781, 31.5076, 39.186, 42.1068, 36.5392),
  b1 = c(-4.5644, -1.8261, -1.9153, -4.8739, -4.1826, -3.6536, -4.6351, -3.1368, -3.3774, -5.8277),
  b2 = c(-5.2276, -3.9667, -2.9237, -7.6383, -2.5675, -5.1732, -1.451, -3.6605, -4.1559, -5.7852),
  b3 = c(6.8522, 10.6505, 10.7839, 6.4292, 8.9568, 8.8313, 4.3675, 7.8938, 9.3794, 8.1801),
  b4 = c(1.3195, 3.6841, 3.3926, 3.126, 3.4957, 2.9057, 3.2898, 2.0038, 3.9755, 1.4375),
  b5 = c(-0.1019, -0.112, -0.1536, -0.0809, -0.0868, -0.0265, -0.07, -0.0899, -0.0994, -0.0786),
  LogPrior = c(-42.25, -42.0222, -42.3597, -42.0782, -42.125, -42.016, -42.1286, -42.169, -41.9415, -42.1678),
  LogLike = c(-863.2, -862.6, -864.8, -862.3, -860, -863.1, -869.4, -860.9, -860.9, -862.4),
  LogPost = c(-905.5, -904.6, -907.1, -904.3, -902.1, -905.1, -911.5, -903.1, -902.9, -904.5)
)

# Print the dataset
cat("=== TOY DATASET ===\n")
print(toy_data)
cat("\n")

# ============================================================================
# TEST PARAMETERS (matching your specification)
# ============================================================================

test_params <- list(
  fixed_effects = "b2",
  covariate_cols = "b3",
  covariate_values = 0.9,
  random_params = "s2",
  multiplication_factor = 1.9,
  by = "EGTPTNUM",
  tolerance_interval_level = 90,
  credible_interval_level = 90,
  prediction_interval_level = 90,
  credibility_of_tolerance_interval = 80
)

cat("=== TEST PARAMETERS ===\n")
cat("Fixed Effects:", test_params$fixed_effects, "\n")
cat("Covariate Column:", test_params$covariate_cols, "\n")
cat("Covariate Value:", test_params$covariate_values, "\n")
cat("Variances:", test_params$random_params, "\n")
cat("Multiplication Factor:", test_params$multiplication_factor, "\n")
cat("By Identifier:", test_params$by, "\n")
cat("Tolerance Interval Level:", test_params$tolerance_interval_level, "%\n")
cat("Credibility of Tolerance:", test_params$credibility_of_tolerance_interval, "%\n")
cat("\n")

# ============================================================================
# EXPECTED RESULTS (from your specification)
# ============================================================================

expected_results <- data.frame(
  By = 60,
  Median = 3.1095,
  CI_Lower = -0.5961,
  CI_Upper = 6.2586,
  PI_Lower = -15.4312,
  PI_Upper = 21.7400,
  TI_Lower = -18.7388,
  TI_Upper = 24.3898
)

cat("=== EXPECTED RESULTS ===\n")
print(expected_results)
cat("\n")

# ============================================================================
# TEST 1: calculate_predicted_distributions() with Covariates
# ============================================================================

cat("=== TEST 1: calculate_predicted_distributions() ===\n")

test_that("calculate_predicted_distributions works with covariates", {

  # Call the function
  pred_dist <- calculate_predicted_distributions(
    data = toy_data,
    fixed_effects = test_params$fixed_effects,
    random_params = test_params$random_params,
    by = test_params$by,
    tolerance_level = test_params$tolerance_interval_level,
    multiplication_factor = test_params$multiplication_factor,
    covariate_cols = test_params$covariate_cols,
    covariate_values = test_params$covariate_values,
    verbose = TRUE
  )

  # Print results
  cat("\n--- Predicted Distributions Output ---\n")
  print(pred_dist)
  cat("\n")

  # Assertions
  expect_true(is.data.frame(pred_dist), "Output should be a data frame")
  expect_true(all(c("By", "mean", "median", "lower_quantile", "upper_quantile", "variance") %in% names(pred_dist)),
              "Output should contain all required columns")
  expect_equal(nrow(pred_dist), 10, label ="Should have 10 rows (one per iteration)")
  expect_true(all(pred_dist$By == 60),"All By values should be 60")

  # Check that mean includes both covariate and fixed effects
  # Manual calculation for first row:
  # mean = (b3 * 0.9) + b2 = (6.8522 * 0.9) + (-5.2276) = 6.167 - 5.2276 = 0.9394
  expected_mean_row1 <- (toy_data$b3[1] * 0.9) + toy_data$b2[1]
  expect_equal(round(pred_dist$mean[1], 2), round(expected_mean_row1, 2),
               tolerance = 0.1,
               label = "Mean should include covariate effect")

  # Check variance calculation
  # variance = s2 * multiplication_factor = 74.1817 * 1.9 = 140.945
  expected_variance_row1 <- toy_data$s2[1] * test_params$multiplication_factor
  expect_equal(round(pred_dist$variance[1], 2), round(expected_variance_row1, 2),
               tolerance = 0.1,
               label = "Variance should be multiplied correctly")

  cat("✓ TEST 1 PASSED: calculate_predicted_distributions() works correctly with covariates\n\n")

  return(pred_dist)
})

# ============================================================================
# TEST 2: bayesian_statistical_intervals() with Covariates
# ============================================================================

cat("=== TEST 2: bayesian_statistical_intervals() ===\n")

test_that("bayesian_statistical_intervals works with covariates", {

  # Call the function
  result <- bayesian_statistical_intervals(
    data = toy_data,
    fixed_effects = test_params$fixed_effects,
    random_params = test_params$random_params,
    by = test_params$by,
    tolerance_interval_level = test_params$tolerance_interval_level,
    multiplication_factor = test_params$multiplication_factor,
    credible_interval_level = test_params$credible_interval_level,
    prediction_interval_level = test_params$prediction_interval_level,
    credibility_of_tolerance_interval = test_params$credibility_of_tolerance_interval,
    log_normal = FALSE,
    covariate_cols = test_params$covariate_cols,
    covariate_values = test_params$covariate_values
  )

  # Extract data and summary
  intervals_data <- result$data
  summary_text <- result$summary

  # Print results
  cat("\n--- Bayesian Statistical Intervals Output ---\n")
  print(intervals_data)
  cat("\n--- Summary ---\n")
  cat(summary_text, "\n\n")

  # Assertions on structure
  expect_true(is.list(result), "Output should be a list")
  expect_true(all(c("data", "summary") %in% names(result)), "Output should contain 'data' and 'summary'")
  expect_true(is.data.frame(intervals_data), "Data component should be a data frame")
  expect_true(is.character(summary_text), "Summary should be a character string")

  # Check required columns
  required_cols <- c("By", "Median", "CI_Lower", "CI_Upper", "PI_Lower", "PI_Upper",
                     "TI_Lower", "TI_Upper", "ReducibleUpper", "ReducibleLower")
  expect_true(all(required_cols %in% names(intervals_data)),
              "Output should contain all required interval columns")

  # Check that we have exactly one row (one By group)
  expect_equal(nrow(intervals_data), 1, label = "Should have 1 row for By = 60")
  expect_equal(intervals_data$By[1], 60, label ="By value should be 60")

  # ============================================================================
  # COMPARE WITH EXPECTED RESULTS
  # ============================================================================

  cat("=== COMPARISON WITH EXPECTED RESULTS ===\n")

  comparison <- data.frame(
    Metric = c("Median", "CI_Lower", "CI_Upper", "PI_Lower", "PI_Upper", "TI_Lower", "TI_Upper"),
    Expected = c(expected_results$Median, expected_results$CI_Lower, expected_results$CI_Upper,
                 expected_results$PI_Lower, expected_results$PI_Upper,
                 expected_results$TI_Lower, expected_results$TI_Upper),
    Actual = c(intervals_data$Median, intervals_data$CI_Lower, intervals_data$CI_Upper,
               intervals_data$PI_Lower, intervals_data$PI_Upper,
               intervals_data$TI_Lower, intervals_data$TI_Upper),
    Difference = NA,
    Within_Tolerance = NA
  )

  comparison$Difference <- comparison$Actual - comparison$Expected
  comparison$Within_Tolerance <- abs(comparison$Difference) < 0.5  # Allow 0.5 unit tolerance

  print(comparison)
  cat("\n")

  # Tolerance-based assertions (allowing for numerical precision differences)
  tolerance <- 0.5  # Allow 0.5 unit difference

  expect_equal(intervals_data$Median, expected_results$Median,
               tolerance = tolerance,
               label = "Median should match expected value")

  expect_equal(intervals_data$CI_Lower, expected_results$CI_Lower,
               tolerance = tolerance,
               label = "CI Lower should match expected value")

  expect_equal(intervals_data$CI_Upper, expected_results$CI_Upper,
               tolerance = tolerance,
               label = "CI Upper should match expected value")

  expect_equal(intervals_data$PI_Lower, expected_results$PI_Lower,
               tolerance = tolerance,
               label = "PI Lower should match expected value")

  expect_equal(intervals_data$PI_Upper, expected_results$PI_Upper,
               tolerance = tolerance,
               label = "PI Upper should match expected value")

  expect_equal(intervals_data$TI_Lower, expected_results$TI_Lower,
               tolerance = tolerance,
               label = "TI Lower should match expected value")

  expect_equal(intervals_data$TI_Upper, expected_results$TI_Upper,
               tolerance = tolerance,
               label = "TI Upper should match expected value")

  # Check that intervals are properly ordered
  expect_true(intervals_data$TI_Lower < intervals_data$PI_Lower,
              "TI Lower should be less than PI Lower")
  expect_true(intervals_data$PI_Lower < intervals_data$CI_Lower,
              "PI Lower should be less than CI Lower")
  expect_true(intervals_data$CI_Lower < intervals_data$Median,
              "CI Lower should be less than Median")
  expect_true(intervals_data$Median < intervals_data$CI_Upper,
              "Median should be less than CI Upper")
  expect_true(intervals_data$CI_Upper < intervals_data$PI_Upper,
              "CI Upper should be less than PI Upper")
  expect_true(intervals_data$PI_Upper < intervals_data$TI_Upper,
              "PI Upper should be less than TI Upper")

  # Check reducible uncertainty is calculated
  expect_true(!is.na(intervals_data$ReducibleUpper), "Reducible Upper should be calculated")
  expect_true(!is.na(intervals_data$ReducibleLower), "Reducible Lower should be calculated")
  expect_true(intervals_data$ReducibleUpper >= 0 && intervals_data$ReducibleUpper <= 1,
              "Reducible Upper should be between 0 and 1")
  expect_true(intervals_data$ReducibleLower >= 0 && intervals_data$ReducibleLower <= 1,
              "Reducible Lower should be between 0 and 1")

  # Check that formula includes covariate
  expect_true(grepl("0.9 \\* b3", summary_text) || grepl("b3", summary_text),
              "Summary should mention the covariate")
  expect_true(grepl("b2", summary_text),
              "Summary should mention the fixed effect")

  cat("✓ TEST 2 PASSED: bayesian_statistical_intervals() works correctly with covariates\n\n")

  return(result)
})

# ============================================================================
# TEST 3: Backward Compatibility (Without Covariates)
# ============================================================================

cat("=== TEST 3: Backward Compatibility (No Covariates) ===\n")

test_that("Functions work without covariates (backward compatibility)", {

  # Call without covariates
  result_no_cov <- bayesian_statistical_intervals(
    data = toy_data,
    fixed_effects = test_params$fixed_effects,
    random_params = test_params$random_params,
    by = test_params$by,
    tolerance_interval_level = test_params$tolerance_interval_level,
    multiplication_factor = test_params$multiplication_factor,
    credible_interval_level = test_params$credible_interval_level,
    prediction_interval_level = test_params$prediction_interval_level,
    credibility_of_tolerance_interval = test_params$credibility_of_tolerance_interval,
    log_normal = FALSE
    # NO covariate_cols or covariate_values
  )

  cat("\n--- Results WITHOUT Covariates ---\n")
  print(result_no_cov$data)
  cat("\n")

  # Should still work and produce valid results
  expect_true(is.list(result_no_cov), "Should return valid results without covariates")
  expect_equal(nrow(result_no_cov$data), 1, info = "Should have 1 row")

  # Results should be different from covariate version
  result_with_cov <- bayesian_statistical_intervals(
    data = toy_data,
    fixed_effects = test_params$fixed_effects,
    random_params = test_params$random_params,
    by = test_params$by,
    tolerance_interval_level = test_params$tolerance_interval_level,
    multiplication_factor = test_params$multiplication_factor,
    credible_interval_level = test_params$credible_interval_level,
    prediction_interval_level = test_params$prediction_interval_level,
    credibility_of_tolerance_interval = test_params$credibility_of_tolerance_interval,
    log_normal = FALSE,
    covariate_cols = test_params$covariate_cols,
    covariate_values = test_params$covariate_values
  )

  expect_false(isTRUE(all.equal(result_no_cov$data$Median, result_with_cov$data$Median)),
               "Results should differ when covariates are included")

  cat("✓ TEST 3 PASSED: Backward compatibility maintained\n\n")
})

# ============================================================================
# TEST 4: Multiple Covariates
# ============================================================================

cat("=== TEST 4: Multiple Covariates ===\n")

test_that("Functions work with multiple covariates", {

  result_multi_cov <- bayesian_statistical_intervals(
    data = toy_data,
    fixed_effects = "b2",
    random_params = "s2",
    by = "EGTPTNUM",
    tolerance_interval_level = 90,
    multiplication_factor = 1.9,
    credible_interval_level = 90,
    prediction_interval_level = 90,
    credibility_of_tolerance_interval = 80,
    log_normal = FALSE,
    covariate_cols = c("b3", "b4"),  # Two covariates
    covariate_values = c(0.9, 1.5)
  )

  cat("\n--- Results WITH Multiple Covariates (b3=0.9, b4=1.5) ---\n")
  print(result_multi_cov$data)
  cat("\n--- Summary ---\n")
  cat(result_multi_cov$summary, "\n\n")

  expect_true(is.list(result_multi_cov), "Should handle multiple covariates")
  expect_equal(nrow(result_multi_cov$data), 1, info = "Should have 1 row")

  # Check that formula includes both covariates
  expect_true(grepl("b3", result_multi_cov$summary) && grepl("b4", result_multi_cov$summary),
              "Summary should mention both covariates")

  cat("✓ TEST 4 PASSED: Multiple covariates work correctly\n\n")
})

# ============================================================================
# TEST 5: Error Handling
# ============================================================================

cat("=== TEST 5: Error Handling ===\n")

test_that("Functions handle errors appropriately", {

  # Test 5a: Mismatched covariate lengths
  expect_error(
    bayesian_statistical_intervals(
      data = toy_data,
      fixed_effects = "b2",
      random_params = "s2",
      by = "EGTPTNUM",
      tolerance_interval_level = 90,
      covariate_cols = c("b3", "b4"),
      covariate_values = c(0.9)  # Only one value for two columns
    ),
    "Length of 'covariate_cols' must match length of 'covariate_values'"
  )
  cat("✓ Test 5a PASSED: Detects mismatched covariate lengths\n")

  # Test 5b: Non-existent covariate column
  expect_error(
    bayesian_statistical_intervals(
      data = toy_data,
      fixed_effects = "b2",
      random_params = "s2",
      by = "EGTPTNUM",
      tolerance_interval_level = 90,
      covariate_cols = "nonexistent_column",
      covariate_values = 1.0
    ),
    "Covariate columns not found in data"
  )
  cat("✓ Test 5b PASSED: Detects non-existent covariate columns\n")

  # Test 5c: Non-numeric covariate values
  expect_error(
    bayesian_statistical_intervals(
      data = toy_data,
      fixed_effects = "b2",
      random_params = "s2",
      by = "EGTPTNUM",
      tolerance_interval_level = 90,
      covariate_cols = "b3",
      covariate_values = "not_a_number"
    ),
    "'covariate_values' must be numeric"
  )
  cat("✓ Test 5c PASSED: Detects non-numeric covariate values\n\n")
})

# ============================================================================
# FINAL SUMMARY
# ============================================================================
# ============================================================================
# FINAL SUMMARY
# ============================================================================

cat("\n")
cat(paste(rep("=", 70), collapse = ""), "\n")
cat("=== ALL TESTS COMPLETED SUCCESSFULLY ===\n")
cat(paste(rep("=", 70), collapse = ""), "\n")
cat("\n")
cat("Summary:\n")
cat("✓ calculate_predicted_distributions() works with covariates\n")
cat("✓ bayesian_statistical_intervals() works with covariates\n")
cat("✓ Results match expected values within tolerance\n")
cat("✓ Backward compatibility maintained (works without covariates)\n")
cat("✓ Multiple covariates supported\n")
cat("✓ Error handling works correctly\n")
cat("\n")
cat("The updated functions are ready for production use!\n")