# Predict ---------------------------------------

#' @export
tidypredict_fit.lm <- function(model) {
  parsedmodel <- parse_model(model)
  build_fit_formula(parsedmodel)
}

build_fit_formula <- function(parsedmodel) {
  parsed_f <- map(
    parsedmodel$terms,
    ~ {
      if (.x$is_intercept == 0) {
        cols <- map(.x$fields, lm_constructor)
        cols <- reduce_multiplication(cols)
        expr((!!cols * !!.x$coef))
      } else {
        if (.x$coef == 0) {
          NULL
        } else {
          expr(!!.x$coef)
        }
      }
    }
  )
  parsed_f <- purrr::discard(parsed_f, is.null)

  f <- reduce_addition(parsed_f)

  if (!is.null(parsedmodel$general$offset)) {
    f <- expr_addition(f, parsedmodel$general$offset)
  }

  if (parsedmodel$general$is_glm == 1) {
    link <- parsedmodel$general$link
    f <- apply_inverse_link(f, link)
  }
  f
}

apply_inverse_link <- function(f, link) {
  switch(
    link,
    "identity" = f,
    "logit" = expr(1 - 1 / (1 + exp(!!f))),
    "log" = expr(exp(!!f)),
    "inverse" = expr(1 / (!!f)),
    "1/mu^2" = expr(1 / sqrt(!!f)),
    # Bowling et al. approximation to normal CDF (max error ~0.014%)
    "probit" = expr(1 / (1 + exp(-0.07056 * (!!f)^3 - 1.5976 * (!!f)))),
    "cloglog" = expr(1 - exp(-exp(!!f))),
    "sqrt" = expr((!!f)^2),
    cli::cli_abort("Link {.val {link}} is not supported.")
  )
}

# Parse model --------------------------------------

#' @export
parse_model.lm <- function(model) parse_model_lm(model)

parse_model_lm <- function(model, call = rlang::caller_env()) {
  acceptable_formula(model)

  coefs <- as.numeric(model$coefficients)
  labels <- names(model$coefficients)
  vars <- names(attr(model$terms, "dataClasses"))
  qr <- NULL
  if (!is.null(model$qr)) {
    qr <- tryCatch(
      qr.solve(qr.R(model$qr)),
      error = function(cnd) {
        if (grepl("singular matrix", cnd$message)) {
          cli::cli_abort(
            c(
              x = "Unable to calculate inverse of QR decomposition.",
              i = "This is likely happening because the predictors contain a 
              linear combination of predictors. Please remove and try again."
            ),
            call = call
          )
        }
        stop(cnd)
      }
    )
  }

  pm <- list()
  pm$general$model <- class(model)[[1]]
  pm$general$version <- 2
  pm$general$type <- "regression"
  pm$general$residual <- model$df.residual

  if (length(summary(model)$sigma^2) > 0) {
    pm$general$sigma2 <- summary(model)$sigma^2
  }
  if (!is.null(model$family$family)) {
    pm$general$family <- model$family$family
  }
  if (!is.null(model$family$link)) {
    pm$general$link <- model$family$link
  }
  if (!is.null(model$call$offset)) {
    pm$general$offset <- model$call$offset
  }
  pm$general$is_glm <- 0
  if (class(model)[[1]] == "glm") {
    pm$general$is_glm <- 1
  }
  terms <- map(
    seq_len(length(labels)),
    ~ {
      list(
        label = labels[.x],
        coef = coefs[.x],
        is_intercept = ifelse(labels[.x] == "(Intercept)", 1, 0),
        fields = parse_label_lm(labels[.x], vars),
        qr = parse_qr_lm(labels[.x], qr)
      )
    }
  )
  pm$terms <- terms
  as_parsed_model(pm)
}

parse_label_lm <- function(label, vars) {
  all_items <- NULL
  items <- strsplit(label, "\\:")[[1]]
  for (i in seq_len(length(items))) {
    item <- list(
      type = "ordinary",
      col = items[i]
    )
    cat_match <- map_lgl(vars, ~ .x == substr(items[i], 1, nchar(.x)))
    if (
      any(cat_match) &&
        any(vars[cat_match] != items[i]) &&
        !(items[i] %in% vars)
    ) {
      cat_match_vars <- vars[cat_match]
      sole_cat_match <- cat_match_vars[rank(-nchar(cat_match_vars))][[1]]
      item <- list(
        type = "conditional",
        col = sole_cat_match,
        val = substr(items[i], nchar(sole_cat_match) + 1, nchar(items[i])),
        op = "equal"
      )
    }
    all_items <- c(all_items, list(item))
  }
  all_items
}

parse_qr_lm <- function(label, qr) {
  qrs <- qr[label == rownames(qr)]
  qrs <- set_names(
    as.list(qrs),
    paste0("qr_", 1:length(qrs))
  )
}

# Intervals -----------------------------------------------

#' @export
tidypredict_interval.lm <- function(model, interval = 0.95) {
  parsedmodel <- parse_model(model)
  te_interval_lm(parsedmodel, interval)
}

get_qr_lm <- function(qr_name, parsedmodel) {
  q <- map(
    parsedmodel$terms,
    ~ {
      cqr <- .x$qr[qr_name][[1]]

      if (.x$is_intercept == 0) {
        cols <- map(.x$fields, lm_constructor)
        cols <- reduce_multiplication(cols)
        if (cqr != 0) {
          expr_multiplication(cols, cqr)
        }
      } else {
        expr(!!cqr)
      }
    }
  )
  f <- reduce_addition(q[!map_lgl(q, is.null)])

  expr(((!!f)) * ((!!f)) * !!parsedmodel$general$sigma2)
}

te_interval_lm <- function(parsedmodel, interval = 0.95) {
  qr_names <- names(parsedmodel$terms[[1]]$qr)
  qrs_map <- map(
    qr_names,
    ~ get_qr_lm(.x, parsedmodel)
  )
  qrs <- reduce_addition(qrs_map)
  tfrac <- qt(1 - (1 - 0.95) / 2, parsedmodel$general$residual)
  expr(!!tfrac * sqrt((!!qrs) + (!!parsedmodel$general$sigma2)))
}

# Helpers -------------------------------------------------

lm_constructor <- function(x) {
  f <- NULL
  if (x$type == "ordinary") {
    f <- expr(!!as.name(x$col))
  }
  if (x$type == "conditional") {
    f <- expr(ifelse(!!as.name(x$col) == !!x$val, 1, 0))
  }
  if (x$type == "operation") {
    if (x$op == "morethan") {
      f <- expr(ifelse(
        !!as.name(x$col) > !!x$val,
        !!as.name(x$col) - !!x$val,
        0
      ))
    }
    if (x$op == "lessthan") {
      f <- expr(ifelse(
        !!as.name(x$col) < !!x$val,
        !!x$val - !!as.name(x$col),
        0
      ))
    }
  }
  f
}
