functions to regress out random effects

#' @export 

regress_out_one_gene <- function(formula_full, formula_reduced, design, y, common_nUMI) {

    tryCatch({

        glmer_res <- lme4::glmer(

            formula_full, cbind(design, y), "poisson", 

            control = lme4::glmerControl(calc.derivs = FALSE, optimizer = "nloptwrap")

        )   

        ## simulate all samples to the same read depth 

        design$logUMI <- log(common_nUMI)

        ypred <- lme4:::predict.merMod(

            object = glmer_res, 

            newdata = design, 

            type = 'response',

            re.form = formula_reduced

        ) 

        return(ypred)        

    }, error = function(e) {

        return(rep())

    })

}

#' @export 

regress_out.presto <- function(obj, formula_full, formula_reduced, features = NULL, do_par = TRUE, common_nUMI=1e6) {

    ## TODO: check that formula_reduced only has RHS 

    if (do_par) {

        plan(multicore)

        it_fxn <- furrr::future_map

    }

    else {

        it_fxn <- purrr::map

    }

    if (is.null(features)) {

        features <- rownames(counts_mat)    

    }

    lres <- it_fxn(features, function(feature_use) {

        regress_out_one_gene(

            formula_full = formula_full, 

            formula_reduced = formula_reduced,

            design = obj$meta_data, 

            y = obj$counts_mat[feature_use, ], 

            common_nUMI

        )

    })

    obj$corr_mat <- Reduce(Matrix::rbind2, lres)

    rownames(obj$corr_mat) <- features

    colnames(obj$corr_mat) <- colnames(obj$counts_mat)

    return(obj)

}

#' @export 

find_markers_glmm_single_gene <- function(dge_formula, design, y, main_effect, nsim) {

    ## Estimate model 

    ## TODO: check that feature and group are the same 

    ## TODO: instead of feature and group, make generic 

    ## TODO: check that key_names uniquely define rows

    message('CAUTION: if gene names have underscores, pasrsing will be wrong')

    res <- switch(

        type, 

        'fixedeffects' = { 

            ## prepare data

            tests <- models[[1]] %>% 

                tidyr::unite(key, all_of(key_names), sep = '_') %>% 

                with(key)

            beta_mat <- bind_rows(models, .id = 'temp_group_var')[, c('beta', 'temp_group_var', key_names)] %>% 

                tidyr::spread(temp_group_var, beta, fill=0) %>% 

                tidyr::unite(key, all_of(key_names), sep = "_") %>% 

                tibble::column_to_rownames('key') %>% as.matrix()

            sd_mat <- bind_rows(models, .id = 'temp_group_var')[, c('sigma', 'temp_group_var', key_names)] %>% 

                tidyr::spread(temp_group_var, sigma, fill=0) %>% 

                tidyr::unite(key, all_of(key_names), sep = "_") %>% 

                tibble::column_to_rownames('key') %>% as.matrix()

            tests <- intersect(rownames(beta_mat), rownames(sd_mat))

            beta_mat <- beta_mat[tests, ]

            sd_mat <- sd_mat[tests, ]

#             tests <- models[[1]] %>% 

#                 tidyr::unite(key, all_of(key_names), sep = '_') %>% 

#                 with(key)

            beta_mat <- models %>% map('beta') %>% dplyr::bind_cols(testid = tests) %>% 

                tibble::column_to_rownames('testid') %>% 

                as.matrix()

#             beta_mat <- models %>% map('beta') %>% dplyr::bind_cols(testid = tests) %>% 

#                 tibble::column_to_rownames('testid') %>% 

#                 as.matrix()

            sd_mat <- models %>% map('sigma') %>% dplyr::bind_cols(testid = tests) %>% 

                tibble::column_to_rownames('testid') %>% 

                as.matrix()

#             sd_mat <- models %>% map('sigma') %>% dplyr::bind_cols(testid = tests) %>% 

#                 tibble::column_to_rownames('testid') %>% 

#                 as.matrix()

            ## summary stats

            compute_fixed_effects(beta_mat, sd_mat) %>% 

                tidyr::separate(test, key_names)

                tidyr::separate(test, key_names, sep='_')

        },

        NULL ## default