context("gp")

skip_on_cran()

test_that("gp_to_s is working properly for unidimensional gps", {
  # All true gp() terms should be changed to s() with k = k+1
  formula <- y ~ s(series) + gp(banana, k = 3, scale = FALSE) +
    infect:you + gp(hardcourt, k = 3)
  dat <- data.frame(y = rnorm(10),
                    series = rnorm(10),
                    banana = rnorm(10),
                    infect = rnorm(10),
                    you = rnorm(10),
                    hardcourt = rnorm(10),
                    gp = rnorm(10))

  # Check that the brms_mock formula is correctly remade
  gp_atts <- mvgam:::get_gp_attributes(formula,
                  data = dat,
                  family = gaussian())

  # scale should be passed; gr is always false and cmc is always true
  expect_true(identical(
    attr(terms(attr(gp_atts, 'gp_formula')),
         'term.labels')[1],
    "gp(banana, k = 3, cov = \"exp_quad\", iso = TRUE, scale = FALSE, c = 1.25, gr = FALSE, cmc = TRUE)")
    )

  expect_true(identical(
    attr(terms(attr(gp_atts, 'gp_formula')),
         'term.labels')[2],
    "gp(hardcourt, k = 3, cov = \"exp_quad\", iso = TRUE, scale = TRUE, c = 1.25, gr = FALSE, cmc = TRUE)")
  )

  expect_equal(attr(terms(mvgam:::gp_to_s(formula,
                                          data = dat,
                                          family = gaussian()),
                          keep.order = TRUE),
                    'term.labels'),
               attr(terms(formula(y ~ s(series) +
                                    s(banana, k = 4) +
                                    infect:you +
                                    s(hardcourt, k = 4)),
                          keep.order = TRUE),
                    'term.labels'))

  # Characters that match to 'gp' should not be changed
  formula <- y ~ gp(hardcourt, k = 3) + s(gp, k = 3)
  expect_equal(attr(terms(mvgam:::gp_to_s(formula,
                                          data = dat,
                                          family = gaussian()), keep.order = TRUE),
                    'term.labels'),
               attr(terms(formula(y ~ s(hardcourt, k = 4) + s(gp, k = 3)),
                          keep.order = TRUE),
                    'term.labels'))
})

test_that("gp_to_s is working properly for multidimensional gps", {
  # All true gp() terms should be changed to s() with k = k+1
  formula <- y ~ s(series) + gp(banana, hardcourt, k = 3, iso = FALSE,
                                c = 1.33, cov = 'matern52') +
    infect:you
  dat <- data.frame(y = rnorm(10),
                    series = rnorm(10),
                    banana = rnorm(10),
                    infect = rnorm(10),
                    you = rnorm(10),
                    hardcourt = rnorm(10),
                    gp = rnorm(10))
  gp_atts <- mvgam:::get_gp_attributes(formula,
                                       data = dat,
                                       family = gaussian())

  expect_true(identical(
    attr(terms(attr(gp_atts, 'gp_formula')),
         'term.labels')[1],
    "gp(banana, hardcourt, k = 3, cov = \"matern52\", iso = FALSE, scale = TRUE, c = 1.33, gr = FALSE, cmc = TRUE)")
  )

  expect_equal(attr(terms(mvgam:::gp_to_s(formula,
                                          data = dat,
                                          family = gaussian()),
                          keep.order = TRUE),
                    'term.labels'),
               attr(terms(formula(y ~ s(series) +
                                    ti(banana, hardcourt, k = 3, mc = c(0,0)) +
                                    infect:you),
                          keep.order = TRUE),
                    'term.labels'))
})

test_that("unidimensional gp for observation models working properly", {
  gaus_data$data_train$y[is.na(gaus_data$data_train$y)] <- 0
  mod <- mvgam(formula = y ~ s(series, bs = 're') +
                 gp(time, by = series, k = 10, c = 5/4) +
                 year:season,
               data = gaus_data$data_train,
               family = gaussian(),
               run_model = FALSE)

  expect_true(
    any(grepl('b[b_idx_gp_time_byseriesseries_3] = sqrt(spd_gp_exp_quad(l_gp_time_byseriesseries_3',
              mod$model_file, fixed = TRUE))
  )

  # Gp data structures should be in the model_data
  expect_true("l_gp_time_byseriesseries_1" %in% names(mod$model_data))
  expect_true("b_idx_gp_time_byseriesseries_1" %in% names(mod$model_data))
  expect_true("k_gp_time_byseriesseries_1" %in% names(mod$model_data))

  # These should match to the eigenvalues and eigenfunctions created by
  # a similar brms call
  brms_dat <- suppressWarnings(brms::make_standata(y ~ s(series, bs = 're') +
                                                     gp(time, by = series, k = 10,
                                                        c = 5/4) +
                                                     year:season,
                                                   data = gaus_data$data_train,
                                                   family = gaussian()))
  # Eigenvalues should be identical
  expect_true(all.equal(as.vector(brms_dat$slambda_1_1),
                        as.vector(mod$model_data$l_gp_time_byseriesseries_1)))

  # Eigenfunctions will be nearly identical
  row_s1 <- which(gaus_data$data_train$series == 'series_1' &
                    !is.na(gaus_data$data_train$y))
  col_s1 <- grep('gp(time):seriesseries_1', names(coef(mod$mgcv_model)),
                 fixed = TRUE)

  expect_true(identical(dim(brms_dat$Xgp_1_1),
                        dim(mod$model_data$X[row_s1,col_s1])))

  expect_true(all(unlist(lapply(seq_len(NCOL(brms_dat$Xgp_1_1)), function(x){
    cor(brms_dat$Xgp_1_1[,x], mod$model_data$X[row_s1,col_s1][,x])
  }), use.names = FALSE) > 0.99))

  # The mgcv model formula should contain s() in place of gp()
  expect_equal(attr(terms(mod$mgcv_model$formula, keep.order = TRUE),
                    'term.labels'),
               attr(terms(formula(y ~ s(time, by = series, k = 11) +
                                    year:season +
                                    s(series, bs = "re")),
                          keep.order = TRUE),
                    'term.labels'))
})

test_that("multidimensional gp for observation models working properly", {
  gaus_data$data_train$y[is.na(gaus_data$data_train$y)] <- 0
  mod <- mvgam(y ~ s(series, bs = 're') +
                 gp(time, year, k = 4, cov = 'matern32'),
               data = gaus_data$data_train,
               family = gaussian(),
               run_model = FALSE)

  expect_true(
    any(grepl('b[b_idx_gp_timeby_year_] = sqrt(spd_gp_matern32(l_gp_timeby_year_',
              mod$model_file, fixed = TRUE))
  )

  # Gp data structures should be in the model_data
  expect_true("l_gp_timeby_year_" %in% names(mod$model_data))
  expect_true("b_idx_gp_timeby_year_" %in% names(mod$model_data))
  expect_true("k_gp_timeby_year_" %in% names(mod$model_data))

  # These should match to the eigenvalues and eigenfunctions created by
  # a similar brms call
  brms_dat <- suppressWarnings(brms::make_standata(
    y ~ s(series, bs = 're') +
      gp(time, year, k = 4, gr = FALSE),
    data = gaus_data$data_train,
    family = gaussian())
    )

  # Eigenvalues should be identical
  expect_true(all.equal(brms_dat$slambda_1,
                        mod$model_data$l_gp_timeby_year_))

  # Eigenfunctions will be nearly identical
  col_s1 <- grep('gp(time,year)', names(coef(mod$mgcv_model)),
                 fixed = TRUE)

  expect_true(identical(dim(brms_dat$Xgp_1),
                        dim(mod$model_data$X[,col_s1])))

  expect_true(all(unlist(lapply(seq_len(NCOL(brms_dat$Xgp_1)), function(x){
    cor(brms_dat$Xgp_1[,x], mod$model_data$X[,col_s1][,x])
  }), use.names = FALSE) > 0.99))

  # The mgcv model formula should contain s() in place of gp()
  expect_equal(attr(terms(mod$mgcv_model$formula, keep.order = TRUE),
                    'term.labels'),
               attr(terms(formula(y ~ ti(time, year, k = 4, mc = c(0, 0)) +
                                    s(series, bs = "re")),
                          keep.order = TRUE),
                    'term.labels'))
})

test_that("noncentring with gp terms working properly", {
  mod <- mvgam(y ~ s(series, bs = 're') +
                 s(season, bs = 'cc', k = 8) +
                 gp(time, by = series, k = 10),
               trend_model = RW(),
               data = gaus_data$data_train,
               newdata = gaus_data$data_test,
               family = gaussian(),
               noncentred = TRUE,
               run_model = FALSE)

  # Model file should have the non-centred trend parameterisation now
  expect_true(
    any(grepl(trimws("trend = trend_raw .* rep_matrix(sigma', rows(trend_raw));"),
              trimws(mod$model_file),
              fixed = TRUE))
  )

  expect_true(
    any(grepl(trimws("trend[2 : n, s] += trend[1 : (n - 1), s];"),
              trimws(mod$model_file),
              fixed = TRUE))
  )

  # Gp data structures should be in the model_data
  expect_true("l_gp_time_byseriesseries_1" %in% names(mod$model_data))
  expect_true("b_idx_gp_time_byseriesseries_1" %in% names(mod$model_data))
  expect_true("k_gp_time_byseriesseries_1" %in% names(mod$model_data))
})

test_that("unidimensional gp for process models working properly", {
  mod <- mvgam(y ~ s(series, bs = 're'),
               trend_formula = ~
                 gp(time, by = trend, k = 10) +
                 year:season,
               data = beta_data$data_train,
               family = betar(),
               trend_model = AR(),
               run_model = FALSE)

  # Model file should have prior lines for gp terms
  expect_true(any(grepl('// prior for gp(time):trendtrend1_trend...',
                        mod$model_file, fixed = TRUE)))

  expect_true(any(grepl("b_trend[b_trend_idx_gp_time_bytrendtrend1] = sqrt(spd_gp_exp_quad(",
                        mod$model_file, fixed = TRUE)))

  # Gp data structures should be in the model_data
  expect_true("l_gp_trend_time_bytrendtrend1" %in% names(mod$model_data))
  expect_true("b_trend_idx_gp_time_bytrendtrend1" %in% names(mod$model_data))
  expect_true("k_gp_trend_time_bytrendtrend1" %in% names(mod$model_data))

  # These should match to the eigenvalues and eigenfunctions created by
  # a similar brms call
  brms_dat <- suppressWarnings(brms::make_standata(y ~ gp(time, by = series, k = 10,
                                                        c = 5/4),
                                                   data = beta_data$data_train,
                                                   family = gaussian()))
  # Eigenvalues should be identical
  expect_true(all.equal(as.vector(brms_dat$slambda_1_1),
                        as.vector(mod$model_data$l_gp_trend_time_bytrendtrend1)))

  # Eigenfunctions will be nearly identical
  row_s1 <- mod$model_data$ytimes_trend[,1]
  col_s1 <- grep('gp(time):trendtrend1', names(coef(mod$trend_mgcv_model)),
                 fixed = TRUE)

  expect_true(identical(dim(brms_dat$Xgp_1_1),
                        dim(mod$model_data$X_trend[row_s1,col_s1])))

  expect_true(max(abs(brms_dat$Xgp_1_1 -
                        mod$model_data$X_trend[row_s1,col_s1])) < 0.01)

  # The mgcv model formula should contain s() in place of gp()
  expect_equal(attr(terms(mod$trend_mgcv_model$formula, keep.order = TRUE),
                    'term.labels'),
               attr(terms(formula(y ~ s(time, by = series, k = 11) +
                                    year:season),
                          keep.order = TRUE),
                    'term.labels'))
})

test_that("multidimensional gp for process models working properly", {
  mod <- mvgam(y ~ s(series, bs = 're'),
               trend_formula = ~
                 gp(time, season, k = 10, iso = FALSE),
               data = beta_data$data_train,
               family = betar(),
               trend_model = AR(),
               run_model = FALSE)

  # Model file should have prior lines for gp terms
  expect_true(any(grepl('// prior for gp(time,season)_trend...',
                        mod$model_file, fixed = TRUE)))

  expect_true(any(grepl("b_trend[b_trend_idx_gp_timeby_season_] = sqrt(spd_gp_exp_quad(",
                        mod$model_file, fixed = TRUE)))

  expect_true(any(grepl("array[1] vector<lower=0>[2] rho_gp_trend_timeby_season_;",
                        mod$model_file, fixed = TRUE)))

  expect_true(any(grepl("rho_gp_trend_timeby_season_[1][1] ~ inv_gamma",
                        mod$model_file, fixed = TRUE)))

  expect_true(any(grepl("rho_gp_trend_timeby_season_[1][2] ~ inv_gamma",
                        mod$model_file, fixed = TRUE)))

  # Gp data structures should be in the model_data
  expect_true("l_gp_trend_timeby_season_" %in% names(mod$model_data))
  expect_true("b_trend_idx_gp_timeby_season_" %in% names(mod$model_data))
  expect_true("k_gp_trend_timeby_season_" %in% names(mod$model_data))

  # These should match to the eigenvalues and eigenfunctions created by
  # a similar brms call
  brms_dat <- suppressWarnings(brms::make_standata(y ~ gp(time, season, k = 10,
                                                          c = 5/4, cmc = TRUE,
                                                          gr = FALSE),
                                                   data = beta_data$data_train,
                                                   family = gaussian()))
  # Eigenvalues should be identical
  expect_true(all.equal(brms_dat$slambda_1,
                        mod$model_data$l_gp_trend_timeby_season_))

  # Eigenfunctions will be nearly identical
  col_s1 <- grep('gp(time,season)', names(coef(mod$trend_mgcv_model)),
                 fixed = TRUE)

  expect_true(identical(dim(brms_dat$Xgp_1),
                        dim(mod$model_data$X_trend[,col_s1])))

  expect_true(all(unlist(lapply(seq_len(NCOL(brms_dat$Xgp_1)), function(x){
    cor(brms_dat$Xgp_1[,x], mod$model_data$X_trend[,col_s1][,x])
  }), use.names = FALSE) > 0.99))

  # The mgcv model formula should contain s() in place of gp()
  expect_equal(attr(terms(mod$trend_mgcv_model$formula, keep.order = TRUE),
                    'term.labels'),
               attr(terms(formula(y ~ ti(time, season, k = 10, mc = c(0, 0))),
                          keep.order = TRUE),
                    'term.labels'))
})