set.seed(1234)

rmse <- function(theta, theta_star) {
  sqrt(sum((theta - theta_star)^2) / sum(theta_star^2))
}

## Common parameters
nbNodes <- c(100, 120)
blockProp <- list(row = c(.5, .5), col = c(1 / 3, 1 / 3, 1 / 3)) # group proportions
nbBlocks <- sapply(blockProp, length)

test_that("BipartiteSBM_fit 'Bernoulli' model, undirected, no covariate", {
  ## BIPARTITE UNDIRECTED BERNOULLI SBM
  means <- matrix(c(0.05, 0.95, 0.4, 0.75, 0.15, 0.6), 2, 3) # connectivity matrix
  connectParam <- list(mean = means)

  ## Basic construction - check for wrong specifications
  mySampler <- BipartiteSBM$new("bernoulli", nbNodes, blockProp, connectParam)
  mySampler$rMemberships(store = TRUE)
  mySampler$rEdges(store = TRUE)

  ## Construction----------------------------------------------------------------
  mySBM <- BipartiteSBM_fit$new(mySampler$networkData, "bernoulli")
  expect_error(BipartiteSBM_fit$new(SamplerBernoulli$networkData, "bernouilli"))

  ## Checking class
  expect_true(inherits(mySBM, "SBM"))
  expect_true(inherits(mySBM, "BipartiteSBM"))
  expect_true(inherits(mySBM, "BipartiteSBM_fit"))

  ## Checking field access and format prior to estimation
  ## parameters
  expect_equal(mySBM$modelName, "bernoulli")
  expect_equal(unname(mySBM$nbNodes), nbNodes)
  expect_equal(mySBM$nbDyads, nbNodes[1] * nbNodes[2])
  expect_true(is.matrix(mySBM$connectParam$mean))

  ## covariates
  expect_equal(mySBM$covarEffect, numeric(0))
  expect_equal(mySBM$nbCovariates, 0)
  expect_equal(mySBM$covarList, list())
  expect_equal(mySBM$covarParam, numeric(0))

  ## S3 methods
  expect_equal(coef(mySBM, "connectivity"), mySBM$connectParam)
  expect_equal(coef(mySBM, "block"), mySBM$blockProp)
  expect_equal(coef(mySBM, "covariates"), mySBM$covarParam)

  ## Estimation-----------------------------------------------------------------
  BM_out <- mySBM$optimize(estimOptions = list(verbosity = 0))
  mySBM$setModel(4)

  expect_equal(mySBM$nbConnectParam, unname(nbBlocks[1] * nbBlocks[2]))
  expect_equal(mySBM$penalty, unname(nbBlocks[1] * nbBlocks[2] * log(nbNodes[1] * nbNodes[2]) + (nbBlocks[1] - 1) * log(nbNodes[1]) + (nbBlocks[2] - 1) * log(nbNodes[2])))
  expect_equal(mySBM$entropy, -sum(mySBM$probMemberships[[1]] * log(mySBM$probMemberships[[1]]))
  - sum(mySBM$probMemberships[[2]] * log(mySBM$probMemberships[[2]])))

  ## Expectation
  expect_equal(dim(mySBM$expectation), nbNodes)
  expect_true(all(mySBM$expectation >= 0, na.rm = TRUE))
  expect_true(all(mySBM$expectation <= 1, na.rm = TRUE))
  expect_null(mySBM$connectParam$var)

  ## blocks
  expect_equal(mySBM$nbBlocks, nbBlocks)
  expect_equivalent(dim(mySBM$probMemberships[[1]]), c(nbNodes[1], nbBlocks[1]))
  expect_equivalent(dim(mySBM$probMemberships[[2]]), c(nbNodes[2], nbBlocks[2]))
  expect_equal(sort(unique(mySBM$memberships[[1]])), 1:nbBlocks[1])
  expect_equal(sort(unique(mySBM$memberships[[2]])), 1:nbBlocks[2])

  ## S3 methods
  expect_equal(coef(mySBM, "connectivity"), mySBM$connectParam)
  expect_equal(coef(mySBM, "block"), mySBM$blockProp)
  expect_equal(coef(mySBM, "covariates"), mySBM$covarParam)
  expect_equal(mySBM$predict(), predict(mySBM))
  expect_equal(fitted(mySBM), predict(mySBM))

  ## correctness
  expect_lt(rmse(sort(mySBM$connectParam$mean), sort(means)), .2)
  expect_lt(1 - aricode::ARI(mySBM$memberships[[1]], mySampler$memberships[[1]]), .2)
  expect_lt(1 - aricode::ARI(mySBM$memberships[[2]], mySampler$memberships[[2]]), .2)

  ## prediction wrt BM
  for (Q in mySBM$storedModels$indexModel) {
    pred_bm <- BM_out$prediction(Q = Q)
    mySBM$setModel(Q - 1)
    pred_sbm <- predict(mySBM)
    expect_lt(rmse(pred_bm, pred_sbm), 1e-12)
  }
})

test_that("BipartiteSBM_fit 'Poisson' model, undirected, no covariate", {
  ## SIMPLE DIRECTED POISSON SBM
  means <- matrix(c(10, 5, 7, 15, 20, 8), 2, 3) # connectivity matrix
  connectParam <- list(mean = means)

  ## Basic construction - check for wrong specifications
  mySampler <- BipartiteSBM$new("poisson", nbNodes, blockProp, connectParam)
  mySampler$rMemberships(store = TRUE)
  mySampler$rEdges(store = TRUE)

  ## Construction----------------------------------------------------------------
  mySBM <- BipartiteSBM_fit$new(mySampler$networkData, "poisson")
  expect_error(BipartiteSBM_fit$new(SamplerBernoulli$networkData, "poison"))

  ## Checking class
  expect_true(inherits(mySBM, "SBM"))
  expect_true(inherits(mySBM, "BipartiteSBM"))
  expect_true(inherits(mySBM, "BipartiteSBM_fit"))

  ## Checking field access and format prior to estimation
  ## parameters
  expect_equal(mySBM$modelName, "poisson")
  expect_equal(unname(mySBM$nbNodes), nbNodes)
  expect_equal(mySBM$nbDyads, nbNodes[1] * nbNodes[2])
  expect_true(is.matrix(mySBM$connectParam$mean))

  ## covariates
  expect_equal(mySBM$covarEffect, numeric(0))
  expect_equal(mySBM$nbCovariates, 0)
  expect_equal(mySBM$covarList, list())
  expect_equal(mySBM$covarParam, numeric(0))

  ## S3 methods
  expect_equal(coef(mySBM, "connectivity"), mySBM$connectParam)
  expect_equal(coef(mySBM, "block"), mySBM$blockProp)
  expect_equal(coef(mySBM, "covariates"), mySBM$covarParam)

  ## Estimation-----------------------------------------------------------------
  BM_out <- mySBM$optimize(estimOptions = list(verbosity = 0))
  mySBM$setModel(4)

  ## Expectation
  expect_equal(dim(mySBM$expectation), nbNodes)
  expect_true(all(mySBM$expectation >= 0, na.rm = TRUE))
  expect_null(mySBM$connectParam$var)

  ## blocks
  expect_equal(mySBM$nbBlocks, nbBlocks)
  expect_equivalent(dim(mySBM$probMemberships[[1]]), c(nbNodes[1], nbBlocks[1]))
  expect_equivalent(dim(mySBM$probMemberships[[2]]), c(nbNodes[2], nbBlocks[2]))
  expect_equal(sort(unique(mySBM$memberships[[1]])), 1:nbBlocks[1])
  expect_equal(sort(unique(mySBM$memberships[[2]])), 1:nbBlocks[2])

  ## S3 methods
  expect_equal(coef(mySBM, "connectivity"), mySBM$connectParam)
  expect_equal(coef(mySBM, "block"), mySBM$blockProp)
  expect_equal(coef(mySBM, "covariates"), mySBM$covarParam)
  expect_equal(mySBM$predict(), predict(mySBM))
  expect_equal(fitted(mySBM), predict(mySBM))

  ## correctness
  expect_lt(rmse(sort(mySBM$connectParam$mean), sort(means)), 1e-1)
  expect_lt(1 - aricode::ARI(mySBM$memberships[[1]], mySampler$memberships[[1]]), 1e-1)
  expect_lt(1 - aricode::ARI(mySBM$memberships[[2]], mySampler$memberships[[2]]), 1e-1)

  ## prediction wrt BM
  for (Q in mySBM$storedModels$indexModel) {
    pred_bm <- BM_out$prediction(Q = Q)
    mySBM$setModel(Q - 1)
    pred_sbm <- predict(mySBM)
    expect_lt(rmse(pred_bm, pred_sbm), 1e-12)
  }
})

test_that("BipartiteSBM_fit 'Gaussian' model, undirected, no covariate", {
  ## SIMPLE UNDIRECTED GAUSSIAN SBM
  means <- matrix(c(0.05, 0.95, 0.4, 0.98, 0.15, 0.6), 2, 3) # connectivity matrix
  connectParam <- list(mean = means, var = .1)

  ## Basic construction - check for wrong specifications
  mySampler <- BipartiteSBM$new("gaussian", nbNodes, blockProp, connectParam)
  mySampler$rMemberships(store = TRUE)
  mySampler$rEdges(store = TRUE)

  ## Construction----------------------------------------------------------------
  mySBM <- BipartiteSBM_fit$new(mySampler$networkData, "gaussian")
  expect_error(BipartiteSBM_fit$new(SamplerBernoulli$networkData, "groß"))

  ## Checking class
  expect_true(inherits(mySBM, "SBM"))
  expect_true(inherits(mySBM, "BipartiteSBM"))
  expect_true(inherits(mySBM, "BipartiteSBM_fit"))

  ## Checking field access and format prior to estimation
  ## parameters
  expect_equal(mySBM$modelName, "gaussian")
  expect_equal(unname(mySBM$nbNodes), nbNodes)
  expect_equal(mySBM$nbDyads, nbNodes[1] * nbNodes[2])
  expect_true(is.matrix(mySBM$connectParam$mean))

  ## covariates
  expect_equal(mySBM$covarEffect, numeric(0))
  expect_equal(mySBM$nbCovariates, 0)
  expect_equal(mySBM$covarList, list())
  expect_equal(mySBM$covarParam, numeric(0))

  ## S3 methods
  expect_equal(coef(mySBM, "connectivity"), mySBM$connectParam)
  expect_equal(coef(mySBM, "block"), mySBM$blockProp)
  expect_equal(coef(mySBM, "covariates"), mySBM$covarParam)

  ## Estimation-----------------------------------------------------------------
  BM_out <- mySBM$optimize(estimOptions = list(verbosity = 0))
  mySBM$setModel(4)

  ## Expectation
  expect_equal(dim(mySBM$expectation), nbNodes)
  expect_gt(mySBM$connectParam$var, 0)

  ## blocks
  expect_equal(mySBM$nbBlocks, nbBlocks)
  expect_equivalent(dim(mySBM$probMemberships[[1]]), c(nbNodes[1], nbBlocks[1]))
  expect_equivalent(dim(mySBM$probMemberships[[2]]), c(nbNodes[2], nbBlocks[2]))
  expect_equal(sort(unique(mySBM$memberships[[1]])), 1:nbBlocks[1])
  expect_equal(sort(unique(mySBM$memberships[[2]])), 1:nbBlocks[2])

  ## S3 methods
  expect_equal(coef(mySBM, "connectivity"), mySBM$connectParam)
  expect_equal(coef(mySBM, "block"), mySBM$blockProp)
  expect_equal(coef(mySBM, "covariates"), mySBM$covarParam)
  expect_equal(mySBM$predict(), predict(mySBM))
  expect_equal(fitted(mySBM), predict(mySBM))

  ## correctness
  expect_lt(rmse(sort(mySBM$connectParam$mean), sort(means)), 1e-1)
  expect_lt(1 - aricode::ARI(mySBM$memberships[[1]], mySampler$memberships[[1]]), 1e-1)
  expect_lt(1 - aricode::ARI(mySBM$memberships[[2]], mySampler$memberships[[2]]), 1e-1)

  ## prediction wrt BM
  for (Q in mySBM$storedModels$indexModel) {
    pred_bm <- BM_out$prediction(Q = Q)
    mySBM$setModel(Q - 1)
    pred_sbm <- predict(mySBM)
    expect_lt(rmse(pred_bm, pred_sbm), 1e-12)
  }
})

test_that("active bindings are working in the class", {
  A <- matrix(rbinom(200, 1, .2), 20, 10)

  myBipartite <- BipartiteSBM_fit$new(incidenceMatrix = A, model = "bernoulli", dimLabels = c("Actor", "Stuff"))

  tau1 <- matrix(runif(20 * 2), 20, 2)
  tau1 <- tau1 / rowSums(tau1)
  tau2 <- matrix(runif(10 * 3), 10, 3)
  tau2 <- tau2 / rowSums(tau2)
  myBipartite$probMemberships <- list(tau1, tau2)
  myBipartite$blockProp <- list(colMeans(tau1), colMeans(tau2))
  myBipartite$connectParam <- list(mean = matrix(runif(3 * 2), 3, 2))

  expect_equal(unname(myBipartite$nbNodes), c(20, 10))

  expect_equal(myBipartite$memberships[[1]], 1 + (tau1[, 1] < .5) * 1)
  expect_equal(dim(myBipartite$connectParam$mean), c(3, 2))
  expect_equal(length(myBipartite$blockProp), 2)
})