context("RW,AR,VAR,CAR")

skip_on_cran()

test_that("ma and cor options should work for trends other than VAR", {
  test <- mvgam(y ~ s(series, bs = 're') +
                       s(season, bs = 'cc') - 1,
                     trend_model = AR(p = 1, ma = TRUE),
                     data = gaus_data$data_train,
                     family = gaussian(),
                     run_model = FALSE)
  expect_true(inherits(test, 'mvgam_prefit'))
  expect_true(any(grepl('vector<lower=-1,upper=1>[n_series]theta;',
                        gsub(' ', '', test$model_file), fixed = TRUE)))
  expect_true(attr(test$model_data, 'trend_model') == 'AR1')

  # Correlation works
  test <- mvgam(y ~ s(series, bs = 're') +
                       s(season, bs = 'cc') - 1,
                     trend_model = AR(p = 1, cor = TRUE),
                     data = gaus_data$data_train,
                     family = gaussian(),
                     run_model = FALSE)
  expect_true(inherits(test, 'mvgam_prefit'))
  expect_true(any(grepl('error[i]~multi_normal_cholesky(trend_zeros,L_Sigma);',
                        gsub(' ', '', test$model_file), fixed = TRUE)))
  expect_true(attr(test$model_data, 'trend_model') == 'AR1')

  test <- mvgam(y ~ s(series, bs = 're') +
                       s(season, bs = 'cc') - 1,
                     trend_model = RW(ma = TRUE),
                     data = gaus_data$data_train,
                     family = gaussian(),
                     run_model = FALSE)
  expect_true(inherits(test, 'mvgam_prefit'))
  expect_true(any(grepl('vector<lower=-1,upper=1>[n_series]theta;',
                        gsub(' ', '', test$model_file), fixed = TRUE)))
  expect_true(attr(test$model_data, 'trend_model') == 'RW')

  test <- mvgam(y ~ s(series, bs = 're') +
                       s(season, bs = 'cc') - 1,
                     trend_model = RW(cor = TRUE),
                     data = gaus_data$data_train,
                     family = gaussian(),
                     run_model = FALSE)
  expect_true(inherits(test, 'mvgam_prefit'))
  expect_true(any(grepl('error[i]~multi_normal_cholesky(trend_zeros,L_Sigma);',
                        gsub(' ', '', test$model_file), fixed = TRUE)))
  expect_true(attr(test$model_data, 'trend_model') == 'RW')
})

test_that("VARMAs are set up correctly", {
  var <- mvgam(y ~ s(series, bs = 're') +
                   s(season, bs = 'cc') - 1,
                 trend_model = VAR(),
                 data = gaus_data$data_train,
                 family = gaussian(),
                 run_model = FALSE)
  expect_true(inherits(var, 'mvgam_prefit'))

  var <- SW(mvgam(y ~ s(series, bs = 're') +
                 gp(time, c = 5/4, k = 20) - 1,
               trend_model = VAR(),
               data = gaus_data$data_train,
               family = gaussian(),
               run_model = FALSE))
  expect_true(inherits(var, 'mvgam_prefit'))

  varma <- SW(mvgam(y ~ s(series, bs = 're') +
                   s(season, bs = 'cc') - 1,
                 trend_model = 'VARMA',
                 data = gaus_data$data_train,
                 family = gaussian(),
                 run_model = FALSE))

  expect_true(any(grepl('// unconstrained ma inverse partial autocorrelations',
                        varma$model_file, fixed = TRUE)))

  varma <- mvgam(y ~ s(series, bs = 're'),
                 trend_formula = ~ gp(time, by = trend, c = 5/4, k = 15),
                 trend_model = VAR(ma = TRUE),
                 data = gaus_data$data_train,
                 family = gaussian(),
                 run_model = FALSE)

  expect_true(any(grepl('// unconstrained ma inverse partial autocorrelations',
                        varma$model_file, fixed = TRUE)))
})

test_that("hierarchical cors are set up correctly", {
  # Simulate hierarchical data
  simdat1 <- sim_mvgam(trend_model = VAR(cor = TRUE),
                       prop_trend = 0.95,
                       n_series = 3,
                       mu = c(1, 2, 3))
  simdat2 <- sim_mvgam(trend_model = VAR(cor = TRUE),
                       prop_trend = 0.95,
                       n_series = 3,
                       mu = c(1, 2, 3))
  simdat3 <- sim_mvgam(trend_model = VAR(cor = TRUE),
                       prop_trend = 0.95,
                       n_series = 3,
                       mu = c(1, 2, 3))
  all_dat <- rbind(simdat1$data_train %>%
                     dplyr::mutate(region = 'qld'),
                   simdat2$data_train %>%
                     dplyr::mutate(region = 'nsw'),
                   simdat3$data_train %>%
                     dplyr::mutate(region = 'vic')) %>%
    dplyr::mutate(species = gsub('series', 'species', series),
                  species = as.factor(species),
                  region = as.factor(region)) %>%
    dplyr::arrange(series, time) %>%
    dplyr::select(-series)

  # Check that all model types can be constructed

  # Random Walk
  mod <- mvgam(formula = y ~ species,
               trend_model = RW(gr = region, subgr = species),
               data = all_dat,
               run_model = FALSE)
  expect_true(inherits(mod, 'mvgam_prefit'))
  expect_true(any(grepl('Sigma[group_inds[g], group_inds[g]] = multiply_lower_tri_self_transpose(L_Sigma_group[g]);',
                        mod$model_file, fixed = TRUE)))
  expect_true(any(grepl('// hierarchical process error correlations',
                        mod$model_file, fixed = TRUE)))

  mod <- mvgam(formula = y ~ -1,
               trend_formula = ~ species,
               trend_model = RW(gr = region, subgr = species),
               data = all_dat,
               run_model = FALSE)
  expect_true(inherits(mod, 'mvgam_prefit'))
  expect_true(any(grepl('Sigma[group_inds[g], group_inds[g]] = multiply_lower_tri_self_transpose(L_Sigma_group[g]);',
                        mod$model_file, fixed = TRUE)))
  expect_true(any(grepl('// hierarchical process error correlations',
                        mod$model_file, fixed = TRUE)))

  # AR
  mod <- mvgam(formula = y ~ species,
               trend_model = AR(gr = region, subgr = species, p = 2),
               data = all_dat,
               run_model = FALSE)
  expect_true(inherits(mod, 'mvgam_prefit'))
  expect_true(any(grepl('Sigma[group_inds[g], group_inds[g]] = multiply_lower_tri_self_transpose(L_Sigma_group[g]);',
                        mod$model_file, fixed = TRUE)))
  expect_true(any(grepl('// hierarchical process error correlations',
                        mod$model_file, fixed = TRUE)))

  mod <- mvgam(formula = y ~ -1,
               trend_formula = ~ species,
               trend_model = AR(gr = region, subgr = species, p = 3),
               data = all_dat,
               run_model = FALSE)
  expect_true(inherits(mod, 'mvgam_prefit'))
  expect_true(any(grepl('Sigma[group_inds[g], group_inds[g]] = multiply_lower_tri_self_transpose(L_Sigma_group[g]);',
                        mod$model_file, fixed = TRUE)))
  expect_true(any(grepl('// hierarchical process error correlations',
                        mod$model_file, fixed = TRUE)))

  # VAR
  mod <- mvgam(formula = y ~ species,
               trend_model = VAR(gr = region, subgr = species),
               data = all_dat,
               run_model = FALSE)
  expect_true(inherits(mod, 'mvgam_prefit'))
  expect_true(any(grepl('Sigma[group_inds[g], group_inds[g]] = multiply_lower_tri_self_transpose(L_Sigma_group[g]);',
                        mod$model_file, fixed = TRUE)))
  expect_true(any(grepl('// derived group-level VAR covariance matrices',
                        mod$model_file, fixed = TRUE)))

  mod <- mvgam(formula = y ~ -1,
               trend_formula = ~ species,
               trend_model = VAR(gr = region, subgr = species),
               data = all_dat,
               run_model = FALSE)
  expect_true(inherits(mod, 'mvgam_prefit'))
  expect_true(any(grepl('Sigma[group_inds[g], group_inds[g]] = multiply_lower_tri_self_transpose(L_Sigma_group[g]);',
                        mod$model_file, fixed = TRUE)))
  expect_true(any(grepl('// derived group-level VAR covariance matrices',
                        mod$model_file, fixed = TRUE)))
})

test_that("site variable must be numeric for ZMVN", {
  site_dat <- data.frame(site = as.character(rep(1:10, 4)),
                         species = as.factor(c(sort(rep(letters[1:4], 9)),
                                               'q', 'q', 'q', 'q')),
                         y = rpois(40, 3))

  trend_model <- ZMVN(unit = site, subgr = species)
  expect_error(mvgam:::validate_series_time(data = site_dat,
                                            trend_model = trend_model),
               'Variable "site" must be either numeric or integer type')

})

test_that("Each subgroup must exist within each site for ZMVN", {
  site_dat <- data.frame(site = rep(1:10, 4),
                         species = as.factor(c(sort(rep(letters[1:4], 9)),
                                               'q', 'q', 'q', 'q')),
                         y = rpois(40, 3))

  trend_model <- ZMVN(unit = site, subgr = species)
  expect_error(mvgam:::validate_series_time(data = site_dat,
                                             trend_model = trend_model),
               'One or more series in data is missing observations for one or more timepoints')

  # Should work if all species were recorded in all sites (even if NA)
  site_dat <- data.frame(site = rep(1:10, 4),
                         species = as.factor(c(sort(rep(letters[1:4], 10)))),
                         y = c(NA, rpois(39, 3)))

  trend_model <- ZMVN(unit = site, subgr = species)
  expect_no_error(mvgam:::validate_series_time(data = site_dat,
                                            trend_model = trend_model))

  mod <- mvgam(formula = y ~ species,
               trend_model = ZMVN(unit = site, subgr = species),
               data = site_dat,
               run_model = FALSE)
  expect_true(inherits(mod, 'mvgam_prefit'))
  expect_true(any(grepl('trend_raw[i] ~ multi_normal_cholesky(trend_zeros, L_Sigma);',
                        mod$model_file, fixed = TRUE)))
  expect_true(any(grepl('L_Sigma = diag_pre_multiply(sigma, L_Omega);',
                        mod$model_file, fixed = TRUE)))


  mod <- mvgam(formula = y ~ -1,
               trend_formula = ~ species,
               trend_model = ZMVN(unit = site, subgr = species),
               data = site_dat,
               run_model = FALSE)
  expect_true(inherits(mod, 'mvgam_prefit'))
  expect_true(any(grepl('LV_raw[i] ~ multi_normal_cholesky(trend_mus[ytimes_trend[i, 1 : n_lv]]',
                        mod$model_file, fixed = TRUE)))
  expect_true(any(grepl('L_Sigma = diag_pre_multiply(sigma, L_Omega);',
                        mod$model_file, fixed = TRUE)))
})


# Replicate CAR1 example
# Function to simulate CAR1 data with seasonality
sim_corcar1 = function(n = 120,
                       phi = 0.5,
                       sigma = 1,
                       sigma_obs = 0.75){
# Sample irregularly spaced time intervals
time_dis <- c(0, runif(n - 1, -0.1, 1))
time_dis[time_dis < 0] <- 0; time_dis <- time_dis * 5

# Set up the latent dynamic process
x <- vector(length = n); x[1] <- -0.3
for(i in 2:n){
  # zero-distances will cause problems in sampling, so mvgam uses a
  # minimum threshold; this simulation function emulates that process
  if(time_dis[i] == 0){
    x[i] <- rnorm(1, mean = (phi ^ 1e-12) * x[i - 1], sd = sigma)
   } else {
     x[i] <- rnorm(1, mean = (phi ^ time_dis[i]) * x[i - 1], sd = sigma)
   }
 }

# Add 12-month seasonality
 cov1 <- sin(2 * pi * (1 : n) / 12); cov2 <- cos(2 * pi * (1 : n) / 12)
 beta1 <- runif(1, 0.3, 0.7); beta2 <- runif(1, 0.2, 0.5)
 seasonality <- beta1 * cov1 + beta2 * cov2

# Take Gaussian observations with error and return
 data.frame(y = rnorm(n, mean = x + seasonality, sd = sigma),
            season = rep(1:12, 20)[1:n],
            time = cumsum(time_dis))
}

# Sample two time series
dat <- rbind(dplyr::bind_cols(sim_corcar1(phi = 0.65,
                                          sigma_obs = 0.55),
                              data.frame(series = 'series1')),
             dplyr::bind_cols(sim_corcar1(phi = 0.8,
                              sigma_obs = 0.35),
                              data.frame(series = 'series2'))) %>%
       dplyr::mutate(series = as.factor(series))

dat_train <- dat %>%
  dplyr::group_by(series) %>%
  dplyr::arrange(time) %>%
  dplyr::slice_head(n = 110) %>%
  dplyr::ungroup()
dat_test <- dat %>%
  dplyr::group_by(series) %>%
  dplyr::arrange(time) %>%
  dplyr::slice_tail(n = 10)%>%
  dplyr::ungroup()

test_that("CAR1 sets up correctly", {
  # mvgam with CAR(1) trends and series-level seasonal smooths
  mod <- mvgam(formula = y ~ s(season, bs = 'cc',
                               k = 5, by = series),
               trend_model = CAR(),
               data = dat_train,
               family = gaussian(),
               run_model = FALSE)
  expect_true(inherits(mod, 'mvgam_prefit'))
  expect_true(exists('time_dis', mod$model_data))
  expect_true(exists('index..time..index', mod$obs_data))
  expect_true(attr(mod$model_data, 'trend_model') == 'CAR1')
  expect_true(any(grepl('vector<lower=0,upper=1>[n_series]ar1;',
                        gsub(' ', '', mod$model_file), fixed = TRUE)))

  # Will work for regularly-spaced data as well
  mod <- mvgam(formula = y ~ s(season, bs = 'cc',
                               k = 5, by = series),
               trend_model = CAR(),
               data = gaus_data$data_train,
               family = gaussian(),
               run_model = FALSE)
  expect_true(inherits(mod, 'mvgam_prefit'))
  expect_true(exists('time_dis', mod$model_data))
  expect_true(all.equal(mean(mod$model_data$time_dis[2:NROW(mod$model_data$time_dis),]),
                        max(mod$model_data$time_dis[2:NROW(mod$model_data$time_dis),]),
                        min(mod$model_data$time_dis[2:NROW(mod$model_data$time_dis),]),
                        1L))
  expect_true(exists('index..time..index', mod$obs_data))
  expect_true(attr(mod$model_data, 'trend_model') == 'CAR1')
})