major_missing_because_of_controls <- 
  filter(df, !complete.cases(select(df, all_of(controls)))) %>% 
  filter(pop_tot>1 & !is.na(country) & country != "NA")

dropped <- major_missing_because_of_controls$country %>% 
  as.character %>% 
  unique 

major_missing_because_of_controls_pc <- 
  filter(df, !complete.cases(select(df, all_of(controls_deaths_capita)))) %>%
  filter(pop_tot>1 & !is.na(country) & country != "NA")

dropped_pc <- major_missing_because_of_controls_pc$country %>% 
  as.character %>% 
  unique 

results <- function(Y, 
                    X = vars, 
                    Xlab = NA,
                    XX = controls,
                    range = NULL, 
                    data = filter(df, geoid2 != "CN"), 
                    bivariate = FALSE, 
                    relative_dates = FALSE,
                    standardize = FALSE, 
                    drop_missing_c = TRUE,
                    model = "lm",
                    dayfreq = 5){

  if(drop_missing_c) data <- 
        filter(data, complete.cases(select(data, all_of(XX))))

  if(!relative_dates)  
    data$points <- data$elapsed
  
  if(relative_dates) {
    data$points <- data$elapsed_rel
    data <- filter(data, elapsed_rel < max(data$elapsed, na.rm = TRUE) - 80)
    }

  if(is.null(range)) 
    range <- c(ifelse(relative_dates, 0, 60), max(data$points, na.rm = TRUE))
  
  ## If standardize
  if (standardize)
    data <- data %>% 
      mutate_at(vars(one_of(X)),
                list(~./sd(., na.rm = T)))

  out <-

    lapply(seq(range[1], range[2], dayfreq), function(j) {

      lapply(X, function(x) {
      
        ## Formula
        rhs <- ifelse(bivariate, x, paste(unique(c(x, XX)), collapse = " + "))

        if (model == "lm") {
      
          my_formula <- paste0(Y, " ~ ",  rhs)
          
          m <- lm_robust(as.formula(my_formula), 
                         data = filter(data, points == j))
          rsq <- m$r.squared
          n <- m$N
          
          ## partial RSQ
          
          if (!bivariate) {
            f_reduced <- paste0(Y, " ~ ",  paste0(setdiff(XX, x), collapse = '+')) %>% 
              as.formula()
            f_full <- my_formula %>% as.formula()
            
            m_reduced <- lm(f_reduced, data = filter(data[data$points == j & 
                                                            complete.cases(data[, unique(c(XX, x))]), ]))
            m_full <- lm(f_full, data = filter(data[data$points == j & 
                                                      complete.cases(data[, unique(c(XX, x))]), ]))
            
            ## Get partial rsq
            
            rsq <- (sum(m_reduced$residuals^2) - sum(m_full$residuals^2)) / sum(m_reduced$residuals^2)
            
            ## Adjust: Formula is 1-(1-R2)*(DF_Reduced / DF_Full)
            
            rsq <- 1-((1-rsq)*(summary(m_reduced)$df[2] / summary(m_full)$df[2]))
          }
          
          ## Return
          
          m %>% tidy() %>% slice(2) %>% 
            mutate(n = n, rsq = rsq)
          
        } else {
          
          ## Negative binomial model
          my_formula <- paste0(Y, " ~ ",  rhs, '|', rhs)
  
          model <- zeroinfl(formula = as.formula(my_formula), data = data, dist = "negbin")
          
          out <- 
            rbind(
            summary(model)$coefficients$count[2,],
            summary(model)$coefficients$zero[2,]) %>% 
            data.frame() 
          
          names(out) <- c("estimate", "std.error", "statistic", "p.value")
          
          out %>%
            mutate(
              term = paste0(X, c("_count", "_zero")),
              conf.high  = estimate + 1.96*std.error,
              conf.low  = estimate - 1.96*std.error,
              outcome = "Y",
              n = model$n
              )
        }