library(blockmodels)
library(phyloseq)

bm_propagate_taus_Q_next_level <- function(phyloseq_data, first_model, Q, per_taxa_networks = collapse_otu_at_taxo(phyloseq_data), rank_id_start = 2L, target_rank_id = rank_id_start + 1, removed_root = TRUE, sd_noise = 0.001) {
    stopifnot("Q must be provided" = !missing(Q), "pĥyloseq_data must be provided" = !missing(phyloseq_data))
    force(per_taxa_networks)

    result_list <- list()
    # Init
    mean_diff_vec <- c()
    PL_per_nb_vec <- c()
    max_rank_id <- tax_table(phyloseq_data) |> ncol()
    net_id_start <- ifelse(removed_root, rank_id_start - 1, rank_id_start)
    result_list <- append(result_list, first_model)
    # Here we extract the memberships
    current_memberships <- first_model$memberships[[Q]]
    if (!is.matrix(current_memberships$Z1)) {
        nrow_Z1 <- length(current_memberships$Z1)
    } else {
        nrow_Z1 <- nrow(current_memberships$Z1)
    }
    next_cc <- list(Z1 = as.matrix(current_memberships$Z1, nrow = nrow_Z1), Z2 = current_memberships$Z2)

    # Here we dispatch the memberships for the OTU
    ## Propagate the rownames
    rownames(next_cc$Z1) <- rownames(per_taxa_networks[[net_id_start]])
    target_net_id <- ifelse(removed_root, target_rank_id - 1, target_rank_id)
    next_Z1 <- propagate_taus(tau_matrix = next_cc$Z1, physeq = phyloseq_data, taxrank = phyloseq::rank_names(phyloseq_data)[target_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 = sd_noise), 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_networks[[target_net_id]]), real_EM = TRUE
    ) # account for the removed root
    result_list <- append(result_list, next_res)
    mean_diff_vec[paste0(rank_id_start, "->", target_rank_id)] <- mean((next_res$membership$Z1 - next_Z1)^2)
    next_cc <- next_res[["membership"]][c("Z1", "Z2")]
    PL_per_nb_vec[paste0(rank_id_start, "->", target_rank_id)] <- next_res$PL / (nrow(next_res$membership$Z1) * nrow(next_res$membership$Z2))
    return(next_res)
}
#' Use a preinited model to run blockmodels exploration
#'
#' @param bm_model A pre-initialized model object
bm_estimate_with_inits <- function(bm_model, reinitialization_effort = 1) {
    l <- TRUE
    n <- 1
    changing_effort <- FALSE
    while (l) {
        cat("Pass", n, "\n")

        cat("With ascending number of groups\n")
        ra <- bm_model$estim_ascend(reinitialization_effort, changing_effort)

        cat("With descending number of groups\n")
        rb <- bm_model$estim_descend(reinitialization_effort)

        l <- ra || rb
        n <- n + 1
        changing_effort <- FALSE
    }
}

bm_propagate_taus_all_models <- function(phyloseq_data, rank_id_start = 2, target_rank_id = rank_id_start + 1, first_model = NULL, per_taxa_networks = collapse_otu_at_taxo(phyloseq_data), sd_noise = 0.001, ncores = parallelly::availableCores(), verbosity = 2L) {
    removed_root <- FALSE
    # Detect if first rank (higher) is fully collapsed
    if (nrow(per_taxa_networks[[1]]) <= 1L) {
        message("The first network has only one row. Removing it.")
        removed_root <- TRUE
        per_taxa_networks <- per_taxa_networks[-1]
    }
    result_list <- list()
    # Init
    net_id_start <- ifelse(removed_root, rank_id_start - 1, rank_id_start)
    if (missing(first_model)) {
        first_model <- BM_poisson(
            membership_type = "LBM",
            adj = per_taxa_networks[[net_id_start]], # Account for the root
            verbosity = verbosity,
            plotting = "",
            ncores = ncores
        )
        first_model$estimate()
    }

    propagated_models <- lapply(seq_along(first_model$ICL), function(current_q) {
        if (!is.na(first_model$ICL[current_q])) {
            bm_propagate_taus_Q_next_level(phyloseq_data = phyloseq_data, first_model = first_model, per_taxa_networks = per_taxa_networks, Q = current_q, removed_root = removed_root, rank_id_start = rank_id_start, target_rank_id = target_rank_id, sd_noise = sd_noise)
        } else {
            return(NA)
        }
    })

    target_net_id <- ifelse(removed_root, target_rank_id - 1, target_rank_id)

    new_model <- BM_poisson(
        membership_type = "LBM",
        adj = per_taxa_networks[[target_net_id]], # Account for the root
        verbosity = verbosity,
        plotting = "",
        ncores = ncores
    )

    lapply(seq_along(propagated_models), function(current_q) {
        if (any(is.na(propagated_models[[current_q]]))) {
            return(NA)
        }
        new_model$memberships[[current_q]] <<- getRefClass("LBM")(from_cc = propagated_models[[current_q]]$membership)
        new_model$model_parameters[[current_q]] <- propagated_models[[current_q]]$model
        new_model$PL[current_q] <- propagated_models[[current_q]]$PL
        new_model$H[current_q] <- propagated_models[[current_q]]$H
        new_model$ICL[current_q] <- propagated_models[[current_q]]$PL - .5 * (propagated_models[[current_q]]$model$n_parameters *
            log(new_model$data_number()) + new_model$memberships[[current_q]]$ICL_penalty())
    })

    bm_estimate_with_inits(bm_model = new_model)

    return(new_model)
}