source("utils.R")
library(biomformat)
library(phyloseq)
library(R.utils)
library(stringr)
library(sbm)
library(blockmodels)

the_data <- import_biom("data/chaillou/chaillou.biom")

per_taxa_network <- collapse_otu_at_taxo(the_data)

# Detect if first rank (higher) is fully collapsed
if (nrow(per_taxa_network[[1]]) <= 1L) {
    message("The first network has only one row. Removing it.")
    per_taxa_network <- per_taxa_network[-1]
}

# fit_res <- lapply(per_taxa_network, function(network) {
#     withTimeout(
#         expr = {
#             fit <- estimateBipartiteSBM(network, model = "poisson", estimOptions = list(plot = 0))
#         }, timeout = 300,
#         onTimeout = "warning"
#     )
#     if (exists("fit")) {
#         out <- fit
#     } else {
#         out <- NULL
#     }
#     return(fit)
# })

model <- BM_poisson(
    membership_type = "LBM",
    adj = per_taxa_network[[1]],
    verbosity = 0,
    plotting = "",
    ncores = 1L
)
model$estimate()
# Here we extract the memberships
best_model_memberships <- model$memberships[[which.max(model$ICL)]]
next_cc <- list(Z1 = best_model_memberships$Z1, Z2 = best_model_memberships$Z2)

# Here we dispatch the memberships for the OTU
## Propagate the rownames
rank_id <- 3


rownames(next_cc$Z1) <- rownames(per_taxa_network[[rank_id - 2]])
next_Z1 <- propagate_taus(tau_matrix = next_cc$Z1, physeq = the_data, taxrank = phyloseq::rank_names(the_data)[rank_id])

## Noising of the next_Z1
dZ1 <- dim(next_Z1)
nZ1 <- dZ1[1] * dZ1[2]
next_Z1 <- next_Z1 + matrix(rnorm(n = nZ1, mean = 0, sd = 0.001), nrow = dZ1[1], ncol = dZ1[2])

## Normalizing by rows
C_Z1 <- rowSums(next_Z1)
next_Z1 <- next_Z1 / C_Z1

next_cc$Z1 <- next_Z1

next_res <- blockmodels:::dispatcher(
    membership_name = "LBM",
    membership_init = next_cc,
    model_name = "poisson",
    network = list(adjacency = per_taxa_network[[2]]), real_EM = TRUE
)
next_cc <- next_res[["membership"]][c("Z1", "Z2")]
rank_id <- rank_id + 1