ols_w_indicator_gg

compute_panel_ols_ind <- function(data, scales) {

  data$indicator = factor(indicator)
  
  model <- lm(y ~ x + indicator,
              data = data)

  data.frame(x = data$x,
             y = model$fitted.values,
             indicator = data$indicator)

}

#### confint ####################



# StatLmindicator <- ggplot2::ggproto(`_class` = "StatLmindicator",
#                                    `_inherit` = ggplot2::Stat,
#                                    # setup_data = my_setup_data,
#                                    required_aes = c("x", "y", "indicator"),
#                                    compute_panel = compute_panel_ols_ind#,
#                                    # default_aes = aes(fill = after_stat(area))
# )




# For most applications the grouping is set implicitly by mapping one
# or more discrete variables to x, y, colour, fill, alpha, shape, size, and/or linetype

StatLmindicator <- ggplot2::ggproto("StatLmindicator",
                                   ggplot2::Stat,
                                   compute_panel = compute_panel_ols_ind,
                                   required_aes = c("x", "y", "indicator"),
                                   default_aes = aes(group = after_stat(indicator))
                                   )

geom_lm_indicator <- function(mapping = NULL, data = NULL,
                                    position = "identity", na.rm = FALSE,
                                    show.legend = NA,
                                    inherit.aes = TRUE, ...) {
  ggplot2::layer(
    stat = StatLmindicator, # proto object from Step 2
    geom = ggplot2::GeomLine, # inherit other behavior
    data = data,
    mapping = mapping,
    position = position,
    show.legend = show.legend,
    inherit.aes = inherit.aes,
    params = list(na.rm = na.rm, ...)
  )
}


geom_lm_indicator_fitted <- function(mapping = NULL, data = NULL,
                                     position = "identity", na.rm = FALSE,
                                     show.legend = NA,
                                     inherit.aes = TRUE, ...) {
  ggplot2::layer(
    stat = StatLmindicator, # proto object from Step 2
    geom = ggplot2::GeomPoint, # inherit other behavior
    data = data,
    mapping = mapping,
    position = position,
    show.legend = show.legend,
    inherit.aes = inherit.aes,
    params = list(na.rm = na.rm, ...)
  )
}




######### label ##
# cars %>% 
#   rename(x = speed) %>% 
#   rename(y = dist) %>% 
#   mutate(indicator = x >15) %>% 
#   compute_panel_ols_ind_label()

compute_panel_ols_ind_label <- function(data, scales) {
  
  model <- lm(y ~ x + indicator,
              data = data)
  
  data.frame(names = model[[1]] %>% names(),
             coeff = model[[1]]) %>% 
    tibble() %>% 
    slice(-1, -2) %>% 
    mutate(equation = paste0(coeff %>% good_digits(), "*", names)) %>% 
    pull(equation) %>% 
    paste(collapse = " + ") ->
    dummies
  
  
  data.frame(x = mean(data$x),
             y = mean(data$y),
             label = paste0("y = ",
                           model$coefficients[2] %>% good_digits(), 
                           "x + ",
                           dummies,
                           " + ",
                           model$coefficients[1] %>% good_digits() 
                           
             ))
  
}



StatOlsindformula <- ggplot2::ggproto("StatOlsindformula",
                                   ggplot2::Stat,
                                   compute_panel = compute_panel_ols_ind_label,
                                   required_aes = c("x", "y", "indicator")
                                )

#' Write formula for ols linear model
#'
#' @param mapping
#' @param data
#' @param position
#' @param na.rm
#' @param show.legend
#' @param inherit.aes
#' @param ...
#'
#' @return
#' @export
#'
#' @examples
#' library(ggplot2)
#' ggplot(cars) + aes(x = speed, y = dist, indicator = dist > 40) +
#' geom_point() + geom_lm_indicator() + geom_lm_ind_formula()
geom_lm_indicator_formula <- function(mapping = NULL, data = NULL,
                            position = "identity", na.rm = FALSE, show.legend = NA,
                            inherit.aes = TRUE, ...) {
  ggplot2::layer(
    stat = StatOlsindformula, geom = ggplot2::GeomLabel, data = data, mapping = mapping,
    position = position, show.legend = show.legend, inherit.aes = inherit.aes,
    params = list(na.rm = na.rm, ...)
  )
}




# ggplot(cars) +
#   aes(x = speed, y = dist) +
#   aes(indicator = dist > 40) +
#   aes(color = dist > 40) +
#   geom_point() +
#   geom_lm_indicator(color = "blue") +
#   geom_lm_indicator_fitted(color = "blue") +
#   geom_lm_indicator_formula() +
#   NULL


# ggplot(palmerpenguins::penguins) + 
#   aes(x = bill_length_mm ) + 
#   aes(y = bill_depth_mm ) + 
#   geom_point() + aes(color = species) + 
#   aes(indicator = species) +
#   geom_lm_indicator() + 
#   geom_lm_ind_formula() + 
#   geom_lm_indicator_fitted() +
#   geom_lm_formula()