context("Tests for family functions")

test_that("distributions work correctly", {
  fam <- tweedie()
  expect_true(inherits(fam, 'family'))
  expect_true(inherits(fam, 'extended.family'))
  expect_true(fam$link == 'log')

  fam <- student_t()
  expect_true(inherits(fam, 'family'))
  expect_true(inherits(fam, 'extended.family'))
  expect_true(fam$link == 'identity')

  fam <- nmix()
  expect_true(inherits(fam, 'family'))
  expect_true(inherits(fam, 'extended.family'))
  expect_true(fam$link == 'log')
})

test_that("nmix predictions work correctly", {
  set.seed(1)
  Xp <- matrix(rnorm(100), ncol = 10, nrow = 10)
  attr(Xp, 'model.offset') <- 0
  family <- 'nmix'
  betas <- rnorm(10)
  latent_lambdas <- runif(10, 2, 5)
  cap <- rep(12, 10)
  family_pars <- list()
  expect_true(all(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = latent_lambdas,
                                        cap = cap,
                                        type = 'link',
                                        family_pars = family_pars) > 0))

  expect_true(all(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = latent_lambdas,
                                        cap = cap,
                                        type = 'expected',
                                        family_pars = family_pars) > 0))

  detects <- mvgam:::mvgam_predict(Xp = Xp,
                                   family = family,
                                   betas = betas,
                                   latent_lambdas = latent_lambdas,
                                   cap = cap,
                                   type = 'detection',
                                   family_pars = family_pars)
  expect_true(all(detects > 0) & all(detects < 1))

  truth <- rpois(10, 6)
  expect_no_error(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = latent_lambdas,
                                        cap = cap,
                                        type = 'link',
                                        density = TRUE,
                                        truth = truth,
                                        family_pars = family_pars))

  # Test that the residual calculation is working
  preds <- runif(10, 0.1, 10)
  N <- rpois(10, 20)
  p <- runif(10, 0.3, 0.6)
  resids <- mvgam:::ds_resids_nmix(truth = truth,
                                   fitted = 1,
                                   draw = 1,
                                   N = as.vector(N),
                                   p = as.vector(p))
  expect_true(all(!is.na(resids)))

})

test_that("beta predictions work correctly", {
  set.seed(1)
  Xp <- matrix(rnorm(100), ncol = 10, nrow = 10)
  attr(Xp, 'model.offset') <- 0
  family <- 'beta'
  betas <- rnorm(10)
  family_pars <- list(phi = rep(1, 10))
  expect_no_error(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'link',
                                        family_pars = family_pars))

  expect_no_error(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'variance',
                                        family_pars = family_pars))
  expecteds <- mvgam:::mvgam_predict(Xp = Xp,
                                     family = family,
                                     betas = betas,
                                     latent_lambdas = NULL,
                                     cap = NULL,
                                     type = 'expected',
                                     family_pars = family_pars)

  expect_true(all(expecteds >= 0) & all(expecteds <= 1))

  preds <- mvgam:::mvgam_predict(Xp = Xp,
                                 family = family,
                                 betas = betas,
                                 latent_lambdas = NULL,
                                 cap = NULL,
                                 type = 'response',
                                 family_pars = family_pars)
  expect_true(all(preds >= 0) & all(preds <= 1))

  truth <- runif(10, 0.01, 0.99)
  expect_no_error(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'link',
                                        density = TRUE,
                                        truth = truth,
                                        family_pars = family_pars))

  preds <- runif(10, 0.01, 0.99)
  resids <- mvgam:::ds_resids_beta(truth = truth,
                                   fitted = as.vector(preds),
                                   draw = 1,
                                   precision = rep(1, 10))
  expect_true(all(!is.na(resids)))

})

test_that("beta-binomial predictions work correctly", {
  set.seed(1)
  Xp <- matrix(rnorm(100), ncol = 10, nrow = 10)
  attr(Xp, 'model.offset') <- 0
  family <- 'beta_binomial'
  betas <- rnorm(10)
  family_pars <- list(phi = rep(1, 10),
                      trials = rep(20, 10))
  expect_no_error(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'link',
                                        family_pars = family_pars))

  expect_no_error(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'variance',
                                        family_pars = family_pars))

  expect_true(all(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'expected',
                                        family_pars = family_pars) > 0))

  expect_true(all(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'response',
                                        family_pars = family_pars) >= 0))

  truth <- rpois(10, 8)
  expect_no_error(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'link',
                                        density = TRUE,
                                        truth = truth,
                                        family_pars = family_pars))

  p <- runif(10, 0.3, 0.6)
  resids <- mvgam:::ds_resids_beta_binomial(truth = truth,
                                            fitted = p,
                                            draw = 1,
                                            N = rep(20, 10),
                                            phi = rep(1, 10))
  expect_true(all(!is.na(resids)))
})

test_that("negative binomial predictions work correctly", {
  set.seed(1)
  Xp <- matrix(rnorm(100), ncol = 10, nrow = 10)
  attr(Xp, 'model.offset') <- 0
  family <- 'negative binomial'
  betas <- rnorm(10)
  family_pars <- list(phi = rep(1, 10))
  expect_no_error(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'link',
                                        family_pars = family_pars))

  expect_no_error(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'variance',
                                        family_pars = family_pars))

  expect_true(all(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'expected',
                                        family_pars = family_pars) > 0))

  expect_true(all(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'response',
                                        family_pars = family_pars) >= 0))

  truth <- rpois(10, 8)
  expect_no_error(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'link',
                                        density = TRUE,
                                        truth = truth,
                                        family_pars = family_pars))

  preds <- rpois(10, 8)
  resids <- mvgam:::ds_resids_nb(truth = truth,
                                 fitted = preds,
                                 draw = 1,
                                 size = rep(1, 10))
  expect_true(all(!is.na(resids)))
})

test_that("lognormal predictions work correctly", {
  set.seed(1)
  Xp <- matrix(rnorm(100), ncol = 10, nrow = 10)
  attr(Xp, 'model.offset') <- 0
  family <- 'lognormal'
  betas <- rnorm(10)
  family_pars <- list(sigma_obs = 1)
  expect_no_error(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'link',
                                        family_pars = family_pars))

  expect_no_error(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'variance',
                                        family_pars = family_pars))

  expect_true(all(mvgam:::mvgam_predict(Xp = Xp,
                                    family = family,
                                    betas = betas,
                                    latent_lambdas = NULL,
                                    cap = NULL,
                                    type = 'expected',
                                    family_pars = family_pars) > 0))

  expect_true(all(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'response',
                                        family_pars = family_pars) > 0))

  truth <- runif(10, 0.1, 20)
  expect_no_error(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'link',
                                        density = TRUE,
                                        truth = truth,
                                        family_pars = family_pars))

  preds <- runif(10, 0.1, 20)
  resids <- mvgam:::ds_resids_lnorm(truth = truth,
                                   fitted = preds,
                                   draw = 1,
                                   sigma = rep(1, 10))
  expect_true(all(!is.na(resids)))
})

test_that("gamma predictions work correctly", {
  set.seed(1)
  Xp <- matrix(rnorm(100), ncol = 10, nrow = 10)
  attr(Xp, 'model.offset') <- 0
  family <- 'Gamma'
  betas <- rnorm(10)
  family_pars <- list(shape = 1)
  expect_no_error(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'link',
                                        family_pars = family_pars))

  expect_no_error(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'variance',
                                        family_pars = family_pars))

  expect_true(all(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'expected',
                                        family_pars = family_pars) > 0))

  expect_true(all(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'response',
                                        family_pars = family_pars) > 0))

  truth <- runif(10, 0.1, 20)
  expect_no_error(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'link',
                                        density = TRUE,
                                        truth = truth,
                                        family_pars = family_pars))

  preds <- runif(10, 0.1, 20)
  resids <- mvgam:::ds_resids_gamma(truth = truth,
                                   fitted = preds,
                                   draw = 1,
                                   shape = rep(1, 10))
  expect_true(all(!is.na(resids)))
})

test_that("student-t predictions work correctly", {
  set.seed(1)
  Xp <- matrix(rnorm(100), ncol = 10, nrow = 10)
  attr(Xp, 'model.offset') <- 0
  family <- 'student'
  betas <- rnorm(10)
  family_pars <- list(sigma_obs = 1,
                      nu = 3)
  expect_no_error(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'link',
                                        family_pars = family_pars))

  expect_no_error(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'variance',
                                        family_pars = family_pars))

  expect_no_error(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'expected',
                                        family_pars = family_pars))

  expect_no_error(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'response',
                                        family_pars = family_pars))

  truth <- rnorm(10)
  expect_no_error(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'link',
                                        density = TRUE,
                                        truth = truth,
                                        family_pars = family_pars))

  preds <- rnorm(10)
  resids <- mvgam:::ds_resids_student(truth = truth,
                                     fitted = preds,
                                     draw = 1,
                                     sigma = rep(1, 10),
                                     nu = rep(5, 10))
  expect_true(all(!is.na(resids)))
})

test_that("tweedie predictions work correctly", {
  set.seed(1)
  Xp <- matrix(rnorm(100), ncol = 10, nrow = 10)
  attr(Xp, 'model.offset') <- 0
  family <- 'tweedie'
  betas <- rnorm(10)
  family_pars <- list(phi = 1)
  expect_no_error(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'link',
                                        family_pars = family_pars))

  expect_no_error(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'variance',
                                        family_pars = family_pars))

  expect_no_error(all(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'expected',
                                        family_pars = family_pars) > 0))

  expect_no_error(all(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'response',
                                        family_pars = family_pars) >= 0))

  truth <- rpois(10, 5)
  expect_no_error(mvgam:::mvgam_predict(Xp = Xp,
                                        family = family,
                                        betas = betas,
                                        latent_lambdas = NULL,
                                        cap = NULL,
                                        type = 'link',
                                        density = TRUE,
                                        truth = truth,
                                        family_pars = family_pars))

  preds <- rpois(10, 5)
  resids <- mvgam:::ds_resids_tw(truth = truth,
                                fitted = preds,
                                draw = 1)
  expect_true(all(!is.na(resids)))
})

# Skip actual model setups on CRAN as they take some time
test_that("family setups work correctly", {
  skip_on_cran()
  simdat <- sim_mvgam(family = poisson())
  mod <- mvgam(y ~ -1,
               trend_model = PW(),
               data = simdat$data_train,
               family = poisson(),
               run_model = FALSE)
  expect_true(inherits(mod, 'mvgam_prefit'))
  expect_no_error(capture_output(print(mod)))

  simdat <- sim_mvgam(family = nb(),
                      n_lv = 2)
  mod <- mvgam(y ~ -1,
               trend_model = PW(),
               data = simdat$data_train,
               family = nb(),
               run_model = FALSE)
  expect_true(inherits(mod, 'mvgam_prefit'))
  expect_no_error(capture_output(print(mod)))

  simdat <- sim_mvgam(family = lognormal(),
                      trend_model = VAR())
  mod <- mvgam(y ~ -1,
               trend_model = PW(),
               data = simdat$data_train,
               family = lognormal(),
               run_model = FALSE)
  expect_true(inherits(mod, 'mvgam_prefit'))
  expect_no_error(capture_output(print(mod)))

  simdat <- sim_mvgam(family = bernoulli(),
                      trend_model = VAR(cor = TRUE))
  mod <- mvgam(y ~ -1,
               trend_model = PW(),
               data = simdat$data_train,
               family = bernoulli(),
               run_model = FALSE)
  expect_true(inherits(mod, 'mvgam_prefit'))
  expect_no_error(capture_output(print(mod)))

  simdat <- sim_mvgam(family = student_t(),
                      trend_model = GP())
  mod <- mvgam(y ~ -1,
               trend_model = PW(),
               data = simdat$data_train,
               family = student_t(),
               run_model = FALSE)
  expect_true(inherits(mod, 'mvgam_prefit'))
  expect_no_error(capture_output(print(mod)))

  simdat <- sim_mvgam(family = Gamma(),
                      seasonality = 'shared',
                      trend_model = AR(p = 3))
  mod <- mvgam(y ~ -1,
               trend_model = PW(),
               data = simdat$data_train,
               family = Gamma(),
               run_model = FALSE)
  expect_true(inherits(mod, 'mvgam_prefit'))
  expect_no_error(capture_output(print(mod)))

  simdat <- sim_mvgam(family = betar(),
                      seasonality = 'hierarchical',
                      trend_model = AR(p = 2))
  mod <- mvgam(y ~ -1,
               trend_model = PW(),
               data = simdat$data_train,
               family = betar(),
               run_model = FALSE)
  expect_true(inherits(mod, 'mvgam_prefit'))
  expect_no_error(capture_output(print(mod)))
})