# Find aes ----------------------------------------------------------------

#' @keywords internal
.find_aes <- function(x, model_info = NULL, numeric_as_discrete = 8) {
  # init basic aes
  data <- as.data.frame(x, preserve_names = TRUE)
  data$.group <- 1

  att <- attributes(x)
  aes <- list(y = "Predicted", group = ".group")

  # extract information for labels
  model_data <- .safe(insight::get_data(attributes(x)$model, verbose = FALSE))
  model_response <- attributes(x)$response

  # Find predictors
  by <- att$focal_terms

  # flag for ordinal and alike models
  has_response_levels <- isTRUE(
    model_info$is_ordinal |
      model_info$is_multinomial |
      model_info$is_categorical |
      model_info$is_cumulative
  )

  # multivariate response models? if so, we need one more stratification in "by"
  if (has_response_levels && "Response" %in% colnames(data)) {
    by <- c(by, "Response")
    data$Response <- factor(data$Response, levels = unique(data$Response))
  }

  # mixture models? if so, we need one more stratification in "by"
  if (isTRUE(model_info$is_mixture) && "Class" %in% colnames(data)) {
    by <- c(by, "Class")
    data$Class <- factor(data$Class, levels = unique(data$Class))
  }

  # if we have only few numeric values, we don't want a continuous color scale.
  # check whether we can treat numeric as discrete
  if (!isFALSE(numeric_as_discrete) && is.numeric(numeric_as_discrete)) {
    data[by] <- lapply(data[by], function(v) {
      if (is.numeric(v) && insight::n_unique(v) < numeric_as_discrete) {
        formatted <- insight::format_value(v, protect_integers = TRUE)
        v <- factor(formatted, levels = unique(formatted))
      }
      v
    })
  }

  # Main geom
  # ------------------------------------------------------------------------
  if ("estimate_contrasts" %in% att$class) {
    insight::format_error("Automated plotting is not yet implemented for this class.")
  } else if ("estimate_means" %in% att$class) {
    aes$y <- att$coef_name
  } else if ("estimate_slopes" %in% att$class) {
    # for frequentist models, we have "Slope" as column name, for Bayesian models
    # we have "Median", "Mean" or "MAP"
    valid_y_vars <- intersect(c("Slope", "Median", "Mean", "MAP"), colnames(data))
    # we wouldn't expect that there is more than one of these columns, but we
    # check for that anyway...
    if (length(valid_y_vars) == 0) {
      insight::format_error(
        "Could not find a suitable column for the y-axis. Expected one of: 'Slope', 'Median', 'Mean', 'MAP'."
      )
    }
    aes$y <- valid_y_vars[1]
    if ("Comparison" %in% names(data)) {
      # Insert "Comparison" column as the 2nd by so that it gets plotted as color
      if (length(by) > 1) {
        by[3:(length(by) + 1)] <- by[2:length(by)]
      }
      by[2] <- "Comparison"
    } else if ("p" %in% colnames(data) && length(by) == 1 && is.numeric(data[[by]])) {
      # this is for slopes of two numeric interaction terms (johnson-neymann plots)
      # the slope of the first numeric term is along the y-axis, values for the
      # 2nd numeric term on the x-axis. whenever confidence intervals exclude 0
      # we find the (range of) numeric values for the 2nd term where the interaction
      # is "significant", i.e. p < 0.05. We want to map this to a special aes
      # so we can color the ribbons accordingly
      by <- c(by, "p")
      significant <- data$p < 0.05
      data$p <- "not significant"
      data$p[significant] <- "significant"
    }
  } else if ("estimate_grouplevel" %in% att$class) {
    aes$x <- "Level"
    # find coefficient name, this may differ for Bayesian models
    if (!is.null(att$coef_name) && length(att$coef_name)) {
      aes$y <- att$coef_name
    } else {
      # Default y-axis for frequentist models
      aes$y <- "Coefficient"
    }
    # Set a custom plot type, used later to apply specific geoms (pointrange, coord_flip)
    aes$type <- "grouplevel"

    # --- Setup for faceting (creating multiple plot panels) ---
    facet_by <- NULL
    # Check if there are multiple parameters (e.g., random intercept and random slope)
    if (insight::n_unique(data$Parameter) > 1) {
      facet_by <- c(facet_by, "Parameter")
    }
    # Check if there are multiple grouping factors (e.g., (1|Subject) + (1|School))
    if (insight::n_unique(data$Group) > 1) {
      facet_by <- c(facet_by, "Group")
    }
    # Check for multiple components (for complex distributional models)
    if (insight::n_unique(data$Component) > 1) {
      facet_by <- c(facet_by, "Component")
    }

    # If any faceting is needed, create a new 'facet' column in the data
    if (!is.null(facet_by)) {
      # Start with the group name (e.g., "Subject")
      data$facet <- data$Group
      # Append the parameter name if needed (e.g., "Subject: Days")
      if ("Parameter" %in% facet_by) {
        data$facet <- paste0(data$facet, ": ", data$Parameter)
      }
      # Append the component name if needed (e.g., "Subject: Days (sigma)")
      if ("Component" %in% facet_by) {
        data$facet <- paste0(
          data$facet,
          " (",
          gsub("_", " ", data$Component, fixed = TRUE),
          ")"
        )
      }
      # Tell the main plotting function to use this new column for faceting
      aes$facet <- "facet"
    }
    # Find the columns for confidence intervals (ymin, ymax)
    aes <- .find_aes_ci(aes, data)
    # Return the aesthetics and data, and exit the function
    return(list(aes = aes, data = data))
  }

  # Assign predictors to aes
  if (is.null(by)) {
    by <- att$by
  }
  if (length(by) == 0) {
    insight::format_error(
      "No `by` variable was detected, so nothing to put in the x-axis."
    )
  }

  # first variable for x-axis - decide whether we have a dot- or a line-plot
  # ------------------------------------------------------------------------
  aes$x <- by[1]
  # If x is a not-numeric, make pointrange
  if (is.numeric(data[[by[1]]])) {
    aes$type <- "ribbon"
  } else {
    aes$type <- "pointrange"
  }

  # second variable, mapped to the color-aes
  # ------------------------------------------------------------------------
  if (length(by) > 1) {
    aes$color <- by[2]
    # if by is of length 2, and the p-value column, *and* we have slopes, then
    # # we have a Johnson-Neyman plot. In this case, we need to re-adjust ".group",
    # which must have an own index for each "part" of the ribbons. ".group" now
    # indicates every switch / flip from significant to non-significant and
    # vice versa
    if ("estimate_slopes" %in% att$class && by[2] == "p") {
      group_index <- 1
      for (i in 2:(nrow(data))) {
        if (data$p[i] != data$p[i - 1]) {
          group_index <- group_index + 1
        }
        data$.group[i] <- group_index
      }
      # for johnson-neymann plots, the "group" aes cannot be assigned to the
      # ribbon geom - instead, it must be part of the "ggplot()" function. this
      # can be achieved by adding the aes as "global_aes" attribute to the
      # returned visualisation_recipe
      aes$group <- NULL
    } else {
      data$.group <- paste(data$.group, data[[by[2]]])
    }
  }

  # third variable, mapped to alpha when numeric, or creates a facet for factors
  # ------------------------------------------------------------------------
  if (length(by) > 2) {
    if (is.numeric(data[[by[3]]])) {
      aes$alpha <- by[3]
    } else {
      aes$facet <- stats::as.formula(paste(
        "~",
        paste(utils::tail(by, -2), collapse = " * ")
      ))
    }
    data$.group <- paste(data$.group, data[[by[3]]])
  }

  # more than three variable? create facet grids then
  # ------------------------------------------------------------------------
  if (length(by) > 3) {
    aes$facet <- NULL
    # we have to switch variables 3 and 4, due to regression formula
    remaining <- paste(utils::tail(by, -2))
    aes$grid <- stats::as.formula(paste(
      remaining[2],
      "~",
      paste(setdiff(remaining, remaining[2]), collapse = "*")
    ))
  }

  # CI
  # ------------------------------------------------------------------------
  aes <- .find_aes_ci(aes, data)

  # axis and legend labels
  # ------------------------------------------------------------------------
  if (!is.null(model_data) && !is.null(model_response)) {
    if ("estimate_slopes" %in% att$class) {
      ylab <- att$trend
    } else {
      # response - mapped to the y-axis if not slopes
      ylab <- .safe(attr(model_data[[model_response]], "label", exact = TRUE))
    }
    # fix default y-label, if necessary
    y_prefix <- aes$y
    if (y_prefix == "Predicted") {
      y_prefix <- "Predicted value"
    }
    # set y-label based on labelled data, or variable name
    if (is.null(ylab)) {
      ylab <- paste(y_prefix, "of", model_response)
    } else {
      ylab <- paste(y_prefix, "of", ylab)
    }
    # main predictor - mapped to x-axis
    xlab <- .safe(attr(model_data[[by[1]]], "label", exact = TRUE))
    # first grouping variable (2nd in "by") - mapped to legend
    if (length(by) > 1) {
      colour <- .safe(attr(model_data[[by[[2]]]], "label", exact = TRUE))
    } else {
      colour <- NULL
    }
    aes$labs <- insight::compact_list(list(y = ylab, x = xlab, colour = colour))
  }

  list(aes = aes, data = data)
}


#' @keywords internal
.find_aes_ci <- function(aes, data) {
  ci_lows <- rev(grep("CI_low", names(data), fixed = TRUE, value = TRUE))
  ci_highs <- rev(grep("CI_high", names(data), fixed = TRUE, value = TRUE))
  if (length(ci_lows) > 0) {
    aes$ymin <- ci_lows
    aes$ymax <- ci_highs
  } else {
    aes$ymin <- NA_real_
    aes$ymax <- NA_real_
  }
  aes
}


# Workhorse function ------------------------------------------------------

#' @keywords internal
.visualization_recipe <- function(
  x,
  show_data = TRUE,
  show_residuals = FALSE,
  collapse_by = NULL,
  point = NULL,
  line = NULL,
  pointrange = NULL,
  ribbon = NULL,
  facet = NULL,
  grid = NULL,
  join_dots = FALSE,
  numeric_as_discrete = 8,
  ...
) {
  # init
  response_scale <- attributes(x)$predict
  model_info <- attributes(x)$model_info

  aes <- .find_aes(x, model_info, numeric_as_discrete)
  data <- aes$data
  aes <- aes$aes
  global_aes <- list()
  layers <- list()
  l <- 1

  # preparation of settings / arguments ----------------------------------

  # check whether point-geoms should be connected by lines
  do_not_join <- "grouplevel"
  if (!join_dots) {
    do_not_join <- c(do_not_join, "pointrange", "point")
  }

  # Don't plot raw data if `predict` is not on the response scale
  if (
    !is.null(response_scale) &&
      !response_scale %in% c("prediction", "response", "expectation", "invlink(link)")
  ) {
    show_data <- show_residuals <- FALSE
  }

  # Don't plot raw data for transformed responses with no back-transformation
  transform <- attributes(x)$transform

  if (isTRUE(model_info$is_linear) && !isTRUE(transform)) {
    # add information about response transformation
    trans_fun <- .safe(insight::find_transformation(attributes(x)$model))
    if (!is.null(trans_fun) && all(trans_fun != "identity")) {
      show_data <- show_residuals <- FALSE
    }
  }

  # add raw data as first layer ----------------------------------
  if (show_data && is.null(collapse_by)) {
    layers[[paste0("l", l)]] <- .visualization_recipe_rawdata(x, aes, numeric_as_discrete)
    # Update with additional args
    if (!is.null(point)) {
      layers[[paste0("l", l)]] <- utils::modifyList(layers[[paste0("l", l)]], point)
    }
    l <- l + 1
  }

  # add residual data as next lowest layer
  if (show_residuals && is.null(collapse_by)) {
    layers[[paste0("l", l)]] <- .visualization_recipe_residuals(
      x,
      aes,
      numeric_as_discrete
    )
    # Update with additional args
    if (!is.null(point)) {
      layers[[paste0("l", l)]] <- utils::modifyList(layers[[paste0("l", l)]], point)
    }
    l <- l + 1
  }

  # next possible data layer is collapsing by group ---------------
  if (!is.null(collapse_by) && (show_data || show_residuals)) {
    layers[[paste0("l", l)]] <- .visualization_recipe_collapsed_data(
      x,
      aes,
      collapse_by,
      numeric_as_discrete,
      show_residuals
    )
    # Update with additional args
    if (!is.null(point)) {
      layers[[paste0("l", l)]] <- utils::modifyList(layers[[paste0("l", l)]], point)
    }
    l <- l + 1
  }

  # intercept line for slopes ----------------------------------
  if (inherits(x, "estimate_slopes")) {
    layers[[paste0("l", l)]] <- insight::compact_list(list(
      geom = "hline",
      yintercept = 0,
      alpha = getOption("modelbased_slopes_alpha", 1 / 2),
      linetype = "dashed"
    ))
    l <- l + 1
    global_aes$group <- ".group"
  }

  # Uncertainty -----------------------------------
  if (!identical(ribbon, "none") && aes$type == "ribbon" && is.null(aes$alpha)) {
    for (i in seq_along(aes$ymin)) {
      # base list elements
      aes_list <- list(
        y = aes$y,
        x = aes$x,
        ymin = aes$ymin[i],
        ymax = aes$ymax[i],
        fill = aes$color
      )
      # optionally add group aes, if not in global_aes and not null
      if (!"group" %in% names(global_aes) || !is.null(aes$group)) {
        aes_list$group <- aes$group
      }
      layers[[paste0("l", l)]] <- list(
        geom = "ribbon",
        data = data,
        aes = aes_list,
        alpha = getOption("modelbased_ribbon_alpha", 1 / 4)
      )
      if (!is.null(ribbon)) {
        layers[[paste0("l", l)]] <- utils::modifyList(layers[[paste0("l", l)]], ribbon)
      }
      l <- l + 1
    }
  }

  # Main ----------------------------------

  # connecting lines between point geoms
  if (!aes$type %in% do_not_join) {
    # base list elements
    aes_list <- list(y = aes$y, x = aes$x, color = aes$color, alpha = aes$alpha)
    # optionally add group aes, if not in global_aes and not null
    if (!"group" %in% names(global_aes) || !is.null(aes$group)) {
      aes_list$group <- aes$group
    }
    layers[[paste0("l", l)]] <- list(geom = "line", data = data, aes = aes_list)
    if (!is.null(aes$color) && aes$type %in% c("pointrange", "point")) {
      layers[[paste0("l", l)]]$position <- "dodge"
      layers[[paste0("l", l)]]$width <- 0.2
    }
    if (!is.null(line)) {
      layers[[paste0("l", l)]] <- utils::modifyList(layers[[paste0("l", l)]], line)
    }
    l <- l + 1
  }

  # points with error bars
  if (aes$type %in% c("pointrange", "grouplevel")) {
    layers[[paste0("l", l)]] <- list(
      geom = "pointrange",
      data = data,
      aes = list(
        y = aes$y,
        x = aes$x,
        ymin = aes$ymin,
        ymax = aes$ymax,
        color = aes$color,
        group = aes$group,
        alpha = aes$alpha
      )
    )
    if (!is.null(aes$color)) {
      layers[[paste0("l", l)]]$position <- "dodge"
      layers[[paste0("l", l)]]$width <- 0.2
    }
    if (!is.null(pointrange)) {
      layers[[paste0("l", l)]] <- utils::modifyList(layers[[paste0("l", l)]], pointrange)
    }
    l <- l + 1
  }
  if (aes$type == "grouplevel") {
    layers[[paste0("l", l)]] <- list(geom = "coord_flip")
    l <- l + 1
  }

  # grids and facets ----------------------------------
  if (!is.null(aes$facet)) {
    layers[[paste0("l", l)]] <- list(geom = "facet_wrap", data = data, facets = aes$facet)
    if (!is.null(facet)) {
      layers[[paste0("l", l)]] <- utils::modifyList(layers[[paste0("l", l)]], facet)
    }
    l <- l + 1
  }
  if (!is.null(aes$grid)) {
    layers[[paste0("l", l)]] <- list(
      geom = "facet_grid",
      data = data,
      rows = aes$grid,
      scales = "free_x"
    )
    if (!is.null(grid)) {
      layers[[paste0("l", l)]] <- utils::modifyList(layers[[paste0("l", l)]], facet)
    }
    l <- l + 1
  }

  # add axis and legend labels ----------------------------------
  if (!is.null(aes$labs)) {
    layers[[paste0("l", l)]] <- insight::compact_list(list(
      geom = "labs",
      x = aes$labs$x,
      y = aes$labs$y,
      colour = aes$labs$colour,
      fill = aes$labs$colour
    ))
    l <- l + 1
  }

  # probability scale? ----------------------------------
  if (
    !is.null(response_scale) &&
      response_scale %in%
        c("response", "expectation", "invlink(link)", "prob", "probs") &&
      isTRUE(
        model_info$is_logit |
          model_info$is_binomial |
          model_info$is_orderedbeta |
          model_info$is_beta |
          model_info$is_ordinal |
          model_info$is_multinomial
      )
  ) {
    # nolint
    layers[[paste0("l", l)]] <- list(
      geom = "scale_y_continuous",
      labels = insight::format_value(
        x = pretty(sort(c(data[[aes$ymin]], data[[aes$ymax]]))),
        as_percent = TRUE,
        digits = 0
      ),
      breaks = pretty(sort(c(data[[aes$ymin]], data[[aes$ymax]])))
    )
    l <- l + 1
  }

  # Out
  class(layers) <- unique(c(
    "visualisation_recipe",
    "see_visualisation_recipe",
    class(layers)
  ))
  attr(layers, "data") <- data
  attr(layers, "global_aes") <- insight::compact_list(global_aes)
  layers
}


# Raw data ----------------------------------------------------------------

#' @keywords internal
.visualization_recipe_rawdata <- function(x, aes, numeric_as_discrete = 8) {
  model <- insight::get_model(x)
  rawdata <- insight::get_data(model, verbose = FALSE)

  # Add response to data if not there
  y <- insight::find_response(attributes(x)$model)
  if (!y %in% names(rawdata)) {
    rawdata[y] <- insight::get_response(attributes(x)$model, verbose = FALSE)
  }

  # if we have less than 8 values for the legend, a continuous color scale
  # is used by default - we then must convert values into factors, when we
  # show data or residuals - but we must ensure that the levels are sorted
  # according to the original data grid, thus we need "sort()"
  if (
    !is.null(aes$color) &&
      is.numeric(rawdata[[aes$color]]) &&
      insight::n_unique(rawdata[[aes$color]]) < numeric_as_discrete
  ) {
    new_values <- insight::format_value(rawdata[[aes$color]], protect_integers = TRUE)
    rawdata[[aes$color]] <- factor(
      new_values,
      levels = as.character(sort(as.numeric(unique(new_values))))
    )
  }

  .data_point_geom(model = model, aes = aes, data = rawdata, y = y)
}


# residuals ----------------------------------------------------------------

#' @keywords internal
.visualization_recipe_residuals <- function(x, aes, numeric_as_discrete = 8) {
  model <- insight::get_model(x)
  residual_data <- residualize_over_grid(x, model)

  # if we have less than 8 values for the legend, a continuous color scale
  # is used by default - we then must convert values into factors, when we
  # show data or residuals - but we must ensure that the levels are sorted
  # according to the original data grid, thus we need "sort()"
  if (
    !is.null(aes$color) &&
      is.numeric(residual_data[[aes$color]]) &&
      insight::n_unique(residual_data[[aes$color]]) < numeric_as_discrete
  ) {
    new_values <- insight::format_value(
      residual_data[[aes$color]],
      protect_integers = TRUE
    )
    residual_data[[aes$color]] <- factor(
      new_values,
      levels = as.character(sort(as.numeric(unique(new_values))))
    )
  }

  .data_point_geom(
    model = model,
    aes = aes,
    data = residual_data,
    y = intersect(colnames(residual_data), c("Mean", "Predicted"))[1]
  )
}


# Collapsed data ----------------------------------------------------------------

#' @keywords internal
.visualization_recipe_collapsed_data <- function(
  x,
  aes,
  collapse_by,
  numeric_as_discrete = 8,
  show_residuals = FALSE
) {
  model <- insight::get_model(x)
  rawdata <- collapse_by_group(x, model, collapse_by, residuals = show_residuals)

  # if we have less than 8 values for the legend, a continuous color scale
  # is used by default - we then must convert values into factors, when we
  # show data or residuals - but we must ensure that the levels are sorted
  # according to the original data grid, thus we need "sort()"
  if (
    !is.null(aes$color) &&
      is.numeric(rawdata[[aes$color]]) &&
      insight::n_unique(rawdata[[aes$color]]) < numeric_as_discrete
  ) {
    new_values <- insight::format_value(rawdata[[aes$color]], protect_integers = TRUE)
    rawdata[[aes$color]] <- factor(
      new_values,
      levels = as.character(sort(as.numeric(unique(new_values))))
    )
  }

  .data_point_geom(
    model = model,
    aes = aes,
    data = rawdata,
    y = insight::find_response(model),
    size = 2,
    jitter = 0.1
  )
}


# helpers -----------------------------------------------------------------

.data_point_geom <- function(model, aes, data, y, ...) {
  if (aes$type == "pointrange" && !is.numeric(data[[aes$x]])) {
    geom <- "jitter"
  } else {
    geom <- "point"
  }

  # Default changes for binomial models
  shape <- 16
  stroke <- 0
  if (insight::model_info(model, response = 1)$is_binomial) {
    shape <- "|"
    stroke <- 1
  }

  out <- list(
    geom = geom,
    data = data,
    aes = list(y = y, x = aes$x, color = aes$color, alpha = aes$alpha),
    height = 0,
    shape = shape,
    stroke = stroke
  )

  dots <- list(...)
  if (!is.null(dots$size)) {
    out$size <- dots$size
  }
  if (!is.null(dots$jitter)) {
    out$position <- "jitter"
    out$width <- dots$jitter
  }

  # check if we have matching columns in the raw data - some functions,
  # likes slopes, have mapped these aes to other columns that are not part
  # of the raw data - we set them to NULL
  if (!is.null(aes$color) && !aes$color %in% colnames(data)) {
    out$aes$color <- NULL
  }
  if (!is.null(aes$alpha) && !aes$alpha %in% colnames(data)) {
    out$aes$alpha <- NULL
  }

  # set default alpha, if not mapped by aes
  if (is.null(aes$alpha)) {
    out$alpha <- getOption("modelbased_ribbon_alpha", 1 / 3)
  } else {
    out$alpha <- NULL
  }

  out
}
