context("Get Functions")


# Model simulation --------------------------------------------------------

nw <- network_initialize(n = 100)
nw <- set_vertex_attribute(nw, "group", rep(1:2, each = 50))
formation <- ~edges
target.stats <- 50
dissolution <- ~offset(edges)
duration <- 20
coef.diss <- dissolution_coefs(dissolution, duration)
est <- netest(nw, formation, target.stats, coef.diss, verbose = FALSE)
param <- param.net(inf.prob = 0.3, inf.prob.g2 = 0.15)
init <- init.net(i.num = 10, i.num.g2 = 10)
control <- control.net(type = "SI", nsteps = 10, nsims = 3,
                       verbose = FALSE)
mod <- netsim(est, param, init, control)


# get network -------------------------------------------------------------

test_that("get_network extracts to nD object", {
  skip_on_cran()
  a <- get_network(mod, sim = 2)
  expect_is(a, "networkDynamic")
})

test_that("get_network extracts and collapses to network object", {
  skip_on_cran()
  a <- get_network(mod, collapse = TRUE, at = 3)
  expect_is(a, "network")
})

test_that("get_network yields warning when missing network", {
  skip_on_cran()
  mod$network <- NULL
  expect_error(get_network(mod, sim = 1))
})

test_that("get_network yields warning for incorrect sim", {
  skip_on_cran()
  expect_error(get_network(mod, sim = 5))
})

test_that("get_network error flags", {
  skip_on_cran()
  expect_error(get_network(list(a = 1), 1), "no applicable method for 'get_network'")
  expect_error(get_network(mod, 4), "Specify a single sim between 1 and 3")
  expect_error(get_network(mod, 1, collapse = TRUE), "Specify collapse time")
  expect_error(get_network(mod, 1, collapse = TRUE, at = 200),
               "Specify collapse time")
  expect_error(get_network(mod, 1, 2), "Specify network")
})


# get transmat ------------------------------------------------------------

test_that("get_transmat extracts data frame", {
  skip_on_cran()
  a <- get_transmat(mod, sim = 1)
  expect_is(a, "data.frame")
})

test_that("get_transmat error flags", {
  skip_on_cran()
  expect_error(get_transmat(mod, sim = 5), "Specify sim between 1 and 3")
  expect_error(get_transmat(list(a = 1)), "x must be of class netsim")
  mod$stats$transmat <- NULL
  expect_error(get_transmat(mod, 1), "transmat not saved")
})


# get nwstats -------------------------------------------------------------

test_that("get_nwstats extracts data frame", {
  skip_on_cran()
  a <- get_nwstats(mod, sim = 1)
  expect_is(a, "data.frame")
  expect_equal(get_nwstats(mod, sim = 1:3), get_nwstats(mod))
})

test_that("get_nwstats error flags", {
  skip_on_cran()
  expect_error(get_nwstats(list(a = 1)), "x must be of class netsim")
  expect_error(get_nwstats(mod, sim = 5))
  expect_error(get_nwstats(mod, sim = 1, network = 2),
               "Specify network between 1")
  mod$stats$nwstats <- NULL
  expect_error(get_nwstats(mod), "Network statistics not saved")
})


# get sims ----------------------------------------------------------------

test_that("get_sims extracts simulations", {
  skip_on_cran()
  expect_is(get_sims(mod, sims = 1), "netsim")
  expect_is(get_sims(mod, sims = 2:3), "netsim")
  expect_is(get_sims(mod, sims = "mean", var = "i.num"), "netsim")
  expect_is(get_sims(mod, sims = 1:3), "netsim")
  expect_is(get_sims(mod, sims = 1, var = c("i.num", "s.num")), "netsim")
  expect_is(get_sims(mod, sims = c(1, 3), var = c("i.num", "s.num")), "netsim")
  expect_equal(length(get_sims(mod, sims = 1:3)$diss.stats), 3)
  expect_equal(names(get_sims(mod, sims = 1:3)$diss.stats), c("sim1", "sim2", "sim3"))
  expect_equal(length(get_sims(mod, sims = c(1, 3))$diss.stats), 2)
  expect_equal(names(get_sims(mod, sims = c(1, 3))$diss.stats), c("sim1", "sim2"))
  expect_equal(length(get_sims(mod, sims = 2)$diss.stats), 1)
  expect_equal(names(get_sims(mod, sims = 2)$diss.stats), c("sim1"))
  expect_equal(length(merge(get_sims(mod, sims = 2:3),
                            get_sims(mod, sims = c(1,3)))$diss.stats), 4)
  plot(get_sims(mod, sims = 2), type = "duration")
  plot(get_sims(mod, sims = c(1, 3)), type = "dissolution")
})


test_that("get_sims error flags", {
  skip_on_cran()
  expect_error(get_sims(list(a = 1)), "x must be of class netsim")
  expect_error(get_sims(mod), "Specify sims as a vector")

# get parameter set ------------------------------------------------------------
  nw <- network_initialize(n = 50)

  est <- netest(
    nw, formation = ~edges,
    target.stats = c(25),
    coef.diss = dissolution_coefs(~offset(edges), 10, 0),
    verbose = FALSE
  )

  init <- init.net(i.num = 10)

  my.randoms <- list(
    act.rate = param_random(c(0.25, 0.5, 0.75)),
    dummy.param = function() rbeta(1, 1, 2),
    dummy.strat.param = function() c(
      rnorm(1, 0, 10),
      rnorm(1, 10, 1)
    )
  )

  param <- param.net(
    inf.prob = 0.3,
    dummy = c(0, 1, 2),
    random.params = my.randoms
  )

  control <- control.net(type = "SI", nsims = 3, nsteps = 5, verbose = FALSE)
  mod <- netsim(est, param, init, control)
  d.set <- get_param_set(mod)

  set.colnames <- c(
    "sim",
    "inf.prob",
    "dummy_1",
    "dummy_2",
    "dummy_3",
    "act.rate",
    "vital",
    "dummy.param",
    "dummy.strat.param_1",
    "dummy.strat.param_2",
    "groups"
  )

  expect_is(d.set, "data.frame")
  expect_true(setequal(names(d.set), set.colnames))
  expect_error(get_param_set(control), "`sims` must be of class netsim")
  expect_equal(dim(get_param_set(mod)), c(3, length(set.colnames)))

  mod.sub <- get_sims(mod, sims = 2)
  d.sub <- get_param_set(mod.sub)
  expect_equal(nrow(d.sub), 1)
  expect_equal(d.sub$act.rate, d.set$act.rate[2])
  expect_equal(d.sub$dummy.param, d.set$dummy.param[2])
  expect_equal(d.sub$dummy.strat.param_1, d.set$dummy.strat.param_1[2])
  expect_equal(d.sub$dummy.strat.param_2, d.set$dummy.strat.param_2[2])
})

test_that("get_param_set handles vector random parameters with one simulation", {
  mod <- structure(
    list(
      param = list(
        inf.prob = 0.3,
        dummy.strat.param = c(0.1, 0.2),
        random.params = list(dummy.strat.param = function() NULL),
        random.params.values = list(dummy.strat.param = c(0.1, 0.2))
      ),
      control = list(nsims = 1)
    ),
    class = "netsim"
  )

  d.set <- get_param_set(mod)
  expect_equal(nrow(d.set), 1)
  expect_equal(d.set$dummy.strat.param_1, 0.1)
  expect_equal(d.set$dummy.strat.param_2, 0.2)
})

test_that("get_sims subsets random parameter values", {
  mod <- structure(
    list(
      param = list(
        inf.prob = 0.3,
        act.rate = 0.1,
        dummy.strat.param = c(1, 2),
        random.params = list(
          act.rate = function() NULL,
          dummy.strat.param = function() NULL
        ),
        random.params.values = list(
          act.rate = c(0.1, 0.2, 0.3),
          dummy.strat.param = list(c(1, 2), c(3, 4), c(5, 6))
        )
      ),
      control = list(nsims = 3, save.other = character(0)),
      epi = list(
        i.num = data.frame(matrix(1:3, nrow = 1))
      ),
      stats = list(nwstats = NULL, transmat = NULL),
      run = NULL,
      network = NULL,
      diss.stats = NULL
    ),
    class = "netsim"
  )

  mod.sub <- get_sims(mod, sims = 2)
  d.sub <- get_param_set(mod.sub)

  expect_equal(mod.sub$control$nsims, 1)
  expect_equal(mod.sub$param$random.params.values$act.rate, 0.2)
  expect_equal(mod.sub$param$random.params.values$dummy.strat.param, list(c(3, 4)))
  expect_equal(d.sub$act.rate, 0.2)
  expect_equal(d.sub$dummy.strat.param_1, 3)
  expect_equal(d.sub$dummy.strat.param_2, 4)
})

dxs <- netdx(est, dynamic = FALSE, nsims = 5,
             nwstats.formula = ~edges + nodemix("group", levels2 = TRUE), verbose = FALSE)
dxd <- netdx(est, dynamic = TRUE, nsims = 5, nsteps = 3, verbose = FALSE)

control1 <- control.net(type = "SI", nsteps = 2, nsims = 3,
                        verbose = FALSE, tergmLite = FALSE,
                        resimulate.network = TRUE,
                        nwstats.formula = ~edges + triangle)
mod1 <- netsim(est, param, init, control1)

control2 <- control.net(type = "SI", nsteps = 3, nsims = 4,
                        verbose = FALSE, tergmLite = FALSE,
                        resimulate.network = FALSE)
mod2 <- netsim(est, param, init, control2)

control3 <- control.net(type = "SI", nsteps = 4, nsims = 2,
                        verbose = FALSE, tergmLite = TRUE,
                        resimulate.network = TRUE,
                        nwstats.formula = ~edges +
                                           nodematch("group", diff = TRUE))
mod3 <- netsim(est, param, init, control3)

test_that("get_nwstats with mode = list behaves as expected", {
  skip_on_cran()
  expect_equal(unique(lapply(get_nwstats(dxs, mode = "list"), class)),
               list(c("matrix", "array")))
  expect_equal(unique(lapply(get_nwstats(dxd, mode = "list"), class)),
               list(c("matrix", "array")))
  expect_equal(unique(lapply(get_nwstats(mod1, mode = "list"), class)),
               list(c("matrix", "array")))
  expect_equal(unique(lapply(get_nwstats(mod2, mode = "list"), class)),
               list(c("matrix", "array")))
  expect_equal(unique(lapply(get_nwstats(mod3, mode = "list"), class)),
               list(c("matrix", "array")))

  expect_equal(unique(lapply(get_nwstats(dxs, mode = "list"), dim)),
               list(c(5, 4)))
  expect_equal(unique(lapply(get_nwstats(dxd, mode = "list"), dim)),
               list(c(3, 1)))
  expect_equal(unique(lapply(get_nwstats(mod1, mode = "list"), dim)),
               list(c(2, 2)))
  expect_equal(unique(lapply(get_nwstats(mod2, mode = "list"), dim)),
               list(c(3, 1)))
  expect_equal(unique(lapply(get_nwstats(mod3, mode = "list"), dim)),
               list(c(4, 3)))

  expect_equal(length(get_nwstats(dxs, mode = "list")), 1)
  expect_equal(length(get_nwstats(dxd, mode = "list")), 5)
  expect_equal(length(get_nwstats(mod1, mode = "list")), 3)
  expect_equal(length(get_nwstats(mod2, mode = "list")), 4)
  expect_equal(length(get_nwstats(mod3, mode = "list")), 2)

  expect_equal(length(get_nwstats(dxs, sim = c(1), mode = "list")), 1)
  expect_equal(length(get_nwstats(dxd, sim = c(5,3,1), mode = "list")), 3)
  expect_equal(length(get_nwstats(mod1, sim = c(2,3), mode = "list")), 2)
  expect_equal(length(get_nwstats(mod2, sim = c(1,3,2), mode = "list")), 3)
  expect_equal(length(get_nwstats(mod3, sim = c(1,2), mode = "list")), 2)
})

test_that("get_network_attributes functions as intended", {
  nw <- network.initialize(10, directed = FALSE, bipartite = 4)
  nw %n% "newattr" <- "string"
  expect_equal(get_network_attributes(nw), list(bipartite = 4,
                                                directed = FALSE,
                                                hyper = FALSE,
                                                loops = FALSE,
                                                multiple = FALSE,
                                                n = 10,
                                                newattr = "string"))
})
