context("Test outbreaker")
current_R_version <- gsub("R version ([0-9.]+) .+$", "\\1", R.version.string)

## test output format ##
test_that("outbreaker's output have expected format", {
    ## skip on CRAN
    skip_on_cran()


    ## get data
    data(fake_outbreak)
    x <- fake_outbreak

    ## run outbreaker
    data <- list(dna = x$dna, dates = x$onset, w_dens = x$w)
    config <- list(n_iter = 10, sample_every = 1, find_import = FALSE)
    out <- outbreaker(data, config)


    out_df <- as.data.frame(out)

    data <- list(dates = x$onset, w_dens = x$w)
    out2 <- outbreaker(data, config)

    ## check output
    expect_is(out, "outbreaker_chains")
    expect_is(out_df, "data.frame")
    expect_equal(nrow(out), 10)
    expect_true(!any(is.na(out_df$post)))
    expect_true(all(out_df$post> -1e30))

})




## test convergence in various settings ##
test_that("results ok: DNA + time", {
    ## skip on CRAN
    skip_on_cran()


    ## get data
    data(fake_outbreak)
    x <- fake_outbreak

    ## outbreaker DNA + time ##
    ## analysis
    set.seed(1)
    data <- list(dna = x$dna, dates = x$onset, w_dens = x$w)
    config <- list(n_iter = 5000, sample_every = 50,
                   init_tree = "seqTrack", find_import = TRUE,
                   move_pi = FALSE)

    out <- outbreaker(data = data, config = config)

    ## checks
    out_smry <- summary(out, burnin = 1000)

    ## approx log post values
    expect_true(min(out_smry$post) > -1250)

    ## at least 80% ancestries correct
    temp <- mean(out_smry$tree$from==x$ances, na.rm = TRUE)
    expect_true(temp > .8)

    ## infection date within 3 days on average
    temp <- mean(abs(out_smry$tree$time - x$onset), na.rm = TRUE)
    expect_true(temp < 3.5)

    ## mu within reasonable ranges
    expect_true(out_smry$mu[2] > 1e-5 &&
                out_smry$mu[5] < 2e-4)

})





test_that("results ok: time, no DNA", {
    ## skip on CRAN
    skip_on_cran()


    ## get data
    data(fake_outbreak)
    x <- fake_outbreak

    ## outbreaker time, no DNA ##
    ## analysis
    set.seed(1)

    data <- list(dates = x$onset, w_dens = x$w)
    config <- list(n_iter = 5000, sample_every = 50,
                   init_tree="star", find_import = FALSE,
                   move_kappa = FALSE)

    out_no_dna <- outbreaker(data = data, config = config)

    ## checks
    out_no_dna.smry <- summary(out_no_dna, burnin = 1000)

    ## approx log post values
    expect_true(min(out_no_dna.smry$post) > -90)

    ## infection date within 3 days on average
    temp <- mean(abs(out_no_dna.smry$tree$time - x$onset), na.rm = TRUE)
    expect_true(temp < 3.5)

    ## check that support for ancestries is weak
    sup <- na.omit(out_no_dna.smry$tree$support)
    expect_lt(quantile(sup, .9), .55)
    expect_lt(mean(sup), .35)

})



test_that("results ok: time, ctd, DNA", {
    ## skip on CRAN
    skip_on_cran()


    ## get data
    data(fake_outbreak)
    x <- fake_outbreak

    ## outbreaker time, ctd, no DNA ##
    ## analysis
    ## set the random number generator kind to old R
    suppressWarnings(RNGversion("3.5.2"))
    set.seed(1)

    tTree <- data.frame(i = x$ances,
                        j = 1:length(x$ances))

    ctd <- sim_ctd(tTree, eps = 0.9, lambda = 0.1)

    data <- list(dates = x$onset, w_dens = x$w,
                 ctd = ctd[,1:2], dna = x$dna)
    config <- list(n_iter = 5000, sample_every = 50,
                   init_tree="star", find_import = FALSE,
                   move_kappa = FALSE)

    out_ctd <- outbreaker(data = data, config = config)

    ## checks
    out_ctd.smry <- summary(out_ctd, burnin = 1000)

    ## at least 90% ancestries correct
    temp <- mean(out_ctd.smry$tree$from==x$ances, na.rm = TRUE)
    expect_true(temp > .9)

    ## eps and lambda parameter estimates within reasonable ranges
    quant_eps <- quantile(out_ctd$eps, c(0.25, 0.75))
    expect_true(quant_eps[[1]] > 0.7 &
                quant_eps[[2]] < 0.95)

    quant_lambda <- quantile(out_ctd$lambda, c(0.25, 0.75))
    expect_true(quant_lambda[[1]] > 0.05 &
                quant_lambda[[2]] < 0.35)

    ## approx log post values
    expect_true(min(out_ctd.smry$post) > -1200)

    ## reset the Random number generator kind
    RNGversion(current_R_version)

})




test_that("results ok: easy run, no missing cases, no import", {
    ## skip on CRAN
    skip_on_cran()


    ## get data
    data(fake_outbreak)
    x <- fake_outbreak

    ## outbreaker, no missing cases ##
    ## analysis
    set.seed(1)
    config <-  list(n_iter = 5e3, sample_every = 50,
                    init_tree = "seqTrack", move_kappa = FALSE,
                    move_pi = FALSE, init_pi = 1,
                    find_import = FALSE)
    data <- list(dna = x$dna, dates = x$onset, w_dens = x$w)

    out_no_missing <- outbreaker(data = data, config = config)

    ## checks
    out_no_missing_smry <- summary(out_no_missing, burnin = 3500)

    ## approx log post values
    expect_true(min(out_no_missing_smry$post) > -935)

    ## at least 85% ancestries correct
    temp <- mean(out_no_missing_smry$tree$from==x$ances, na.rm = TRUE)
    expect_true(temp > .85)

    ## infection datewithin 3 days on average
    temp <- mean(abs(out_no_missing_smry$tree$time - x$onset), na.rm = TRUE)
    expect_true(temp < 3.5)

})



test_that("results ok: no missing cases, detect imported cases",
{
    ## skip on CRAN
    skip_on_cran()


    ## get data
    data(fake_outbreak)
    x <- fake_outbreak

    ## outbreaker, no missing cases, detect imported cases ##
    ## analysis
    set.seed(1)
    out_with_import <- outbreaker(data = list(dna = x$dna, dates = x$onset,
                                              w_dens = x$w),
                                  config = list(n_iter = 5000, sample_every = 100,
                                                init_tree="star",
                                                move_kappa = FALSE, move_pi = FALSE,
                                                init_pi = 1, find_import = TRUE)
                                  )

    ## checks
    out_with_import_smry <- summary(out_with_import, burnin = 500)
    out_with_import_smry$tree$from[is.na(out_with_import_smry$tree$from)] <- 0
    x$ances[is.na(x$ances)] <- 0

    ## approx log post values
    expect_true(min(out_with_import_smry$post) > -460)

    ## at least 80% ancestries correct
    temp <- mean(out_with_import_smry$tree$from==x$ances, na.rm = TRUE)
    expect_true(temp >= .80)

    ## infection datewithin 3 days on average
    temp <- mean(abs(out_with_import_smry$tree$time - x$onset), na.rm = TRUE)
    expect_true(temp < 3.5)

})






test_that("results ok: kappa and pi", {
    ## skip on CRAN
    skip_on_cran()


    ## get data
    onset <- c(0, 2, 6, 13)
    w <- c(.25, .5, .25)

    ## outbreaker, no missing cases, detect imported cases ##
    ## analysis
    set.seed(1)

    data <- list(dates = onset, w_dens = w)
    config <- list(n_iter = 5000, sample_every = 50, init_tree = "star",
                   move_kappa = TRUE, move_pi = TRUE, init_pi = 1,
                   find_import = FALSE, max_kappa = 10)


    out <- outbreaker(data, config)

    smry <- summary(out, burnin = 500)

    ## checks
    expect_equal(smry$tree$from, c(NA, 1, 2, 3))
    expect_equal(smry$tree$generations, c(NA, 1, 2, 4))
    expect_true(min(smry$post) > -35)
    expect_true(all(smry$pi[3:4] > 0.5 & smry$pi[3:4] < 0.8))

})






test_that("results ok: kappa from genetic data", {

    ## skip on CRAN
    skip_on_cran()

    dates <- c(0, 5, 15, 20, 30)
    w <- dnorm(1:10, 5, 1)
    dna <- matrix(c("A", "A", "A", "A", "A", "A", "A", "A", "A", "A", "A", "A",
                    "T", "T", "A", "A", "A", "A", "A", "A", "A", "A", "A", "A",
                    "T", "T", "G", "G", "G", "G", "A", "A", "A", "A", "A", "A",
                    "T", "T", "G", "G", "G", "G", "T", "T", "A", "A", "A", "A",
                    "T", "T", "G", "G", "G", "G", "T", "T", "G", "G", "G", "G"),
                  byrow = TRUE, nrow = 5)
    dna <- ape::as.DNAbin(dna)
    data <- outbreaker_data(dates = dates, dna = dna, w_dens = w)
    config <- create_config(init_mu = 2/12, move_mu = FALSE)
    res <- outbreaker(data, config)
    expect_equal(summary(res)$tree$generations,
                 c(NA, 1L, 2L, 1L, 2L))

})









test_that("results ok: outbreaker with custom priors",
{
    ## skip on CRAN
    skip_on_cran()


    ## get data
    data(fake_outbreak)
    x <- fake_outbreak

    data <- list(dna = x$dna, dates = x$onset, w_dens = x$w)

    config <-  list(n_iter = 5e3, sample_every = 50,
                    init_tree = "star", move_kappa = TRUE,
                    move_pi = TRUE, init_pi = 0.9,
                    find_import = FALSE)


    ## plot(function(x) dbeta(x, 50, 25)) # to see the distribution

    f_pi_1 <- function(x) dbeta(x$pi, 100, 1, log = TRUE) # stronger prior
    f_pi_2 <- function(x) 0.0 # flat prior

    set.seed(1)

    out_1 <- outbreaker(data, config, priors = list(pi = f_pi_1))
    out_2 <- outbreaker(data, config, priors = list(pi = f_pi_2))

    smry_1 <- summary(out_1, burnin = 500)
    smry_2 <- summary(out_2, burnin = 500)

    expect_true(smry_1$pi[2] > 0.9)
    expect_true(smry_2$pi[2] > 0.6 && smry_2$pi[5] < 0.9)

})








test_that("results ok: outbreaker with fixed number returning priors and likelihoods",
{
    ## skip on CRAN
    skip_on_cran()


    ## get data

    data(fake_outbreak)
    x <- fake_outbreak

    data <- list(dna = x$dna, dates = x$onset, w_dens = x$w)

    config <-  list(n_iter = 1000, sample_every = 10,
                    init_tree = "star", move_kappa = TRUE,
                    move_pi = TRUE, init_pi = 1,
                    find_import = FALSE)


    ## custom functions

    f1 <- function(x) return(0.0)
    f2 <- function(x, y) return(0.0)
    f3 <- function(x) return(1.123)

    priors1 <- custom_priors(mu = f1, pi = f1)
    priors2 <- custom_priors(mu = f1, pi = f3)

    ll <- custom_likelihoods(genetic = f2,
                             timing_infections = f2,
                             timing_sampling = f2,
                             reporting = f2)


    ## tests
    out1 <- outbreaker(data, config, priors1, ll)
    out2 <- outbreaker(data, config, priors2, ll)

    expect_true(all(out1$post == 0))
    expect_true(all(out2$post == 1.123))

})





## test consensus tree
test_that("test consensus trees", {

    ## skip on CRAN
    skip_on_cran()


    ## get data
    data(fake_outbreak)
    x <- fake_outbreak

    ## outbreaker DNA + time ##
    ## analysis
    set.seed(1)
    data <- list(dna = x$dna, dates = x$onset, w_dens = x$w)
    config <- list(n_iter = 5000, sample_every = 50,
                   init_tree = "seqTrack", find_import = TRUE,
                   move_pi = FALSE)

    out <- outbreaker(data = data, config = config)

    ## checks
    out_smry_mpa <- summary(out, burnin = 1000, method = 'mpa')
    out_smry_dec <- summary(out, burnin = 1000, method = 'decycle')

    ## as there are no cycles, these should be the same
    testthat::expect_equal(out_smry_mpa$tree$from,
                           out_smry_dec$tree$from)

})