#!/usr/bin/env Rscript
utils::globalVariables(".progress")
###############################################################################
# MaAsLin3 fitting

# Copyright (c) 2025 Harvard School of Public Health

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
###############################################################################

# Function to augment data for logistic fitting
augment_data <- function(formula, random_effects_formula, dat_sub) {
    d2 = dat_sub 
    d3 = dat_sub 
    
    d2 = collapse::ftransform(d2, expr = 1)
    d3 = collapse::ftransform(d3, expr = 0)
    
    dat_sub_new <- collapse::rowbind(dat_sub, d2, d3)
    
    formula <- stats::formula(formula)
    
    if (is.null(random_effects_formula)) {
        # No random effects
        p <- ncol(model.matrix(formula, dat_sub)) - 1
    } else {
        # With random effects
        p <- ncol(lme4::lFormula(formula, dat_sub)$X) - 1
    }
    
    # Calculate the weights
    weight_scheme <-
        c(rep(1, nrow(dat_sub)),
        rep(p / (2 * nrow(dat_sub)), 2 * nrow(dat_sub)))
    
    return(list(
        mm_input = dat_sub_new,
        weight_scheme = weight_scheme,
        new_formula = formula
    ))
}

# Deparse formula without line breaks
safe_deparse <- function(formula) {
    paste0(trimws(deparse(formula)), collapse = " ")
}

# Extract a predictor of the form: predictor_type(variable)
extract_special_predictor <- function(formula, predictor_type) {
    match.arg(predictor_type, c("group", "ordered", "strata"))
    groups <-
        regmatches(safe_deparse(formula), gregexpr(
            paste0(predictor_type, "\\((.*?)\\)"),
            safe_deparse(formula)
        ))[[1]]
    groups <-
        gsub("\\)$", "", gsub(paste0("^", predictor_type, "\\("), "", groups))
    formula_tmp <-
        trimws(gsub(
            paste0("(\\s*\\+\\s*)?", predictor_type, "\\(.*?\\)"),
            "",
            safe_deparse(formula)
        ))
    if (substr(formula_tmp, nchar(formula_tmp), nchar(formula_tmp)) == '~') {
        formula_tmp <- paste0(formula_tmp, '1')
    }
    formula <- stats::formula(formula_tmp)
    formula <- stats::formula(gsub("~ \\+", "~", safe_deparse(formula)))
    
    if (predictor_type == "strata" & length(groups) > 1) {
        stop("Only one strata allowed. Please change the formula.")
    }
    
    return(list(formula, predictor_type = groups))
}

# Get all fixed effects, not assuming a model has or can be fit
get_fixed_effects <-
    function(formula,
            random_effects_formula,
            dat_sub,
            groups,
            ordereds,
            feature_specific_covariate_name) {
        names_to_include <- c()
        if (is.null(random_effects_formula)) {
            # Fixed and group effects only
            names_to_include <- colnames(model.matrix(stats::formula(gsub(
                "^expr ", "", safe_deparse(formula)
            )), dat_sub))
            names_to_include <-
                names_to_include[names_to_include != "(Intercept)"]
        } else {
            # Random effects
            patterns <- paste0("(", unlist(reformulas::findbars(stats::formula(
                gsub("^expr ", "", safe_deparse(formula))
            ))), ")")
            
            # Remove all the random effects from the formula
            for (pattern in patterns) {
                fixed_effects_only <- gsub(pattern, "",
                                        paste0(trimws(safe_deparse(
                                            stats::formula(gsub(
                                                "^expr ", "",
                                                safe_deparse(formula)
                                            ))
                                        )), collapse = " "),
                                        fixed = TRUE)
                fixed_effects_only <-
                    gsub("[+ ]+$", "", fixed_effects_only)
                fixed_effects_only <-
                    gsub("\\+\\s*\\++", "+", fixed_effects_only)
                formula <- stats::formula(fixed_effects_only)
            }
            
            if (!is.null(feature_specific_covariate_name)) {
                names_to_include <- colnames(model.matrix(stats::formula(gsub(
                    paste0(feature_specific_covariate_name, " \\+|^expr "), 
                    "", safe_deparse(formula)
                )), dat_sub))
                names_to_include <- c(names_to_include, 
                                        feature_specific_covariate_name)
            } else {
                names_to_include <- colnames(model.matrix(stats::formula(gsub(
                    "^expr ", "", safe_deparse(formula)
                )), dat_sub))
            }
            
            names_to_include <-
                names_to_include[names_to_include != "(Intercept)"]
        }
        
        ordered_levels <- unlist(lapply(ordereds, function(ordered) {
            ordered_no_ticks <- gsub("^`|`$", "", ordered)
            paste0(ordered, levels(dat_sub[[ordered_no_ticks]])[-1])
        }))
        
        names_to_include <-
            c(names_to_include, groups, ordered_levels)
        return(names_to_include)
    }

# Get non-baseline names for all non-numeric columns
get_character_cols <- function(dat_sub) {
    all_factors <- unlist(lapply(colnames(dat_sub), function(col) {
        if (!is.numeric(dat_sub[, col])) {
            if (is.factor(dat_sub[, col])) {
                factor_levels <- levels(dat_sub[, col])
            } else {
                factor_levels <- levels(factor(dat_sub[, col]))
            }
            # All factor levels except baseline
            return(unlist(paste0(col, unique(factor_levels[-1]))))
        }
        return(character(0))
    }))
    
    return(all_factors)
}

# Correct over prevalence and abundance to get q-values
add_qvals <- function(fit_data_abundance, fit_data_prevalence, correction) {
    match.arg(correction, 
        c("BH", "holm", "hochberg", "hommel", "bonferroni", "BY"))
    # Select out p-values and NA if errors
    if (!is.null(fit_data_abundance)) {
        abundance_pvals <- as.numeric(fit_data_abundance$results$pval)
        
        abundance_pvals <- data.table::fifelse(!is.na(fit_data_abundance$results$error),
                                               NA_real_,
                                               abundance_pvals)
    } else {
        abundance_pvals <- c()
    }
    
    if (!is.null(fit_data_prevalence)) {
        prevalence_pvals <- as.numeric(fit_data_prevalence$results$pval)
        
        prevalence_pvals <- data.table::fifelse(!is.na(fit_data_prevalence$results$error),
                                                NA_real_,
                                                prevalence_pvals)
    } else {
        prevalence_pvals <- c()
    }

    # Create and write combined q-vals
    combined_qvals <- as.numeric(p.adjust(c(abundance_pvals, prevalence_pvals), 
                                method = correction))
    
    if (!is.null(fit_data_abundance)) {
        fit_data_abundance$results$qval <- combined_qvals[
            seq_len(length(abundance_pvals))]
    }
    if (!is.null(fit_data_prevalence)) {
        fit_data_prevalence$results$qval <- combined_qvals[
            seq(length(abundance_pvals) + 1, length(abundance_pvals) + 
                    length(prevalence_pvals))]
    }
    
    return(list(
        if (!is.null(fit_data_abundance)) fit_data_abundance$results else NULL,
        if (!is.null(fit_data_prevalence)) fit_data_prevalence$results else NULL
    ))
}

# Combine abundance and prevalence p-values
create_combined_pval <- function(merged_signif, correction) {
    match.arg(correction,
            c("BH", "holm", "hochberg", "hommel", "bonferroni", "BY"))
    # Create a combined p-value
    merged_signif$pval_joint <-
        pbeta(pmin(merged_signif[, "linear"],
                merged_signif[, "logistic"]), 1, 2)
    
    # If NA or model errored, use the p-value of the non-NA
    merged_signif$pval_joint <-
        data.table::fifelse(
            is.na(merged_signif$linear) |
                !is.na(merged_signif$linear_error),
            as.numeric(merged_signif$logistic),
            merged_signif$pval_joint
        )
    # ^ as.numeric() because fifelse checks if class(yes) == class(no) and NA is
    # logical by default
    
    merged_signif$pval_joint <-
        data.table::fifelse(
            is.na(merged_signif$logistic) |
                !is.na(merged_signif$logistic_error),
            as.numeric(merged_signif$linear),
            merged_signif$pval_joint
        )
    
    merged_signif$pval_joint <-
        data.table::fifelse((
            is.na(merged_signif$logistic) |
                !is.na(merged_signif$logistic_error)
        ) & (is.na(merged_signif$linear) |
                 !is.na(merged_signif$linear_error)),
        NA_real_,
        merged_signif$pval_joint)
    
    merged_signif$qval_joint <-
        as.numeric(p.adjust(merged_signif$pval_joint, method = correction))
    
    return(merged_signif)
}

# Put note when abundance effect has likely turned into a prevalence effect
flag_abundance_turned_prevalence <- function(merged_signif,
                                            max_significance) {
    
    tst = !is.na(merged_signif$logistic_coef) &
        !is.na(merged_signif$linear_coef) &
        !is.na(merged_signif$logistic) &
        !is.na(merged_signif$linear)
    
    inner_tst = is.na(merged_signif$logistic_error) &
        is.na(merged_signif$linear_error) &
        merged_signif$linear_qval < max_significance &
        sign(merged_signif$logistic_coef) == 
        sign(merged_signif$linear_coef) &
        abs(merged_signif$linear_coef) > 
        abs(merged_signif$logistic_coef)
    
    # Join and check linear and logistic pieces
    merged_signif$logistic_error <- data.table::fifelse(
        tst,
        data.table::fifelse(
            inner_tst,
            paste0("Prevalence association possibly induced ",
                   "by stronger abundance association"),
            merged_signif$logistic_error
        ),
        merged_signif$logistic_error
    )
    
    return(merged_signif)
}

# Get joint significance for abundance and prevalence associations
add_joint_signif <-
    function(fit_data_abundance,
             fit_data_prevalence,
             new_fit_data_abundance,
             max_significance,
             correction) {
        
        match.arg(correction,
                  c("BH", "holm", "hochberg", "hommel", "bonferroni", "BY"))
        
        # Subset to shared columns
        sel_vec = c("feature", "metadata", "value",
                    "name", "coef", "null_hypothesis",
                    "pval", "qval", "error")
        
        fit_data_prevalence_signif <-
            fit_data_prevalence$results |> collapse::fselect(sel_vec)
        
        colnames(fit_data_prevalence_signif) <-
            c("feature",
            "metadata",
            "value",
            "name",
            "logistic_coef",
            "logistic_null_hypothesis",
            "logistic",
            "logistic_qval",
            "logistic_error")
        
        fit_data_abundance_signif <-
            fit_data_abundance$results |> collapse::fselect(sel_vec)
        
        colnames(fit_data_abundance_signif) <-
            c("feature",
            "metadata",
            "value",
            "name",
            "linear_coef",
            "linear_null_hypothesis",
            "linear",
            "linear_qval",
            "linear_error")
        
        # Join and check linear and logistic pieces
        merged_signif <-
            collapse::join(
                collapse::funique(fit_data_prevalence_signif),
                collapse::funique(fit_data_abundance_signif),
                how = "full",
                verbose = FALSE,
                on = c("feature", "metadata", "value", "name"),
                overid = 2 # Pretty sure 0 would be fine, but 2 to be safe
            )
        
        # Show the difference between the logistic and linear models fit
        if (nrow(merged_signif) != nrow(unique(fit_data_prevalence_signif)) |
            nrow(merged_signif) != nrow(unique(fit_data_abundance_signif))) {
            stop(
                "Merged significance tables have different associations.
                This is likely a MaAsLin 3 issue 
                due to unexpected data or models."
            )
        }
        
        if (!is.null(new_fit_data_abundance)) {
            fit_data_prevalence_signif_tmp <-
                fit_data_prevalence$results |> collapse::fselect(sel_vec)
            
            colnames(fit_data_prevalence_signif_tmp) <-
                c("feature",
                    "metadata",
                    "value",
                    "name",
                    "logistic_coef",
                    "logistic_null_hypothesis",
                    "logistic",
                    "logistic_qval",
                    "logistic_error")
            
            fit_data_abundance_signif_tmp <-
                new_fit_data_abundance$results[, 
                    c("feature", "metadata", "value",
                    "name", "coef", "null_hypothesis",
                    "pval", "qval", "error")]
            
            colnames(fit_data_abundance_signif_tmp) <-
                c("feature",
                    "metadata",
                    "value",
                    "name",
                    "linear_coef",
                    "linear_null_hypothesis",
                    "linear",
                    "linear_qval",
                    "linear_error")
            
            merged_signif_tmp <-
                collapse::join(
                    collapse::funique(fit_data_prevalence_signif_tmp),
                    collapse::funique(fit_data_abundance_signif_tmp),
                    how = "left",
                    verbose = FALSE,
                    overid = 2,
                    on = c("feature", "metadata", "value", "name")
                )
            
            merged_signif_tmp <- 
                flag_abundance_turned_prevalence(merged_signif_tmp,
                                                max_significance)
            
            overlapping_cols <- c("feature",
                                "metadata",
                                "value",
                                "name",
                                "logistic_coef",
                                "logistic_null_hypothesis",
                                "logistic",
                                "logistic_qval")
            
            merged_signif <- collapse::join(merged_signif,
                                            collapse::fselect(merged_signif_tmp,
                                                              c(overlapping_cols, 
                                                                "logistic_error")),
                                            how = "left",
                                            overid = 2,
                                            on = overlapping_cols,
                                            verbose = FALSE,
                                            suffix = c("_x", "_y"))
            
            # I don't think collapse::join can perfectly mimic the .x/.y
            # suffixing of dplyr joins.
            
            merged_signif <- collapse::fmutate(merged_signif,
                logistic_error = data.table::fcoalesce(merged_signif$logistic_error_y,
                                                       merged_signif$logistic_error_x)
                )
            
            collapse::fselect(merged_signif,
                              c("logistic_error_x",
                                "logistic_error_y")) <- NULL
            
        }
        
        merged_signif <- create_combined_pval(merged_signif,
                                            correction)
        
        return(list(
            append_joint(fit_data_abundance, merged_signif, "abundance"),
            append_joint(fit_data_prevalence, merged_signif, "prevalence")
        ))
    }

# Take logistic or linear results and add on the joint significance
append_joint <- function(outputs, merged_signif, association_type) {
    sel_vec = c("feature", "metadata", "value",
                "name", "pval_joint", "qval_joint")
    
    if (association_type == 'abundance') {
        
        merged_signif <- merged_signif |> collapse::fselect(sel_vec)
        
        collapse::setrename(outputs$results,
                            "pval_individual" = "pval",
                            "qval_individual" = "qval")
        
        # merged_signif <- collapse::join(outputs$results,
        #                                 merged_signif,
        #                                 on = sel_vec[1:4],
        #                                 how = "inner",
        #                                 overid = 2,
        #                                 verbose = FALSE) |> 
        #     collapse::roworderv(sel_vec[1:4]) |> 
        #     collapse::colorderv(sel_vec[1:4])
        
        merged_signif <- merge(outputs$results,
                               merged_signif,
                               by = c("feature", "metadata", "value", "name"))
        
        # o = base::order(merged_signif$feature,
        #                 merged_signif$metadata,
        #                 merged_signif$value,
        #                 merged_signif$name) 
        # # ^ have to use base::order to recapitulate the locale-dependent ordering of base::merge
        # 
        # merged_signif = merged_signif |> 
        #     collapse::ss(o) |> 
        #     collapse::colorderv(neworder = sel_vec[1:4])
        # Too slow. Just let the order be a little changed.
        
    } else if (association_type == 'prevalence') {
        merged_signif <- merged_signif |> 
            collapse::fselect(c(sel_vec,
                                "error" = "logistic_error"))
        
        outputs$results$error <- NULL
        
        collapse::setrename(outputs$results, 
                            "pval_individual" = "pval",
                            "qval_individual" = "qval"
        )
        
        # merged_signif <- collapse::join(outputs$results,
        #                                 merged_signif,
        #                                 on = c("feature", "metadata", 
        #                                        "value", "name"),
        #                                 how = "inner",
        #                                 overid = 2,
        #                                 verbose = FALSE) |> 
        #     collapse::roworderv(sel_vec[1:4]) |> 
        #     collapse::colorderv(sel_vec[1:4])
        
        merged_signif <- merge(outputs$results,
                               merged_signif,
                               by = c("feature", "metadata", "value", "name"))
        
    } else {
        stop("Invalid association_type")
    }
    
    merged_signif <- merged_signif |> 
        collapse::roworderv("qval_individual")
    
    return(merged_signif)
}

# All functions for export in parallelization
function_vec <-
    c(
        "augment_data",
        "safe_deparse",
        "extract_special_predictor",
        "get_fixed_effects",
        "get_character_cols",
        "add_joint_signif",
        "append_joint",
        "check_for_zero_one_obs",
        "check_missing_first_factor_level",
        "check_mffl_one",
        "fit_augmented_logistic",
        "non_augmented",
        "run_group_models",
        "run_ordered_models",
        "fitting_wrap_up"
    )

optimizers <-
    c('nloptwrap', 'nlminbwrap', 'bobyqa', 'Nelder_Mead')
optCtrlList <- list(list(maxeval = 100000),
                    list(maxit = 1500),
                    list(maxfun = 100000),
                    list(maxfun = 100000))

# Check that a formula is provided
check_formulas_valid <- function(formula, random_effects_formula) {
    if (is.null(random_effects_formula)) {
        if (is.null(formula)) {
            logging::logerror(paste("Both formula and 
                                    random_effects_formula are null"))
            stop()
        }
    }
}

# Prepare ranef, modeling, and summary functions for the linear model
choose_ranef_model_summary_funs_linear <- function(random_effects_formula) {
    if (is.null(random_effects_formula)) {
        # Fixed effects only
        ranef_function <- NULL
        
        model_function <-
            function(form,
                    data,
                    weight_scheme = NULL,
                    na.action) {
                return(stats::lm(
                    stats::formula(form),
                    data = data,
                    na.action = na.action
                ))
            }
        
        summary_function <- function(fit, names_to_include) {
            lm_summary <- summary(fit)$coefficients
            
            store_names <- rownames(lm_summary)
            
            if (!all(names_to_include %in% store_names)) {
                # If deficient rank, make sure all rownames are included
                rows_to_add <-
                    names_to_include[!(names_to_include %in% store_names)]
                lm_summary <-
                    rbind(lm_summary, matrix(rep(
                        NaN, 4 * length(rows_to_add)
                    ), nrow = length(rows_to_add)))
                rownames(lm_summary) <-
                    c(store_names, rows_to_add)
            }
            
            para <- as.data.frame(lm_summary) |> 
                ftail(-1) |> 
                collapse::fselect(-3)
            
            para$name <- rownames(lm_summary)[-1]
            
            return(para)
        }
    } else {
        # Random effects
        ranef_function <- lme4::ranef
        
        model_function <-
            function(form,
                    data,
                    weight_scheme = NULL,
                    na.action) {
                index <- 1
                
                while (index < length(optimizers)) {
                    tryCatch({
                        return(
                            lmerTest::lmer(
                                formula = form,
                                data = data,
                                na.action = na.action,
                                control = lme4::lmerControl(
                                    optimizer = optimizers[index],
                                    optCtrl = optCtrlList[[index]]
                                )
                            )
                        )
                    }, warning = function(w) {
                        'warning'
                    }, error = function(e) {
                        'error'
                    })
                    
                    # Something warned or errored if here
                    index <- index + 1
                }
                
                return(
                    lmerTest::lmer(
                        formula = form,
                        data = data,
                        na.action = na.action,
                        control = lme4::lmerControl(
                            optimizer = optimizers[index],
                            optCtrl = optCtrlList[[index]])
                    )
                )
            }
        summary_function <- function(fit, names_to_include) {
            summ = lmerTest:::summary.lmerModLmerTest # Scoped calls - this will be running on daemons

            lm_summary <- summ(fit)$coefficients # Couldn't figure out how to get a scoped call for coef.summary.merMod...

            store_names <- rownames(lm_summary)
            if (!all(names_to_include %in% store_names)) {
                # If deficient rank, make sure all rownames are included
                rows_to_add <-
                    names_to_include[!(names_to_include %in% store_names)]
                lm_summary <-
                    rbind(lm_summary, matrix(rep(
                        NaN, 5 * length(rows_to_add)
                    ), nrow = length(rows_to_add)))
                rownames(lm_summary) <-
                    c(store_names, rows_to_add)
            }
            
            para <- as.data.frame(lm_summary) |> 
                ftail(-1) |> 
                collapse::fselect(-c(3:4))
            # This yields slightly different rownames than [-1, -c(3:4)], but I
            # don't think that matters.
            
            para$name <- rownames(lm_summary)[-1]
            
            return(para)
        }
    }
    return(list("ranef_function" = ranef_function, 
                "model_function" = model_function, 
                "summary_function" = summary_function))
}

# Prepare ranef, modeling, and summary functions for the logistic model
choose_ranef_model_summary_funs_logistic <- function(random_effects_formula,
                                                    strata,
                                                    augment) {
    if (is.null(random_effects_formula)) {
        ranef_function <- NULL
        if (length(strata) > 0) {
            if (augment) {
                model_function <-
                    function(formula,
                            data,
                            weight_scheme = NULL,
                            na.action) {
                        formula <-
                            stats::formula(paste0(
                                safe_deparse(formula),
                                ' + strata(',
                                strata,
                                ')'
                            ))
                        
                        weight_sch_current <- weight_scheme
                        assign("weight_sch_current",
                            weight_scheme,
                            envir = environment(formula))
                        
                        clogit_out <- tryCatch({
                            fit1 <- survival::clogit(
                                stats::formula(formula),
                                data = data,
                                method = "breslow",
                                control = survival::coxph.control(
                                    iter.max = 1000),
                                na.action = na.action,
                                weights = weight_sch_current,
                                robust = FALSE
                            ) # Robust SE seem to be worse with weighting...
                        }, warning = function(w) {
                            'warning'
                        }, error = function(e) {
                            'error'
                        })
                        
                        if (is.character(clogit_out)) {
                            fit1 <- survival::clogit(
                                stats::formula(formula),
                                data = data,
                                method = "breslow",
                                control = survival::coxph.control(
                                    iter.max = 1000),
                                na.action = na.action,
                                weights = weight_sch_current,
                                robust = FALSE
                            ) # Robust SE seem to be worse with weighting...
                            return(fit1)
                        } else  {
                            return(clogit_out)
                        }
                    }
            } else {
                model_function <-
                    function(formula,
                            data,
                            weight_scheme = NULL,
                            na.action) {
                        clogit_out <- tryCatch({
                            fit1 <- survival::clogit(
                                stats::formula(formula),
                                data = data,
                                method = "breslow",
                                control = survival::coxph.control(
                                    iter.max = 1000),
                                na.action = na.action,
                                robust = FALSE
                            ) # Robust SE seem to be worse 
                            # with weighting...)
                        }, warning = function(w) {
                            'warning'
                        }, error = function(e) {
                            'error'
                        })
                        
                        if (is.character(clogit_out)) {
                            fit1 <- survival::clogit(
                                stats::formula(formula),
                                data = data,
                                method = "breslow",
                                control = survival::coxph.control(
                                    iter.max = 1000),
                                na.action = na.action,
                                robust = FALSE
                            ) # Robust SE seem to be 
                            # worse with weighting...)
                            return(fit1)
                        } else  {
                            return(clogit_out)
                        }
                    }
            }
            summary_function <-
                function(fit, names_to_include) {
                    summ = survival:::summary.coxph # Scoped calls - this will be running on daemons
                    
                    lm_summary <- summ(fit)$coefficients
                    
                    store_names <- rownames(lm_summary)
                    if (!all(names_to_include %in% store_names)) {
                        # If deficient rank, 
                        # make sure all rownames are included
                        rows_to_add <-
                            names_to_include[
                                !(names_to_include %in% store_names)]
                        lm_summary <-
                            rbind(lm_summary,
                                matrix(
                                    rep(
                                        NaN,
                                        ncol(lm_summary) * 
                                            length(rows_to_add)
                                    ),
                                    nrow = length(rows_to_add)
                                ))
                        rownames(lm_summary) <-
                            c(store_names, rows_to_add)
                    }
                    
                    if ('robust se' %in% colnames(lm_summary)) {
                        para <-
                            as.data.frame(lm_summary) |> 
                            collapse::fselect(-c(2, 4, 5))
                        # Don't actually use robust SE
                    } else {
                        para <- as.data.frame(lm_summary) |> 
                            collapse::fselect(-c(2, 4))
                    }
                    
                    para$name <- rownames(lm_summary)
                    return(para)
                }
        } else {
            # Fixed effects only
            if (augment) {
                model_function <-
                    function(formula,
                            data,
                            weight_scheme = NULL,
                            na.action) {
                        weight_sch_current <- weight_scheme
                        assign("weight_sch_current",
                            weight_scheme,
                            envir = environment(formula))
                        
                        glm_out <- glm(
                            formula = stats::formula(formula),
                            family = 'binomial',
                            data = data,
                            weights = weight_sch_current,
                            na.action = na.action,
                        )
                        
                        return(glm_out)
                    }
            } else {
                model_function <-
                    function(formula,
                            data,
                            weight_scheme = NULL,
                            na.action) {
                        return(
                            glm(
                                stats::formula(formula),
                                data = data,
                                family = 'binomial',
                                na.action = na.action,
                            )
                        )
                    }
            }
            summary_function <-
                function(fit, names_to_include) {
                    lm_summary <- summary(fit)$coefficients
                    store_names <- rownames(lm_summary)
                    if (!all(names_to_include %in% store_names)) {
                        # If deficient rank, 
                        # make sure all rownames are included
                        rows_to_add <-
                            names_to_include[
                                !(names_to_include %in% store_names)]
                        lm_summary <-
                            rbind(lm_summary,
                                matrix(
                                    rep(NaN, 4 * length(rows_to_add)),
                                    nrow = length(rows_to_add)
                                ))
                        rownames(lm_summary) <-
                            c(store_names, rows_to_add)
                    }
                    para <- as.data.frame(lm_summary) |> 
                        ftail(-1) |> 
                        collapse::fselect(-3)
                    
                    para$name <- rownames(lm_summary)[-1]
                    return(para)
                }
        }
    } else {
        # Random effects
        ranef_function <- lme4::ranef
        
        if (augment) {
            model_function <-
                function(formula,
                        data,
                        weight_scheme = NULL,
                        na.action) {
                    weight_sch_current <- weight_scheme
                    assign("weight_sch_current",
                        weight_scheme,
                        envir = environment(formula))
                    
                    index <- 1
                    
                    while (index < length(optimizers)) {
                        glm_out <- tryCatch({
                            withCallingHandlers({
                                # Catch non-integer # successes first
                                fit1 <- lme4::glmer(
                                    stats::formula(formula),
                                    data = data,
                                    family = 'binomial',
                                    na.action = na.action,
                                    weights = weight_sch_current,
                                    control = lme4::glmerControl(
                                        optimizer = optimizers[index],
                                        optCtrl = optCtrlList[[index]]
                                    )
                                )
                            }, warning = function(w) {
                                if (w$message == 
                                "non-integer #successes in a binomial glm!") {
                                    # Still worked
                                    invokeRestart("muffleWarning")
                                }
                            })
                        }, warning = function(w) {
                            'warning'
                        }, error = function(e) {
                            'error'
                        })
                        
                        # Something warned or errored if here
                        if (is.character(glm_out)) {
                            index <- index + 1
                        } else  {
                            break
                        }
                    }
                    
                    if (is.character(glm_out)) {
                        withCallingHandlers({
                            # Catch non-integer # successes first
                            fit1 <- lme4::glmer(
                                stats::formula(formula),
                                data = data,
                                family = 'binomial',
                                na.action = na.action,
                                weights = weight_sch_current,
                                control = lme4::glmerControl(
                                    optimizer = optimizers[index],
                                    optCtrl = optCtrlList[[index]]
                                )
                            )
                        }, warning = function(w) {
                            if (w$message == 
                                "non-integer #successes in a binomial glm!") {
                                # Still worked
                                invokeRestart("muffleWarning")
                            }
                        })
                        return(fit1)
                    } else  {
                        return(glm_out)
                    }
                }
        } else {
            model_function <-
                function(formula,
                        data,
                        weight_scheme = NULL,
                        na.action) {
                    index <- 1
                    
                    while (index < length(optimizers)) {
                        glm_out <- tryCatch({
                            lme4::glmer(
                                stats::formula(formula),
                                data = data,
                                family = 'binomial',
                                na.action = na.action,
                                control = lme4::glmerControl(
                                    optimizer = optimizers[index],
                                    optCtrl = optCtrlList[[index]]
                                )
                            )
                        }, warning = function(w) {
                            'warning'
                        }, error = function(e) {
                            'error'
                        })
                        
                        # Something warned or errored if here
                        if (is.character(glm_out)) {
                            index <- index + 1
                        } else  {
                            break
                        }
                    }
                    
                    if (is.character(glm_out)) {
                        return(
                            lme4::glmer(
                                stats::formula(formula),
                                data = data,
                                family = 'binomial',
                                na.action = na.action,
                                control = lme4::glmerControl(
                                    optimizer = optimizers[index],
                                    optCtrl = optCtrlList[[index]]
                                )
                            )
                        )
                    } else  {
                        return(glm_out)
                    }
                }
        }
        summary_function <- function(fit, names_to_include) {
            lm_summary <- summary(fit)$coefficients
            
            store_names <- rownames(lm_summary)
            if (!all(names_to_include %in% store_names)) {
                # If deficient rank, make sure all rownames are included
                rows_to_add <-
                    names_to_include[!(names_to_include %in% store_names)]
                lm_summary <-
                    rbind(lm_summary, matrix(rep(
                        NaN, 4 * length(rows_to_add)
                    ), nrow = length(rows_to_add)))
                rownames(lm_summary) <-
                    c(store_names, rows_to_add)
            }
            
            para <- as.data.frame(lm_summary) |> 
                ftail(-1) |> 
                collapse::fselect(-3)
            
            para$name <- rownames(lm_summary)[-1]
            
            return(para)
        }
    }
    return(list("ranef_function" = ranef_function, 
                "model_function" = model_function, 
                "summary_function" = summary_function))
}

# Return blank results if there are zero or one samples
check_for_zero_one_obs <- function(formula,
                                    random_effects_formula,
                                    dat_sub,
                                    groups,
                                    ordereds,
                                    # features,
                                    fn,
                                    x,
                                    model,
                                    feature_specific_covariate_name) {
    match.arg(model, c("linear", "logistic"))
    if (length(unique(dat_sub$expr)) < 2) {
        output <- list()
        
        # List fixed effects that will be included
        names_to_include <-
            get_fixed_effects(formula,
                            random_effects_formula,
                            dat_sub,
                            groups,
                            ordereds,
                            feature_specific_covariate_name)
        
        # Build outputs
        output$para <-
            as.data.frame(matrix(
                NA,
                nrow = length(names_to_include),
                ncol = 3
            ))
        output$para$name <- names_to_include
        
        output$residuals <- NA
        output$fitted <- NA
        if (!(is.null(random_effects_formula)))
            output$ranef <- NA
        output$fit <- NA
        
        output$para = collapse::setColnames(output$para, c('coef', 'stderr' , 'pval', 'name'))
        
        output$para$feature <- fn #colnames(features)[x]
        
        output$para$error <-
            if (model == "logistic") {
                "All logistic values are the same"
            } else {
                "All linear values are the same"
            }
        return(output)
    }
    return (NULL)
}

check_mffl_one = function(v, i, ex, fxf) {
    res = FALSE
    
    if (!is.factor(v)) return(FALSE)
    
    fl = collapse::ffirst(levels(v))
    
    fl_expr = ex[collapse::whichv(v, fl)]
    
    nc_fct = nchar(i)
    
    fxf_nms = substr(fxf, 1, nc_fct) |> collapse::funique()
    
    if (collapse::allNA(fl_expr) && (i %in% fxf_nms)) {
        TRUE
    } else {
        FALSE
    }
}

# If the baseline level is missing for a factor, return blank results
check_missing_first_factor_level <- function(formula,
                                            random_effects_formula,
                                            dat_sub,
                                            groups,
                                            ordereds,
                                            fn,
                                            feature_specific_covariate_name) {
    missing_first_factor_level <- FALSE
    
    fxf <- get_fixed_effects(formula,
                             random_effects_formula,
                             dat_sub,
                             groups,
                             ordereds,
                             feature_specific_covariate_name)
    
    missing_first_factor_level <- any(mapply(FUN = check_mffl_one, 
                                             as.list(dat_sub),
                                             colnames(dat_sub),
                                             MoreArgs = list(ex = dat_sub$expr,
                                                             fxf = fxf)))
    
    if (missing_first_factor_level) {
        output <- list()
        
        # List fixed effects that will be included
        names_to_include <- fxf
        
        # Build outputs
        output$para <-
            as.data.frame(matrix(
                NA,
                nrow = length(names_to_include),
                ncol = 3
            ))
        
        output$para$name <- names_to_include
        
        output$residuals <- NA
        
        output$fitted <- NA
        
        if (!(is.null(random_effects_formula))) output$ranef <- NA
        
        output$fit <- NA
        
        output$para = collapse::setColnames(output$para, c('coef', 'stderr' , 'pval', 'name')) 
        
        output$para$feature <- fn # colnames(features)[x]
        
        output$para$error <-
            "No data points have the baseline factor level"
        return(output)
    }
    
    return (NULL)
}

# Fit the augmented logistic model
fit_augmented_logistic <- function(ranef_function,
                                model_function,
                                formula,
                                random_effects_formula,
                                groups,
                                ordereds,
                                dat_sub,
                                # features,
                                fn,
                                x) {
    warning_message <- NA
    
    error_message <- NA
    
    calling_env <- environment()
    
    mm_input <- NULL
    
    weight_scheme <- NULL
    
    fit1 <- tryCatch({
        withCallingHandlers({
            withCallingHandlers({
                # Catch non-integer # successes first
                formula_new <-
                    stats::formula(paste0(
                        c(
                            safe_deparse(formula),
                            groups,
                            ordereds
                        ),
                        collapse = " + "
                    ))
                
                # Augment data
                augmented_data <-
                    augment_data(formula_new,
                                random_effects_formula,
                                dat_sub)
                mm_input <- augmented_data[["mm_input"]]
                weight_scheme <-
                    augmented_data[["weight_scheme"]]
                formula_new <-
                    augmented_data[["new_formula"]]
                
                # Fit augmented model
                fit1 <-
                    model_function(
                        formula = formula_new,
                        data = mm_input,
                        weight_scheme = weight_scheme,
                        na.action = na.exclude
                    )
                fit1
            }, warning = function(w) {
                if (w$message == 
                    "non-integer #successes in a binomial glm!") {
                    # Still worked
                    invokeRestart("muffleWarning")
                }
            })
        }, warning = function(w) {
            message(sprintf("Feature %s : %s", 
                            fn, w))
            logging::logwarn(paste(
                "Fitting problem for feature",
                x,
                "a warning was issued"
            ))
            
            assign("warning_message",
                conditionMessage(w),
                envir = calling_env)
            invokeRestart("muffleWarning")
        })
    },
    error = function(err) {
        assign("error_message", err$message, envir = calling_env)
        error_obj <-
            structure(list(message = conditionMessage(err)), 
                    class = "try-error")
        return(error_obj)
    })
    
    if (!is.na(error_message)) {
        fit_and_message <- c(list(fit1), list(error_message))
        error_message <- NA
    } else if (!is.na(warning_message)) {
        fit_and_message <- c(list(fit1), list(warning_message))
        warning_message <- NA
    } else {
        fit_and_message <- c(list(fit1), NA)
    }
    
    return(list("fit_and_message" = fit_and_message,
                "weight_scheme" = weight_scheme,
                "mm_input" = mm_input))
}

# Fit linear or non-augmented logistic models
non_augmented <- function(ranef_function,
                        model_function,
                        formula,
                        random_effects_formula,
                        groups,
                        ordereds,
                        dat_sub,
                        # features,
                        fn,
                        x) {

    warning_message <- NA
    
    error_message <- NA
    
    calling_envir <- environment()
    
    fit1 <- tryCatch({
        withCallingHandlers({
            formula_new <-
                stats::formula(paste0(
                    c(safe_deparse(formula), groups, ordereds),
                    collapse = " + "
                ))
            
            fit1 <-
                model_function(formula_new,
                            data = dat_sub,
                            na.action = na.exclude)
            fit1
        }, warning = function(w) {
            message(sprintf("Feature %s : %s", 
                            fn, w))
            logging::logwarn(paste(
                "Fitting problem for feature",
                x,
                "a warning was issued"
            ))
            
            assign("warning_message",
                conditionMessage(w),
                envir = calling_envir)
            invokeRestart("muffleWarning")
        })
    },
    error = function(err) {
        assign("error_message", err$message, envir = calling_envir)
        error_obj <-
            structure(list(message = conditionMessage(err)), 
                    class = "try-error")
        return(error_obj)
    })
    
    if (!is.na(error_message)) {
        fit_and_message <- c(list(fit1), list(error_message))
        error_message <- NA
    } else if (!is.na(warning_message)) {
        fit_and_message <- c(list(fit1), list(warning_message))
        warning_message <- NA
    } else {
        fit_and_message <- c(list(fit1), NA)
    }
    
    return(list("fit_and_message" = fit_and_message,
                "weight_scheme" = NULL,
                "mm_input" = NULL))
}

# Fit models with group effects
run_group_models <- function(ranef_function,
                            model_function,
                            groups,
                            formula,
                            random_effects_formula,
                            model,
                            fit,
                            augment,
                            weight_scheme,
                            dat_sub,
                            output,
                            mm_input) {
    match.arg(model, c("linear", "logistic"))
    
    output$para <- rbind(output$para,
                        setNames(do.call(rbind, lapply(groups, function(group) {
        as.data.frame(setNames(tryCatch({
            withCallingHandlers({
                # Catch non-integer # successes first
                if (is.null(random_effects_formula)) {
                    # Fixed effects
                    if (model == "logistic") {
                        pval_new <- tryCatch({
                            anova(fit, test = 'LRT')[group, 'Pr(>Chi)']
                        },
                        error = function(err) {
                            NA
                        })
                    } else {
                        pval_new <- tryCatch({
                            anova(fit)[group, 'Pr(>F)']
                        },
                        error = function(err) {
                            NA
                        })
                    }
                } else {
                    # Random effects
                    if (model == "logistic") {
                        if (augment) {
                            weight_sch_current <- weight_scheme
                            assign(
                                "weight_sch_current",
                                weight_scheme,
                                envir = environment(formula)
                            )
                            
                            fit_new <-
                                model_function(
                                    formula = stats::update.formula(
                                        stats::formula(fit),
                                        stats::formula(
                                            paste0('~.-', group)
                                        )
                                    ),
                                    data = mm_input,
                                    weight_scheme = 
                                        weight_sch_current,
                                    na.action = na.exclude
                                )
                        } else {
                            fit_new <-
                                model_function(
                                    stats::update.formula(
                                        stats::formula(fit),
                                        stats::formula(
                                            paste0('~.-', group)
                                        )
                                    ),
                                    data = dat_sub,
                                    na.action = na.exclude
                                )
                        }
                        
                        pval_new <-
                            tryCatch({
                                anova(fit_new, fit)[
                                    2, 'Pr(>Chisq)']
                            },
                            error = function(err) {
                                NA
                            })
                    } else {
                group_no_ticks <- gsub("^`|`$", "", group)
                contrast_mat <- matrix(0, ncol = length(
                    lme4::fixef(fit)),
                    nrow = length(levels(dat_sub[[group_no_ticks]])[-1])
                )
                if (length(which(names(lme4::fixef(fit)) %in% paste0(
                    group, levels(dat_sub[[group_no_ticks]])[-1]))) == 0) {
                    stop("group levels not found in fit")
                }
                contrast_mat[seq_along(levels(
                dat_sub[[group_no_ticks]])[-1]),
                which(names(lme4::fixef(fit)) %in% paste0(
                    group, levels(dat_sub[[group_no_ticks]])[-1]))] <-
                diag(1, nrow = length(levels(dat_sub[[group_no_ticks]])[-1]))
                        
                pval_new <- tryCatch({
                    lmerTest::contest(
                        fit, 
                        contrast_mat,
                        rhs = 
                            rep(0, length(levels(
                                dat_sub[[group_no_ticks]])[-1])))[['Pr(>F)']]
                },
                error = function(err) {
                    NA
                })
                    }
                }
                
                tmp_output <- list(NA, NA, pval_new, group)
                tmp_output
            }, warning = function(w) {
                if (w$message == "non-integer #successes in a binomial glm!") {
                    invokeRestart("muffleWarning")
                }
            })
        }, warning = function(w) {
            return(list(NA, NA, NA, group))
        }, error = function(err) {
            return(list(NA, NA, NA, group))
        }), names(output$para)), check.names = FALSE)
    })), names(output$para)))
    
    tmp_rownames <- rownames(output$para)
    
    output$para <- data.frame(lapply(output$para, function(x) {
        if (is.list(x)) unlist(x) else x
    }), check.names = FALSE)
    
    rownames(output$para) <- tmp_rownames
    
    rownames(output$para) <- c(rownames(output$para)[
        -seq((nrow(output$para) - length(groups) + 1), nrow(output$para))], 
        groups)
    return(output)
}

# Fit models with ordered effects
run_ordered_models <- function(ranef_function,
                            model_function,
                            ordereds,
                            fit_and_message,
                            formula,
                            random_effects_formula,
                            model,
                            fit,
                            augment,
                            weight_scheme,
                            dat_sub,
                            output) {
    match.arg(model, c("linear", "logistic"))
    output$para <- rbind(output$para,
                        setNames(do.call(rbind, lapply(ordereds, 
                            function(ordered) {
        ordered_no_ticks <- gsub("^`|`$", "", ordered)
        ordered_levels <- paste0(ordered,
            levels(dat_sub[[ordered_no_ticks]])[-1])
        as.data.frame(setNames(tryCatch({
            withCallingHandlers({
                # Catch non-integer # successes first
                if (is.null(random_effects_formula)) {
                    # Fixed effects
                    if (any(!ordered_levels %in% names(coef(fit)))) {
                        fit_and_message[[
                            length(fit_and_message)]] <-
                            "Error: Some ordered levels are missing"
                        stop("Some ordered levels are missing")
                    }
                    
                    contrast_mat <-
                        matrix(0, ncol = length(coef(fit, complete = FALSE)),
                        nrow = length(levels(dat_sub[[ordered_no_ticks]])[-1])
                        )
                    
                    cols_to_add_1s <-
                        which(names(coef(fit, complete = FALSE)) %in% 
                                ordered_levels)
                    contrast_mat[1, cols_to_add_1s[1]] <- 1
                    for (i in seq_along(cols_to_add_1s[-1])) {
                        contrast_mat[i + 1, cols_to_add_1s[-1][i]] <- 1
                        contrast_mat[i + 1, cols_to_add_1s[i]] <- -1
                    }
                    
                    pvals_new <- vapply(seq(nrow(contrast_mat)), 
                                        function(row_num) {
                        contrast_vec <- t(matrix(contrast_mat[row_num, ]))
                        tryCatch({
                            summary(multcomp::glht(fit, linfct = contrast_vec, 
                rhs = 0, 
                coef. = function(x) { coef(x, complete = FALSE) }))$test$pvalues
                        }, error = function(err) { NA })
                    }, numeric(1))
                    
                    coefs_new <- vapply(seq(nrow(contrast_mat)), 
                                        function(row_num) {
                        contrast_vec <- t(matrix(contrast_mat[row_num, ]))
                        tryCatch({
                            summary(multcomp::glht(fit, linfct = contrast_vec, 
                rhs = 0, 
                coef. = function(x) { 
                    coef(x, complete = FALSE) }))$test$coefficients
                        }, error = function(err) { NA })
                    }, numeric(1))
                    
                    sigmas_new <- vapply(seq(nrow(contrast_mat)), 
                                        function(row_num) {
                        contrast_vec <- t(matrix(contrast_mat[row_num, ]))
                        tryCatch({
                            summary(multcomp::glht(fit, linfct = contrast_vec, 
                rhs = 0, 
                coef. = function(x) { coef(x, complete = FALSE) }))$test$sigma
                        }, error = function(err) { NA })
                    }, numeric(1))
                } else {
                    # Random effects
                    if (any(!ordered_levels %in% names(lme4::fixef(fit)))) {
                        fit_and_message[[length(fit_and_message)]] <-
                            "Error: Some ordered levels are missing"
                        stop("Some ordered levels are missing")
                    }
                    
                    ordered_no_ticks <- gsub("^`|`$", "", ordered)
                    contrast_mat <-
                        matrix(0, ncol = length(lme4::fixef(fit)),
                        nrow = length(levels(dat_sub[[ordered_no_ticks]])[-1])
                        )
                    
                    cols_to_add_1s <-
                        which(names(lme4::fixef(fit)) %in% ordered_levels)
                    contrast_mat[1, cols_to_add_1s[1]] <- 1
                    for (i in seq_along(cols_to_add_1s[-1])) {
                        contrast_mat[i + 1, cols_to_add_1s[-1][i]] <- 1
                        contrast_mat[i + 1, cols_to_add_1s[i]] <- -1
                    }
                    
                    if (model == "logistic") {
        pvals_new <- vapply(seq(nrow(contrast_mat)), 
            function(row_num) {
                contrast_vec <- t(matrix(contrast_mat[row_num, ]))
                tryCatch({
                    summary(multcomp::glht(fit, linfct = contrast_vec, 
                                        rhs = 0))$test$pvalues
                }, error = function(err) { NA })
            }, numeric(1))
                        
        coefs_new <- vapply(seq(nrow(contrast_mat)), 
            function(row_num) {
                contrast_vec <- t(matrix(contrast_mat[row_num, ]))
                tryCatch({
                    summary(multcomp::glht(fit, linfct = contrast_vec, 
                                        rhs = 0))$test$coefficients
                }, error = function(err) { NA })
            }, numeric(1))
                        
        sigmas_new <- vapply(seq(nrow(contrast_mat)), 
            function(row_num) {
                contrast_vec <- t(matrix(contrast_mat[row_num, ]))
                tryCatch({
                    summary(multcomp::glht(fit, linfct = contrast_vec, 
                    rhs = 0))$test$sigma
                }, error = function(err) { NA })
            }, numeric(1))
                    } else {
        pvals_new <- vapply(seq(nrow(contrast_mat)), 
            function(row_num) {
                contrast_vec <- t(matrix(contrast_mat[row_num, ]))
                tryCatch({
                    lmerTest::contest(fit,
                                        matrix(
                                        contrast_vec,
                                        TRUE
                                        ), rhs = 0)[['Pr(>F)']]
                }, error = function(err) { NA })
            }, numeric(1))
        
        coefs_new <- vapply(seq(nrow(contrast_mat)), 
            function(row_num) {
                contrast_vec <- t(matrix(contrast_mat[row_num, ]))
                tryCatch({
                    contrast_vec %*% lme4::fixef(fit)
                }, error = function(err) { NA })
            }, numeric(1))
        
        sigmas_new <- vapply(seq(nrow(contrast_mat)), 
            function(row_num) {
                contrast_vec <- t(matrix(contrast_mat[row_num, ]))
                tryCatch({
                    sqrt((contrast_vec %*% vcov(fit) %*% t(contrast_vec))[1, 1])
                }, error = function(err) { NA })
            }, numeric(1))
            }
                }
                
                tmp_output <-
                    data.frame(coefs_new,
                        sigmas_new,
                        pvals_new,
                        ordered_levels
                        )
                rownames(tmp_output) <- ordered_levels
                tmp_output
            }, warning = function(w) {
                if (w$message == 
                    "non-integer #successes in a binomial glm!") {
                    # Still worked
                    invokeRestart("muffleWarning")
                }
            })
        }, warning = function(w) {
            tmp_output <- data.frame(matrix(nrow = length(ordered_levels),
                                            ncol = 4))
            tmp_output[seq(1,length(ordered_levels)), seq(3)] <- NA
            tmp_output[seq(1,length(ordered_levels)), 4] <- ordered_levels
            rownames(tmp_output)[seq(1,length(ordered_levels))] <-
                ordered_levels
            
            return(tmp_output)
        }, error = function(err) {
            tmp_output <- data.frame(matrix(nrow = length(ordered_levels),
                                            ncol = 4))
            tmp_output[seq(1,length(ordered_levels)), seq(3)] <- NA
            tmp_output[seq(1,length(ordered_levels)), 4] <- ordered_levels
            rownames(tmp_output)[seq(1,length(ordered_levels))] <-
                ordered_levels
            
            return(tmp_output)
        }), names(output$para)), check.names = FALSE)
    })), names(output$para)))
    
    return(output)
}

# Check whether the model fit properly and return modeling outputs
fitting_wrap_up <- function(fit_properly,
                            fit_and_message,
                            output,
                            fit,
                            random_effects_formula,
                            metadata,
                            median_comparison,
                            save_models,
                            formula,
                            dat_sub,
                            groups,
                            ordereds,
                            # features,
                            fn,
                            x,
                            ranef_function,
                            feature_specific_covariate_name) {
    if (fit_properly) {
        output$residuals <- stats::residuals(fit)
        output$fitted <- stats::fitted(fit)
        # ^ It would be better to unname these and put them together in a
        # co-indexed data frame. TODO:
        
        if (!(is.null(random_effects_formula))) {
            # Returns a list with a table for each random effect
            l <- ranef_function(fit)
            
            # Rename rows as random effect labels 
            # if only random intercepts
            rl_chk = length(l) == 1 &
                ncol(l[[1]]) == 1 &
                colnames(l[[1]])[1] == "(Intercept)"
            
            if (rl_chk) {
                
                d <- as.vector(unlist(l))
                
                names(d) <- row.names(l[[1]])
                
                d[setdiff(collapse::funique(metadata[, names(l)]), 
                          names(d))] <- NA
                
                d <- d[collapse::radixorder(collapse::funique(metadata[, names(l)]))]
                
                output$ranef <- d
                
            } else {
                # Otherwise return the random effects list
                output$ranef <- l
            }
        }
        
        if (median_comparison) {
            output$fit <- fit
        } else {
            if (save_models) {
                output$fit <- fit
            } else {
                output$fit <- NA
            }
        }
        
    } else {
        # Fitting issue
        logging::logwarn(paste("Fitting problem for feature",
                            x,
                            "returning NA"))
        
        names_to_include <-
            get_fixed_effects(formula,
                            random_effects_formula,
                            dat_sub,
                            groups,
                            ordereds,
                            feature_specific_covariate_name)
        
        # Store NA values for missing outputs
        output$para <-
            as.data.frame(matrix(
                NA,
                nrow = length(names_to_include),
                ncol = 3
            ))
        
        output$para$name <- names_to_include
        
        output$residuals <- NA
        
        output$fitted <- NA
        
        if (!(is.null(random_effects_formula))) output$ranef <- NA
        
        output$fit <- NA
    }
    
    if (is(output$fit, "merMod")) {
        # Delete the model frame in the object to reduce the size of the output
        # fit. It's just a redundant copy of the metadata. Same for these
        # elements of the reponse module in fit@resp.
        output$fit@frame <- data.frame()
        # For some reason this makes the models_*.rds file a lot larger TODO:
        # identify why or better yet save only what's needed.
        
        # The response slot is big too.
        # output$fit@resp$mu <- vector("numeric")
        # output$fit@resp$y <- vector("numeric")

        # TODO: check if this can/should be done for the other model types like
        # lm/glm/clogit The "y" element of clogit objects looks suspiciously
        # large. It would be best of all to just extract the
        # parameters/statistics of interest and drop everything else. If one
        # really wants access to the full fit object it would be best to save
        # individual feature fits to their own RDS files under fits/models_*/
        # from within func_to_run() and avoid returning it. That would save
        # time, memory, and avoid creating huge models_logistic.rds files.

        # The fits also contain redundant copies of the fitted / residual values
        # too. Those can/should be deleted too probably.

        # OOP - What I want is a banana. What I get is a gorilla holding a banana.
    }
    
    colnames(output$para) <-
        c('coef', 'stderr' , 'pval', 'name')
    
    output$para$feature <- fn # colnames(features)[x]
    
    output$para$error <-
        fit_and_message[[length(fit_and_message)]]
    
    output$para$error[is.na(output$para$error) &
                        is.na(output$para$pval)] <-
        'Fitting error (NA p-value returned from fitting procedure)'
    
    return(output)
}

# Run the median comparison procedure for ordered predictors
run_median_comparison_ordered <- function(paras_sub,
                                        fits,
                                        ordereds,
                                        metadata,
                                        random_effects_formula,
                                        median_comparison_threshold,
                                        metadata_variable,
                                        pvals_new,
                                        cur_median,
                                        model) {
    match.arg(model, c("linear", "logistic"))
    
    ordered <- ordereds[which(startsWith(
        metadata_variable, ordereds))]
    
    use_this_coef <- !is.na(paras_sub$pval) & paras_sub$pval < 0.95
    
    n_coefs <- nrow(paras_sub)
    
    sigmas <- paras_sub$stderr
    
    coefs <- paras_sub$coef
    
    sigma_sq_med <- var(coefs[use_this_coef], na.rm=TRUE)
    
    # Variance from asymptotic distribution
    sd_median <- sqrt(0.25 * 2 * base::pi * sigma_sq_med / sum(use_this_coef))
    
    # MC for covariance
    nsims <- 10000
    
    sim_results <- replicate(nsims, {
        sim_coefs <- rnorm(n_coefs, coefs, sigmas)
        sim_median <- median(sim_coefs[use_this_coef])
        c(sim_median, sim_coefs)
    })
    
    sim_medians <- sim_results[1, ]
    
    all_sims <- sim_results[-1, , drop = FALSE]
    
    cov_adjust <- apply(all_sims, 1, function(x){cov(x, sim_medians)})
    
    # Necessary offsets for contrast testing
    offsets_to_test <- abs(cur_median - coefs) * 
        sqrt((sigmas^2) / (sigmas^2+ sd_median^2 - 2 * cov_adjust)) + coefs
    
    pvals_new <- c(pvals_new,
        vapply(seq(nrow(paras_sub)), function(row_index) {
            feature <- paras_sub$feature[row_index]
            if (is.null(random_effects_formula)) {
                # Fixed effects
                cur_fit <- fits[[feature]]
                
                if (!metadata_variable %in% names(coef(cur_fit))) {
                    return(NA)
                }
                
                if (any(!paste0(ordered, 
                                levels(
                                    metadata[[ordered]])[-1]) %in% 
                        names(coef(cur_fit)))) {
                    return(NA)
                }
                
                mm_variable <-
                    model.matrix(cur_fit)[, metadata_variable]
                if (any(!unique(
                    mm_variable[!is.na(mm_variable)]) %in% c(0, 1))) {
                    median_comparison_threshold_updated <-
                        median_comparison_threshold / 
                        sd(mm_variable)
                } else {
                    median_comparison_threshold_updated <- 
                        median_comparison_threshold
                }
                
                if (is.na(coef(cur_fit, complete = FALSE)[
                    which(names(coef(
                        cur_fit, complete = FALSE
                    )) == metadata_variable)])) {
                    pval_new_current <- NA
                } else if (abs(coef(cur_fit, complete = FALSE)[
                    which(names(coef(
                        cur_fit, complete = FALSE
                    )) == metadata_variable)] -
                    cur_median) < median_comparison_threshold_updated) {
                    pval_new_current <- 1
                } else {
                    contrast_mat <-
                        matrix(0, ncol = length(
                            coef(cur_fit, complete = FALSE)),
                            nrow = length(levels(metadata[[ordered]])[-1])
                        )
                    
                    cols_to_add_1s <-
                        which(names(coef(cur_fit, complete = FALSE)) %in% 
                            paste0(ordered, levels(metadata[[ordered]])[-1]))
                    contrast_mat[1, cols_to_add_1s[1]] <- 1
                    for (i in seq_along(cols_to_add_1s[-1])) {
                        contrast_mat[i + 1, cols_to_add_1s[-1][i]] <- 1
                        contrast_mat[i + 1, cols_to_add_1s[i]] <- -1
                    }
                    
                    contrast_vec <-
                        t(matrix(contrast_mat[which(paste0(
                            ordered, levels(
                                metadata[[ordered]]
                            )[-1]) == metadata_variable),]))
                    pval_new_current <-
                        tryCatch({
                            summary(
                                multcomp::glht(
                                    cur_fit,
                                    linfct = contrast_vec,
                                    rhs = offsets_to_test[row_index],
                                    coef. = function(x) {
                                        coef(x, complete = FALSE)
                                    }
                                )
                            )$test$pvalues
                        },
                        error = function(err) {
                            NA
                        })
                }
                
                return(pval_new_current)
            } else {
                # Random effects
                cur_fit <- fits[[feature]]
                
                if (!metadata_variable %in% 
                    names(lme4::fixef(cur_fit))) {
                    NA
                }
                
                if (any(!paste0(ordered, 
                                levels(metadata[[
                                    ordered]])[-1]) %in% 
                        names(lme4::fixef(cur_fit)))) {
                    NA
                }
                
                mm_variable <-
                    model.matrix(cur_fit)[, metadata_variable]
                if (any(
                    !unique(mm_variable[!is.na(mm_variable)]) %in% 
                    c(0, 1))) {
                    median_comparison_threshold_updated <-
                        median_comparison_threshold / 
                        sd(mm_variable)
                } else {
                    median_comparison_threshold_updated <- 
                        median_comparison_threshold
                }
                
                contrast_mat <-
                    matrix(
                        0,
                        ncol = length(lme4::fixef(cur_fit)),
                        nrow = length(levels(
                            metadata[[ordered]])[-1])
                    )
                cols_to_add_1s <-
                    which(names(lme4::fixef(cur_fit)) %in% 
                        paste0(ordered, levels(
                        metadata[[ordered]])[-1]))
                contrast_mat[1, cols_to_add_1s[1]] <- 1
                for (i in seq_along(cols_to_add_1s[-1])) {
                    contrast_mat[i + 1, cols_to_add_1s[-1][i]] <- 1
                    contrast_mat[i + 1, cols_to_add_1s[i]] <- -1
                }
                contrast_vec <-
                    t(matrix(
                        contrast_mat[
                            which(paste0(
                                ordered, levels(metadata[[
                                    ordered]])[-1]) == 
                                    metadata_variable),]))
                
                if (model == "logistic") {
                    if (is.na(lme4::fixef(cur_fit)[
                        which(names(lme4::fixef(
                            cur_fit
                        )) == metadata_variable)])) {
                        pval_new_current <- NA
                    } else if (abs(lme4::fixef(cur_fit)[
                        which(names(lme4::fixef(
                            cur_fit
                        )) == metadata_variable)] -
                        cur_median) < 
                        median_comparison_threshold_updated) {
                        pval_new_current <- 1
                    } else {
                        pval_new_current <-
                            tryCatch({
                                summary(
                                    multcomp::glht(
                                        cur_fit,
                                        linfct = contrast_vec,
                                        rhs = offsets_to_test[row_index]
                                    )
                                )$test$pvalues
                            },
                            error = function(err) {
                                NA
                            })
                    } 
                    return(pval_new_current)
                } else {
                    if (is.na(lme4::fixef(cur_fit)[
                        which(names(lme4::fixef(
                            cur_fit
                        )) == metadata_variable)])) {
                        pval_new_current <- NA
                    } else if (abs(lme4::fixef(cur_fit)[
                        which(names(lme4::fixef(
                            cur_fit
                        )) == metadata_variable)] -
                        cur_median) < 
                        median_comparison_threshold_updated) {
                        pval_new_current <- 1
                    } else {
                        pval_new_current <-
                            tryCatch({
                                lmerTest::contest(cur_fit,
                                            matrix(contrast_vec, TRUE),
                                            rhs = offsets_to_test[row_index])[[
                                                'Pr(>F)']]
                            },
                            error = function(err) {
                                NA
                            })
                    }
                    
                    return(pval_new_current)
                }
            }
        }, numeric(1)))
    
    return(pvals_new)
}

# Run the median comparison procedure for non-ordered predictors
run_median_comparison_general <- function(paras_sub,
                                        fits,
                                        metadata,
                                        random_effects_formula,
                                        median_comparison_threshold,
                                        metadata_variable,
                                        pvals_new,
                                        cur_median,
                                        model) {
    match.arg(model, c("linear", "logistic"))
    use_this_coef <- !is.na(paras_sub$pval) & paras_sub$pval < 0.95
    
    n_coefs <- nrow(paras_sub)
    sigmas <- paras_sub$stderr
    coefs <- paras_sub$coef
    sigma_sq_med <- var(coefs[use_this_coef], na.rm=TRUE)
    
    # Variance from asymptotic distribution
    sd_median <- sqrt(0.25 * 2 * base::pi * sigma_sq_med / sum(use_this_coef))

    # MC for covariance
    nsims <- 10000
    sim_results <- replicate(nsims, {
        sim_coefs <- rnorm(n_coefs, coefs, sigmas)
        sim_median <- median(sim_coefs[use_this_coef])
        c(sim_median, sim_coefs)
    })
    
    sim_medians <- sim_results[1, ]
    all_sims <- sim_results[-1, , drop = FALSE]
    cov_adjust <- apply(all_sims, 1, function(x){cov(x, sim_medians)})
    
    # Necessary offsets for contrast testing
    offsets_to_test <- abs(cur_median - coefs) * 
        sqrt((sigmas^2) / (sigmas^2+ sd_median^2 - 2 * cov_adjust)) + coefs
    
    pvals_new <- c(pvals_new,
        vapply(seq(nrow(paras_sub)), function(row_index) {
            feature <- paras_sub$feature[row_index]
            if (is.null(random_effects_formula)) {
                # Fixed effects
                cur_fit <- fits[[feature]]
                
                if (!metadata_variable %in% names(coef(cur_fit))) {
                    return(NA)
                }
                
                mm_variable <-
                    model.matrix(cur_fit)[, metadata_variable]
                if (any(
                    !unique(mm_variable[!is.na(mm_variable)]) %in% 
                    c(0, 1))) {
                    median_comparison_threshold_updated <-
                        median_comparison_threshold / 
                        sd(mm_variable)
                } else {
                    median_comparison_threshold_updated <- 
                        median_comparison_threshold
                }
                
                contrast_vec <-
                    rep(0, length(coef(
                        cur_fit, complete = FALSE
                    )))
                contrast_vec[which(names(coef(
                    cur_fit, complete = FALSE
                )) == metadata_variable)] <- 1
                
                if (is.na(coef(cur_fit, complete = FALSE)[
                    which(names(coef(
                        cur_fit, complete = FALSE
                    )) == metadata_variable)])) {
                    pval_new_current <- NA
                } else if (abs(coef(cur_fit, complete = FALSE)[
                    which(names(coef(
                        cur_fit, complete = FALSE
                    )) == metadata_variable)] -
                    cur_median) < median_comparison_threshold_updated) {
                    pval_new_current <- 1
                } else {
                    pval_new_current <-
                        tryCatch({
                            summary(
                                multcomp::glht(
                                    cur_fit,
                                    linfct = matrix(contrast_vec, 
                                                    TRUE),
                                    rhs = offsets_to_test[row_index],
                                    coef. = function(x) {
                                        coef(x, complete = FALSE)
                                    }
                                )
                            )$test$pvalues[1]
                        },
                        error = function(err) {
                            NA
                        })
                }
                
                return(pval_new_current)
            } else {
                # Random effects
                cur_fit <- fits[[feature]]
                
                if (!metadata_variable %in% 
                    names(lme4::fixef(cur_fit))) {
                    return(NA)
                }
                
                mm_variable <-
                    model.matrix(cur_fit)[, metadata_variable]
                if (any(!unique(
                    mm_variable[!is.na(mm_variable)]) %in% 
                    c(0, 1))) {
                    median_comparison_threshold_updated <-
                        median_comparison_threshold / 
                        sd(mm_variable)
                } else {
                    median_comparison_threshold_updated <- 
                        median_comparison_threshold
                }
                
                contrast_vec <- rep(0, length(lme4::fixef(cur_fit)))
                contrast_vec[which(names(lme4::fixef(cur_fit)) == 
                            metadata_variable)] <- 1
                
                if (model == "logistic") {
                    if (is.na(lme4::fixef(cur_fit)[
                        which(names(lme4::fixef(
                            cur_fit
                        )) == metadata_variable)])) {
                        pval_new_current <- NA
                    } else if (abs(lme4::fixef(cur_fit)[
                        which(names(lme4::fixef(
                            cur_fit
                        )) == metadata_variable)] -
                        cur_median) < 
                        median_comparison_threshold_updated) {
                        pval_new_current <- 1
                    } else {
                        pval_new_current <-
                            tryCatch({
                                summary(
                                    multcomp::glht(
                                        cur_fit,
                                        linfct = matrix(
                                            contrast_vec, TRUE),
                                        rhs = offsets_to_test[row_index]
                                    )
                                )$test$pvalues[1]
                            },
                            error = function(err) {
                                NA
                            })
                    }
                    
                    return(pval_new_current)
                } else {
                    if (is.na(lme4::fixef(cur_fit)[
                        which(names(lme4::fixef(
                            cur_fit
                        )) == metadata_variable)])) {
                        pval_new_current <- NA
                    } else if (abs(lme4::fixef(cur_fit)[
                        which(names(lme4::fixef(
                            cur_fit
                        )) == metadata_variable)] -
                        cur_median) < 
                        median_comparison_threshold_updated) {
                        pval_new_current <- 1
                    } else {
                        pval_new_current <-
                            tryCatch({
                                lmerTest::contest(
                                    cur_fit,
                                    matrix(contrast_vec, TRUE),
                                    rhs = 
                                    offsets_to_test[row_index])[['Pr(>F)']]
                            },
                            error = function(err) {
                                NA
                            })
                    }
                    
                    return(pval_new_current)
                }
            }
        }, numeric(1)))
    return(pvals_new)
}

# Run the median comparison for all variables
run_median_comparison <- function(paras,
                                fits,
                                ordereds,
                                groups,
                                metadata,
                                random_effects_formula,
                                median_comparison_threshold,
                                subtract_median,
                                model) {
    match.arg(model, c("linear", "logistic"))
    logging::loginfo("Performing tests against medians")
    
    
    if (length(ordereds) > 0) {
        ordered_levels <- unlist(lapply(ordereds, function(ordered) {
            unlist(paste0(ordered, levels(metadata[[ordered]])[-1]))
        }))
    } else {
        ordered_levels <- c()
    }
    
    final_paras <-
        paras[!is.na(paras$error) | paras$name %in% groups,]
    final_paras$null_hypothesis <- rep(NA, nrow(final_paras))
    paras <-
        paras[is.na(paras$error) & !paras$name %in% groups,]
    
    process_metadata_variable <- function(metadata_variable,
                                        paras,
                                        fits,
                                        ordered_levels,
                                        ordereds,
                                        metadata,
                                        random_effects_formula, 
                                        median_comparison_threshold, 
                                        subtract_median, 
                                        model) {
        
        paras_sub <- paras[paras$name == metadata_variable, ]
        cur_median <- median(paras_sub$coef[!is.na(paras_sub$pval) & 
                                        paras_sub$pval < 0.95], na.rm = TRUE)
        
        if (is.na(cur_median)) {
            pvals_new <- rep(NA, nrow(paras_sub))
            coefs_new <- paras_sub$coef
        } else {
            pvals_new <- vector()
            if (subtract_median) {
                coefs_new <- paras_sub$coef - cur_median
            } else {
                coefs_new <- paras_sub$coef
            }
            
            if (metadata_variable %in% ordered_levels) {
                pvals_new <- run_median_comparison_ordered(
                    paras_sub, 
                    fits, 
                    ordereds, 
                    metadata, 
                    random_effects_formula, 
                    median_comparison_threshold, 
                    metadata_variable, 
                    pvals_new, 
                    cur_median, 
                    model
                )
            } else {
                pvals_new <- run_median_comparison_general(
                    paras_sub, 
                    fits, 
                    metadata, 
                    random_effects_formula, 
                    median_comparison_threshold, 
                    metadata_variable, 
                    pvals_new, 
                    cur_median, 
                    model
                )
            }
        }
        
        paras_sub$error <- data.table::fifelse(
            is.na(pvals_new) & !is.na(paras_sub$pval), 
            "P-value became NA in median comparison, try rerunning without the median comparison", 
            paras_sub$error
        )
        
        paras_sub$pval <- pvals_new
        paras_sub$coef <- coefs_new
        paras_sub$null_hypothesis <- cur_median
        
        return(paras_sub)
    }
    
    final_paras_list <- lapply(unique(paras$name), function(metadata_variable) {
        process_metadata_variable(metadata_variable,
                                paras,
                                fits,
                                ordered_levels,
                                ordereds,
                                metadata,
                                random_effects_formula,
                                median_comparison_threshold,
                                subtract_median,
                                model)
    })
    
    final_paras <- collapse::rowbind(final_paras, 
                                     collapse::rowbind(final_paras_list))

    paras <- final_paras
    return(paras)
}

# fit the data using the model selected and applying the correction
fit.model <- function(features,
                    metadata,
                    model,
                    formula = NULL,
                    random_effects_formula = NULL,
                    correction = "BH",
                    save_models = FALSE,
                    small_random_effects = FALSE,
                    augment = FALSE,
                    cores = 1,
                    median_comparison = FALSE,
                    median_comparison_threshold = 0,
                    subtract_median = FALSE,
                    feature_specific_covariate = NULL,
                    feature_specific_covariate_name = NULL,
                    feature_specific_covariate_record = NULL) {
    match.arg(model, c("linear", "logistic"))
    match.arg(correction,
            c("BH", "holm", "hochberg", "hommel", "bonferroni", "BY"))
    check_formulas_valid(formula, random_effects_formula)
    formula <- stats::as.formula(formula, env = baseenv())

    extract_out <- extract_special_predictor(formula, 'group')
    formula <- extract_out[[1]]
    groups <- extract_out[[2]]

    extract_out <- extract_special_predictor(formula, 'ordered')
    formula <- extract_out[[1]]
    ordereds <- extract_out[[2]]

    extract_out <- extract_special_predictor(formula, 'strata')
    formula <- extract_out[[1]]
    strata <- extract_out[[2]]
    
    if (small_random_effects & model == 'logistic') {
        fixed_part <- lme4::nobars(formula)
        random_terms <- reformulas::findbars(formula)
        if (length(random_terms) > 0) {
            grouping <- unique(
                vapply(random_terms, function(x) deparse(x[[3]]), 
                    FUN.VALUE = character(length(random_terms))))
            grouping <- data.table::fifelse(grouping == make.names(grouping),
                grouping, paste0('`', grouping, '`'))
            new_formula <- as.formula(
                paste(safe_deparse(fixed_part), "+", 
                    paste(grouping, collapse = " + ")),
                env = environment(formula)
            )
            environment(new_formula) <- environment(formula)
            formula <- new_formula
        } else {
            stop("small_random_effects=TRUE but no random effects specified")
        }
        random_effects_formula <- NULL
    }
    
    if (length(strata) > 0 & !is.null(random_effects_formula)){
        stop(
            "Strata and random effects cannot be combined. Please only use
            random effects if you have multiple grouping categories."
        )
    }
    
    if (length(strata) > 0 & model == 'linear') {
        formula <-
            stats::formula(paste0(safe_deparse(formula), ' + (1 | ', strata, ')'))
        random_effects_formula <- formula
    }
    
    ################
    # Linear Model #
    ################
    
    if (model == "linear")
        fun_list <- choose_ranef_model_summary_funs_linear(
            random_effects_formula)
    
    ##################
    # Logistic Model #
    ##################
    
    if (model == "logistic") {
        fun_list <- choose_ranef_model_summary_funs_logistic(
            random_effects_formula,
            strata = strata,
            augment = augment)
    }
    
    ranef_function <- fun_list$ranef_function
    model_function <- fun_list$model_function
    summary_function <- fun_list$summary_function
    
    #######################################
    # Init cluster for parallel computing #
    #######################################
    
    # cluster <- NULL
    # if (cores > 1) {
    #     logging::loginfo("Creating cluster of %s R processes", cores)
    #     cluster <- parallel::makeCluster(cores)
        # parallel::clusterExport(cluster, c(ls(), function_vec),
                                # envir = environment())
    # }
    
    ##############################
    # Apply per-feature modeling #
    ##############################
   
    fit_vars = all.vars(formula) |> tail(-1)
    
    small_meta = metadata |>
        collapse::fselect(fit_vars)
    
    if (mirai::daemons_set()) {
        mirai::everywhere({}, 
                          metadata = small_meta,
                          random_effects_formula = random_effects_formula,
                          groups = groups,
                          ordereds = ordereds,
                          model = model,
                          feature_specific_covariate = feature_specific_covariate,
                          feature_specific_covariate_name = feature_specific_covariate_name,
                          formula = formula,
                          summary_function = summary_function,
                          ranef_function = ranef_function,
                          model_function = model_function,
                          augment = augment,
                          median_comparison = median_comparison,
                          get_fixed_effects = get_fixed_effects,
                          safe_deparse = safe_deparse,
                          save_models = save_models,
                          ftail = ftail,
                          check_for_zero_one_obs = check_for_zero_one_obs,
                          check_missing_first_factor_level = check_missing_first_factor_level,
                          fit_augmented_logistic = fit_augmented_logistic,
                          non_augmented = non_augmented,
                          run_group_models = run_group_models,
                          run_ordered_models = run_ordered_models,
                          get_character_cols = get_character_cols,
                          fitting_wrap_up = fitting_wrap_up,
                          optimizers = optimizers,
                          optCtrlList = optCtrlList)
        
        mirai::everywhere({
            requireNamespace("stats")
            requireNamespace("methods")
        })
    }
    
    func_to_run <- function(fv, fn, fi) {
        # Extract Features One by One
        featuresVector <- fv 
        
        logging::loginfo("Fitting model to feature number %d, %s",
                         fi,
                         fn)
        
        # Make fitting matrix of features and metadata
        if (!is.null(feature_specific_covariate)) {
            covariateVector <- feature_specific_covariate[, fi]
            
            dat_sub <-
                data.frame(
                    expr = as.numeric(featuresVector),
                    feature_specific_covariate = covariateVector,
                    metadata, 
                    check.names = FALSE
                )
            
            colnames(dat_sub)[colnames(dat_sub) == 
                            'feature_specific_covariate'] <-
                feature_specific_covariate_name
        } else {
            dat_sub <- data.frame(expr = 
                as.numeric(featuresVector), metadata, 
                check.names = FALSE)
        }
        
        # 0 or 1 observations
        zero_one_out <- check_for_zero_one_obs(formula,
            random_effects_formula,
            dat_sub,
            groups,
            ordereds,
            # features, # TODO: avoid handing the full matrix to all daemons... Will need to fix functions like this one
            fn = fn,
            fi,
            model,
            feature_specific_covariate_name)
        
        if (!is.null(zero_one_out)) {
            return(zero_one_out)
        }
        
        # Missing first factor level
        check_out <- check_missing_first_factor_level(formula,
            random_effects_formula,
            dat_sub,
            groups,
            ordereds,
            fn = fn,
            feature_specific_covariate_name)
        
        if (!is.null(check_out)) {
            return(check_out)
        }
        
        # Augment logistic fitting
        if (augment &
            model == "logistic" &
            length(unique(featuresVector)) >= 2) {
            fitting_out <- fit_augmented_logistic(
                ranef_function,
                model_function,
                formula,
                random_effects_formula,
                groups,
                ordereds,
                dat_sub,
                # features,
                fn = fn,
                fi)
        } else { # linear or non-augmented logistic
            fitting_out <- non_augmented(
                ranef_function,
                model_function,
                formula,
                random_effects_formula,
                groups,
                ordereds,
                dat_sub,
                # features,
                fn = fn,
                fi)
        }
        
        fit_and_message <- fitting_out[["fit_and_message"]]
        weight_scheme <- fitting_out[["weight_scheme"]]
        mm_input <- fitting_out[["mm_input"]]
        
        fit <- fit_and_message[[1]]
        
        # Gather Output
        output <- list()
        
        # Check for fitting errors and add on special predictors
        low_n_error <- FALSE
        if (all(!inherits(fit, "try-error"))) {
            names_to_include <-
                get_fixed_effects(formula,
                                             random_effects_formula,
                                             dat_sub,
                                             character(0),
                                             character(0),
                                             feature_specific_covariate_name)
            # Suppress warnings about variance-covariance matrix calculation
            fit_properly <- FALSE
            
            withCallingHandlers({
                tryCatch({
                    output$para <-
                        summary_function(fit, c('(Intercept)', 
                                            names_to_include))
                    fit_properly <- TRUE
                }, error = function(e) {
                    return()
                })
            }, warning = function(w) {
                invokeRestart("muffleWarning")
            })
            
            if (fit_properly) {
                output$para <-
                    output$para[names_to_include, , drop = FALSE]
                
                n_uni_cols <- nrow(output$para)
                
                if (length(groups) > 0) {
                    output <- run_group_models(ranef_function,
                                            model_function,
                                            groups,
                                            formula,
                                            random_effects_formula,
                                            model,
                                            fit,
                                            augment,
                                            weight_scheme,
                                            dat_sub,
                                            output,
                                            mm_input)
                }
                
                if (length(ordereds) > 0) {
                    output <- run_ordered_models(ranef_function,
                                                model_function,
                                                ordereds,
                                                fit_and_message,
                                                formula,
                                                random_effects_formula,
                                                model,
                                                fit,
                                                augment,
                                                weight_scheme,
                                                dat_sub,
                                                output)
                }
                
                # Check whether summaries are correct
                names_to_include <-
                    get_fixed_effects(formula,
                                    random_effects_formula,
                                    dat_sub,
                                    groups,
                                    ordereds,
                                    feature_specific_covariate_name)
                if (any(!(names_to_include %in% rownames(output$para)))) {
                    # Don't worry about dropped factor levels
                    missing_names <- names_to_include[
                        !(names_to_include %in% rownames(output$para))]
                    
                    character_cols <- get_character_cols(dat_sub)
                    
                    if (!all(missing_names %in% character_cols)) {
                        fit_properly <- FALSE
                        
                        fit_and_message[[length(fit_and_message)]] <-
                            "Metadata dropped during fitting (rank deficient)"
                    } else {
                        fit_properly <- TRUE
                    }
                } else {
                    # No errors, summaries are correct
                    fit_properly <- TRUE
                }
                
                # Prevents individual group levels from ending up in results
                output$para <-
                    output$para[rownames(output$para) %in% names_to_include,]
            }
        } else {
            # Fit issue occurred
            fit_properly <- FALSE
        }
        
        output <- fitting_wrap_up(fit_properly,
                                fit_and_message,
                                output,
                                fit,
                                random_effects_formula,
                                metadata,
                                median_comparison,
                                save_models,
                                formula,
                                dat_sub,
                                groups,
                                ordereds,
                                # features,
                                fn = fn,
                                fi,
                                ranef_function,
                                feature_specific_covariate_name)
       
        
        return(output)
    }
    
    env_objects <- ls(environment(func_to_run))
    for (obj in env_objects) {
        size <- utils::object.size(get(obj, envir = environment(func_to_run)))
        logging::logdebug(paste0("Object: ", obj, ", Size: ", size))
    }
    size <- utils::object.size(func_to_run)
    logging::logdebug(paste0("Object: ", "func_to_run", ", Size: ", size))

    feat_list = as.list(as.data.frame(features))
    feat_nm = colnames(features)
    feat_i = seq_len(ncol(features))
    
    # purrr::list_transpose() would be a better way to do this V.
    map_input = data.frame(fv = I(feat_list), # feature value, name, index
                           fn = feat_nm, 
                           fi = feat_i)
    
    if (mirai::daemons_set()) {
        outputs <- mirai::mirai_map(map_input, func_to_run)[.progress]
        # ^ TODO, figure out how to ensure this writes to the log file
    } else {
        outputs <- mapply(func_to_run, 
                          map_input$fv,
                          map_input$fn,
                          map_input$fi,
                          SIMPLIFY = FALSE)
        # mirai_map() deliberately doesn't fall back to serial without daemons set: https://github.com/r-lib/mirai/issues/397
        
        # To get progress here:
        # lightest: error out and tell user to set daemons with mirai::require_daemons()
        # alternative: utils::txtProgressBar + a loop
        # alternative 2: reintroduce pbapply dependency (it's pretty light)
    }
   
    # Check to run once on a daemon: 
    # print(mirai::mirai(func_to_run(map_input$fv[[1]],
    #                   map_input$fn[1],
    #                   map_input$fi[1]),
    #                   func_to_run = func_to_run,
    #                   map_input = map_input)[])
    
    # cli::cli_alert("First result: ")
    # print(outputs[[1]])
    # outputs <-
    #     pbapply::pblapply(seq_len(ncol(features)), cl = cluster, func_to_run)
    
    # # stop the cluster
    # if (!is.null(cluster))
    #     parallel::stopCluster(cluster)
    
    # bind the results for each feature
    paras <- collapse::rowbind(lapply(outputs, function(x) {
        return(x$para)
    }))
   
    residuals <-
        do.call(rbind, lapply(outputs, function(x) {
            return(x$residuals)
        }))
   
    row.names(residuals) <- colnames(features)
    
    
    fitted <-
        do.call(rbind, lapply(outputs, function(x) {
            return(x$fitted)
        }))
    
    row.names(fitted) <- colnames(features)
    
    fits <-
        lapply(outputs, function(x) {
            return(x$fit)
        })
    names(fits) <- colnames(features)
    
    # Adjust the significance levels based on tests against the median
    if (median_comparison) {
        paras <- run_median_comparison(paras,
                                    fits,
                                    ordereds,
                                    groups,
                                    metadata,
                                    random_effects_formula,
                                    median_comparison_threshold,
                                    subtract_median,
                                    model)
    } else {
        paras$null_hypothesis <- 0
    }
    
    # Return NULL rather than empty object if fits aren't saved
    if (all(is.na(fits)) | !save_models) {
        fits <- NULL
    }
    
    if (!(is.null(random_effects_formula))) {
        
        l <- lapply(outputs, function(x) {
            return(x$ranef)
        })
        
        ranef <- do.call(rbind, l) |> 
            collapse::setRownames(colnames(features))
            
    }
    
    #####################################################
    # Determine the metadata names from the model names #
    #####################################################
    
    metadata_names <- colnames(metadata)
    if (!is.null(feature_specific_covariate)) {
        metadata_names <-
            union(metadata_names, feature_specific_covariate_name)
    }
    # order the metadata names by decreasing length
    metadata_names_ordered <-
        metadata_names[order(nchar(metadata_names), decreasing = TRUE)]
    # find the metadata name based on the match
    # to the beginning of the string
    extract_metadata_name <- function(name) {
        tmp_val <- metadata_names_ordered[mapply(startsWith,
                                                name,
                                                metadata_names_ordered)][1]
        if (is.na(tmp_val)) {
            metadata_names_ordered[mapply(grepl,
                                        metadata_names_ordered,
                                        name)][1]
        } else {
            return(tmp_val)
        }
    }
    paras$metadata <-
        unlist(lapply(paras$name, extract_metadata_name))
    # compute the value as the model contrast minus metadata
    paras$value <-
        mapply(function(x, y) {
            if (x == y)
                x
            else
                gsub("^\\:", "", sub(x, "", y))
        }, paras$metadata, paras$name)
    # Trim ticks if metadata had ticks
    paras$value <- gsub('^``', '', paras$value)
    
    if (!is.null(feature_specific_covariate_record)) {
        if (!feature_specific_covariate_record) {
            paras <-
                paras[!grepl(feature_specific_covariate_name, paras$name),]
        }
    }
    
    if (small_random_effects & model == 'logistic') {
        paras <- paras[!paras$metadata %in% trimws(grouping, whitespace = '`'),]
    }
    
    ##############################
    # Sort by decreasing q-value #
    ##############################
    
    paras <- paras[order(paras$pval, decreasing = FALSE), ]
    
    paras <- collapse::colorderv(paras, c('feature', 'metadata', 'value'))
    paras$model <- model
    rownames(paras) <- NULL
    
    if (!(is.null(random_effects_formula))) {
        return(
            list(
                "results" = paras,
                "residuals" = residuals,
                "fitted" = fitted,
                "ranef" = ranef,
                "fits" = fits
            )
        )
    } else {
        return(
            list(
                "results" = paras,
                "residuals" = residuals,
                "fitted" = fitted,
                "ranef" = NULL,
                "fits" = fits
            )
        )
    }
}
