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-data")

# 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)


message("Will use SEQ")
r2_model_seq <- BM_poisson(
    membership_type = "LBM",
    adj = per_taxa_networks[[2]], # Account for the root
    verbosity = 6,
    plotting = "",
    autosave = here(tmp_folder, "r2_seq.Rds"),
    ncores = 1L
)
r2_model_seq$estimate()
r3_model_seq <- bm_propagate_taus_all_models(phyloseq_data = the_data, rank_id_start = 2, target_rank_id = 3, per_taxa_networks = per_taxa_networks, first_model = r2_model_seq, ncores = 1, autosave = here(tmp_folder, "r3_seq.Rds"))
r4_model_seq <- bm_propagate_taus_all_models(phyloseq_data = the_data, rank_id_start = 3, per_taxa_networks = per_taxa_networks, first_model = r3_model_seq, ncores = 1, autosave = here(tmp_folder, "r4_seq.Rds"))
r5_model_seq <- bm_propagate_taus_all_models(phyloseq_data = the_data, rank_id_start = 4, per_taxa_networks = per_taxa_networks, first_model = r4_model_seq, ncores = 1, autosave = here(tmp_folder, "r5_seq.Rds"))
r6_model_seq <- bm_propagate_taus_all_models(phyloseq_data = the_data, rank_id_start = 4, per_taxa_networks = per_taxa_networks, first_model = r5_model_seq, ncores = 1, autosave = here(tmp_folder, "r6_seq.Rds"))
r7_model_seq <- bm_propagate_taus_all_models(phyloseq_data = the_data, rank_id_start = 4, per_taxa_networks = per_taxa_networks, first_model = r4_model_seq, ncores = 1, autosave = here(tmp_folder, "r7_seq.Rds"))

out_seq <- list(
    models = list(
        Rank2 = r2_model_seq,
        Rank3 = r3_model_seq,
        Rank4 = r4_model_seq,
        Rank5 = r5_model_seq,
        Rank6 = r6_model_seq,
        Rank7 = r7_model_seq
    )
)

saveRDS(out_seq, here(named_data_folder, "seq.Rds"))