library(sbm)
library(phyloseq)
library(biomformat)
library(here)
source("utils.R")

options(parallelly.availableCores.custom = function() {
    ncores <- min(parallelly::availableCores(), 64)
    max(1L, ncores)
})

message(paste("Number of cores available:", parallelly::availableCores()))
args <- commandArgs(trailingOnly = TRUE)
if (length(args) == 0) {
    author <- "mach"
} else {
    author <- args[1]
}

message("For author: ", author)
the_data <- import_biom(here("data", author, paste0(author, ".biom")))

epoch <- as.integer(Sys.time())

per_taxa_networks <- collapse_otu_at_taxo(the_data)
per_taxa_networks <- per_taxa_networks[-1]

lbm_list <- lapply(seq_along(per_taxa_networks), function(idx) {
    defineSBM(per_taxa_networks[[idx]], model = "poisson", dimLabels = c(row = paste0("Rank", idx + 1), col = "Sample"))
})

lbm_list_covariates <- lapply(seq_along(per_taxa_networks), function(idx) {
    covar_sd <- log(rep(1, nrow(per_taxa_networks[[idx]])) %*% t(colSums(per_taxa_networks[[idx]])))
    defineSBM(per_taxa_networks[[idx]], model = "poisson", dimLabels = c(row = paste0("Rank", idx + 1), col = "Sample"), covariates = list(covar_sd))
})


multipartite_fit <- estimateMultipartiteSBM(listSBM = lbm_list, estimOptions = list(initBM = FALSE))

multipartite_covariates_fit <- estimateMultipartiteSBM(listSBM = lbm_list_covariates, estimOptions = list(initBM = FALSE))

plot(multipartite_covariates_fit)