library(testthat)
library(BioMixModel)


# ---------------------------------------------------------
# Formula validation
# ---------------------------------------------------------

test_that("spatial model validates formula", {
  
  skip_if_not_installed("nlme")
  
  dat <- data.frame(
    y = rnorm(20),
    x = rnorm(20)
  )
  
  expect_error(
    fit_spatial_mix(
      formula = "y ~ x",
      data = dat,
      correlation = nlme::corExp()
    ),
    "formula must be a formula"
  )
})


# ---------------------------------------------------------
# Data validation
# ---------------------------------------------------------

test_that("spatial model validates data", {
  
  skip_if_not_installed("nlme")
  
  expect_error(
    fit_spatial_mix(
      formula = y ~ x,
      data = list(
        y = rnorm(20),
        x = rnorm(20)
      ),
      correlation = nlme::corExp()
    ),
    "data must be a data frame"
  )
})


# ---------------------------------------------------------
# Correlation validation
# ---------------------------------------------------------

test_that("spatial model requires correlation", {
  
  skip_if_not_installed("nlme")
  
  dat <- data.frame(
    y = rnorm(20),
    x = rnorm(20)
  )
  
  expect_error(
    fit_spatial_mix(
      formula = y ~ x,
      data = dat
    ),
    "correlation.*supplied"
  )
})


test_that("spatial model validates correlation type", {
  
  skip_if_not_installed("nlme")
  
  dat <- data.frame(
    y = rnorm(20),
    x = rnorm(20)
  )
  
  expect_error(
    fit_spatial_mix(
      formula = y ~ x,
      data = dat,
      correlation = "invalid"
    ),
    "correlation.*nlme"
  )
})


# ---------------------------------------------------------
# Exponential spatial correlation
# ---------------------------------------------------------

test_that("spatial mixed model fits with exponential correlation", {
  
  skip_if_not_installed("nlme")
  
  set.seed(123)
  
  dat <- data.frame(
    id = factor(rep(1:5, each = 6)),
    y = rnorm(30),
    x = rnorm(30),
    time = rep(seq(0, 1, length.out = 6), 5)
  )
  
  m <- fit_spatial_mix(
    y ~ x,
    data = dat,
    correlation = nlme::corExp(
      form = ~time | id
    ),
    random = ~1 | id
  )
  
  expect_s3_class(m, "biomix_spatial")
  expect_s3_class(m, "biomix")
  expect_true(is.list(m))
  expect_true(inherits(m$fit, "lme"))
  expect_true(is.finite(AIC(m$fit)))
  expect_equal(m$model_type, "Spatial mixed model")
})


# ---------------------------------------------------------
# Gaussian spatial correlation
# ---------------------------------------------------------

test_that("spatial mixed model fits with Gaussian correlation", {
  
  skip_if_not_installed("nlme")
  
  set.seed(456)
  
  dat <- data.frame(
    id = factor(rep(1:5, each = 6)),
    y = rnorm(30),
    x = rnorm(30),
    time = rep(seq(0, 1, length.out = 6), 5)
  )
  
  m <- fit_spatial_mix(
    y ~ x,
    data = dat,
    correlation = nlme::corGaus(
      form = ~time | id
    ),
    random = ~1 | id
  )
  
  expect_s3_class(m, "biomix_spatial")
  expect_s3_class(m, "biomix")
  expect_true(inherits(m$fit, "lme"))
  expect_true(is.finite(AIC(m$fit)))
})


# ---------------------------------------------------------
# Spherical spatial correlation
# ---------------------------------------------------------

test_that("spatial mixed model fits with spherical correlation", {
  
  skip_if_not_installed("nlme")
  
  set.seed(789)
  
  dat <- data.frame(
    id = factor(rep(1:5, each = 6)),
    y = rnorm(30),
    x = rnorm(30),
    time = rep(seq(0, 1, length.out = 6), 5)
  )
  
  m <- fit_spatial_mix(
    y ~ x,
    data = dat,
    correlation = nlme::corSpher(
      form = ~time | id
    ),
    random = ~1 | id
  )
  
  expect_s3_class(m, "biomix_spatial")
  expect_s3_class(m, "biomix")
  expect_true(inherits(m$fit, "lme"))
  expect_true(is.finite(AIC(m$fit)))
})


# ---------------------------------------------------------
# Multiple correlation structures
# ---------------------------------------------------------

test_that("spatial model supports multiple correlation structures", {
  
  skip_if_not_installed("nlme")
  
  set.seed(123)
  
  dat <- data.frame(
    id = factor(rep(1:5, each = 6)),
    y = rnorm(30),
    x = rnorm(30),
    time = rep(seq(0, 1, length.out = 6), 5)
  )
  
  cor_structures <- list(
    nlme::corExp(form = ~time | id),
    nlme::corGaus(form = ~time | id),
    nlme::corSpher(form = ~time | id)
  )
  
  for (cor_structure in cor_structures) {
    
    m <- fit_spatial_mix(
      y ~ x,
      data = dat,
      correlation = cor_structure,
      random = ~1 | id
    )
    
    expect_s3_class(m, "biomix_spatial")
    expect_s3_class(m, "biomix")
    expect_true(inherits(m$fit, "lme"))
  }
})


# ---------------------------------------------------------
# Additional arguments
# ---------------------------------------------------------

test_that("spatial model passes additional arguments", {
  
  skip_if_not_installed("nlme")
  
  set.seed(123)
  
  dat <- data.frame(
    id = factor(rep(1:5, each = 6)),
    y = rnorm(30),
    x = rnorm(30),
    time = rep(seq(0, 1, length.out = 6), 5)
  )
  
  m <- fit_spatial_mix(
    y ~ x,
    data = dat,
    correlation = nlme::corExp(
      form = ~time | id
    ),
    random = ~1 | id,
    method = "ML"
  )
  
  expect_s3_class(m, "biomix_spatial")
  expect_equal(m$fit$method, "ML")
})


# ---------------------------------------------------------
# Stored components
# ---------------------------------------------------------

test_that("spatial model stores supplied components", {
  
  skip_if_not_installed("nlme")
  
  set.seed(123)
  
  dat <- data.frame(
    id = factor(rep(1:4, each = 5)),
    y = rnorm(20),
    x = rnorm(20),
    time = rep(seq(0, 1, length.out = 5), 4)
  )
  
  cor_structure <- nlme::corExp(
    form = ~time | id
  )
  
  m <- fit_spatial_mix(
    y ~ x,
    data = dat,
    correlation = cor_structure,
    random = ~1 | id
  )
  
  expect_equal(m$formula, y ~ x)
  expect_identical(m$data, dat)
  expect_true(inherits(m$correlation, "corExp"))
  expect_equal(m$model_type, "Spatial mixed model")
})