source("utils.R")
source("utils-bm-seq.R")
library(biomformat)
library(phyloseq)
library(R.utils)
library(stringr)
library(sbm)
library(blockmodels)
library(here)
library(microbenchmark)

args <- commandArgs(trailingOnly = TRUE)
mode <- args[1]
nb_cores <- as.integer(args[2])

options(parallelly.availableCores.custom = function() {
    max(1L, nb_cores)
})

result_folder <- here("results", "lbm-seq")

# the_data <- import_biom("data/mach/kinetic.biom")
the_data <- import_biom("data/chaillou/chaillou.biom")
data_name <- "chaillou"
epoch <- as.integer(Sys.time())

tmp_folder <- here(result_folder, data_name, paste0("tmp", epoch))
named_data_folder <- here(result_folder, data_name)

if (!dir.exists(tmp_folder)) {
    dir.create(tmp_folder, recursive = TRUE)
}

per_taxa_networks <- collapse_otu_at_taxo(the_data)

prev_model <- NULL

sequential <- FALSE

if (mode == "seq" || mode == "para") {
    sequential <- TRUE
}

ncores <- ifelse(mode == "para", parallelly::availableCores(), 1L)

model_list <- list()
net_ids <- seq_along(per_taxa_networks)[-1]


for (idx in net_ids) {
    if (!sequential || is.null(prev_model)) {
        current_model <- BM_poisson(
            membership_type = "LBM",
            adj = per_taxa_networks[[idx]], # Account for the root
            verbosity = 6,
            plotting = "",
            autosave = here(tmp_folder, paste0("r", idx, "_", mode, ".Rds")),
            ncores = ncores
        )
        current_model$estimate()
    } else {
        # We transfer
        current_model <- bm_propagate_taus_all_models(phyloseq_data = the_data, rank_id_start = idx - 1, target_rank_id = idx, per_taxa_networks = per_taxa_networks, first_model = prev_model, ncores = ncores, autosave = here(tmp_folder, paste0("r", idx, "_", mode, ".Rds")))
    }
    prev_model <- current_model
    model_list <- append(model_list, current_model)
}

names(model_list) <- paste0("Rank", net_ids)
saveRDS(model_list, here(named_data_folder, paste0(mode, "_", epoch, ".Rds")))