#' @title Interaction Analysis Functions for tidylearn
#' @name tidylearn-interactions
#' @description Functions for testing, visualizing, and analyzing interactions
#' @importFrom stats lm anova
#' @importFrom dplyr filter select mutate
#' @importFrom ggplot2 ggplot aes geom_line geom_point
#'   facet_wrap labs theme_minimal
NULL

#' Test for significant interactions between variables
#'
#' @param data A data frame containing the data
#' @param formula A formula specifying the base model without interactions
#' @param var1 First variable to test for interactions
#' @param var2 Second variable to test for interactions
#'   (if NULL, tests var1 with all others)
#' @param all_pairs Logical; whether to test all variable pairs
#' @param categorical_only Logical; whether to only test categorical variables
#' @param numeric_only Logical; whether to only test numeric variables
#' @param mixed_only Logical; whether to only test numeric-categorical pairs
#' @param alpha Significance level for interaction tests
#' @return A data frame with one row per tested interaction pair, containing
#'   columns \code{var1}, \code{var2}, \code{p_value}, \code{significant}
#'   (logical), \code{delta_r2} (change in R-squared), and
#'   \code{f_statistic}, sorted by \code{p_value} ascending.
#' @examples
#' \donttest{
#' results <- tl_test_interactions(mtcars, mpg ~ wt + hp + cyl,
#'   var1 = "wt", var2 = "hp")
#' }
#' @export
tl_test_interactions <- function(data, formula, var1 = NULL, var2 = NULL,
                                 all_pairs = FALSE, categorical_only = FALSE,
                                 numeric_only = FALSE, mixed_only = FALSE,
                                 alpha = 0.05) {
  # Extract predictor variables
  predictors <- all.vars(formula)[-1]

  # Categorize variables
  var_types <- sapply(data[predictors], function(x) {
    if (is.factor(x) || is.character(x)) {
      "categorical"
    } else if (is.numeric(x)) {
      "numeric"
    } else {
      "other"
    }
  })

  # Generate pairs to test
  if (all_pairs) {
    # Test all variable pairs
    pairs <- combn(predictors, 2, simplify = FALSE)
  } else if (!is.null(var1) && !is.null(var2)) {
    # Test specific pair
    pairs <- list(c(var1, var2))
  } else if (!is.null(var1)) {
    # Test var1 with all other variables
    pairs <- lapply(setdiff(predictors, var1), function(v) c(var1, v))
  } else {
    stop("Must specify at least var1 or set all_pairs = TRUE", call. = FALSE)
  }

  # Filter pairs based on variable types
  if (categorical_only) {
    pairs <- pairs[sapply(
      pairs, function(p) all(var_types[p] == "categorical")
    )]
  } else if (numeric_only) {
    pairs <- pairs[sapply(pairs, function(p) all(var_types[p] == "numeric"))]
  } else if (mixed_only) {
    pairs <- pairs[sapply(pairs, function(p) {
      var_types[p[1]] != var_types[p[2]] &&
        all(var_types[p] %in% c("categorical", "numeric"))
    })]
  }

  # Test interactions
  results <- lapply(pairs, function(pair) {
    # Build base model
    base_model <- lm(formula, data = data)

    # Build model with interaction
    int_formula <- update(formula, paste0(". ~ . + ", pair[1], ":", pair[2]))
    int_model <- lm(int_formula, data = data)

    # Perform ANOVA to compare models
    models_comparison <- anova(base_model, int_model)

    # Extract p-value for interaction term
    p_value <- models_comparison$`Pr(>F)`[2]

    # Calculate interaction effect size (change in R-squared)
    r2_base <- summary(base_model)$r.squared
    r2_int <- summary(int_model)$r.squared
    delta_r2 <- r2_int - r2_base

    # Create result row
    data.frame(
      var1 = pair[1],
      var2 = pair[2],
      p_value = p_value,
      significant = p_value < alpha,
      delta_r2 = delta_r2,
      f_statistic = models_comparison$F[2]
    )
  })

  # Combine results
  all_results <- do.call(rbind, results)

  # Sort by significance
  all_results <- all_results[order(all_results$p_value), ]

  all_results
}

#' Plot interaction effects
#'
#' @param model A tidylearn model object
#' @param var1 First variable in the interaction
#' @param var2 Second variable in the interaction
#' @param n_points Number of points to use for continuous variables
#' @param fixed_values Named list of values for other variables in the model
#' @param confidence Logical; whether to show confidence intervals
#' @param ... Additional arguments to pass to predict()
#' @return A \code{\link[ggplot2]{ggplot}} object.
#' @export
tl_plot_interaction <- function(model, var1, var2,
                                n_points = 100,
                                fixed_values = NULL,
                                confidence = TRUE,
                                ...) {
  # Extract data
  data <- model$data

  # Check if variables exist in the model
  formula <- model$spec$formula
  all_vars <- all.vars(formula)[-1] # Remove response variable
  if (!var1 %in% all_vars || !var2 %in% all_vars) {
    stop("Variables not found in model formula", call. = FALSE)
  }

  # Determine variable types
  var1_type <- if (
    is.factor(data[[var1]]) || is.character(data[[var1]])
  ) "categorical" else "numeric"
  var2_type <- if (
    is.factor(data[[var2]]) || is.character(data[[var2]])
  ) "categorical" else "numeric"

  # Create grid of values
  if (var1_type == "categorical") {
    if (is.factor(data[[var1]])) {
      var1_values <- levels(data[[var1]])
    } else {
      var1_values <- unique(data[[var1]])
    }
  } else {
    var1_values <- seq(min(data[[var1]], na.rm = TRUE),
                       max(data[[var1]], na.rm = TRUE),
                       length.out = n_points)
  }

  if (var2_type == "categorical") {
    if (is.factor(data[[var2]])) {
      var2_values <- levels(data[[var2]])
    } else {
      var2_values <- unique(data[[var2]])
    }
  } else {
    var2_values <- seq(min(data[[var2]], na.rm = TRUE),
                       max(data[[var2]], na.rm = TRUE),
                       length.out = n_points)
  }

  # Create grid of all combinations
  grid <- expand.grid(
    var1 = var1_values,
    var2 = var2_values
  )
  names(grid) <- c(var1, var2)

  # Add fixed values for other variables
  all_other_vars <- setdiff(all_vars, c(var1, var2))

  for (v in all_other_vars) {
    if (!is.null(fixed_values) && v %in% names(fixed_values)) {
      # Use provided value
      grid[[v]] <- fixed_values[[v]]
    } else {
      # Use median/mode as default
      if (is.numeric(data[[v]])) {
        grid[[v]] <- median(data[[v]], na.rm = TRUE)
      } else if (is.factor(data[[v]]) || is.character(data[[v]])) {
        # Use most frequent value
        tab <- table(data[[v]])
        grid[[v]] <- names(tab)[which.max(tab)]

        # Convert to factor if necessary
        if (is.factor(data[[v]])) {
          grid[[v]] <- factor(grid[[v]], levels = levels(data[[v]]))
        }
      } else {
        # Use first value
        grid[[v]] <- data[[v]][1]
      }
    }
  }

  # Make predictions
  predictions <- predict(model, grid, ...)

  # Add predictions to grid -- tidylearn predict returns tibble with .pred
  if (is.data.frame(predictions) && ".pred" %in% names(predictions)) {
    grid$prediction <- predictions$.pred
    y_col <- "prediction"
    lower_col <- NULL
    upper_col <- NULL
  } else if (is.data.frame(predictions)) {
    # Non-tidylearn output (e.g., confidence intervals with fit/lwr/upr)
    grid <- cbind(grid, predictions)
    y_col <- if ("fit" %in% names(predictions)) "fit" else names(predictions)[1]
    lower_col <- if ("lwr" %in% names(predictions)) "lwr" else NULL
    upper_col <- if ("upr" %in% names(predictions)) "upr" else NULL
  } else {
    grid$prediction <- predictions
    y_col <- "prediction"
    lower_col <- NULL
    upper_col <- NULL
  }

  # Create plot
  if (var1_type == "numeric" && var2_type == "categorical") {
    # Line plot with x = var1, color = var2
    p <- ggplot2::ggplot(grid, ggplot2::aes(
      x = .data[[var1]], y = .data[[y_col]],
      color = .data[[var2]]
    )) +
      ggplot2::geom_line() +
      ggplot2::labs(
        title = paste("Interaction between", var1, "and", var2),
        x = var1,
        y = "Predicted Value",
        color = var2
      )

    # Add confidence intervals if available
    if (confidence && !is.null(lower_col) && !is.null(upper_col)) {
      p <- p + ggplot2::geom_ribbon(
        ggplot2::aes(
          ymin = .data[[lower_col]],
          ymax = .data[[upper_col]],
          fill = .data[[var2]]
        ),
        alpha = 0.2,
        linetype = 0
      )
    }
  } else if (var1_type == "categorical" && var2_type == "numeric") {
    # Line plot with x = var2, color = var1
    p <- ggplot2::ggplot(grid, ggplot2::aes(
      x = .data[[var2]], y = .data[[y_col]],
      color = .data[[var1]]
    )) +
      ggplot2::geom_line() +
      ggplot2::labs(
        title = paste("Interaction between", var1, "and", var2),
        x = var2,
        y = "Predicted Value",
        color = var1
      )

    # Add confidence intervals if available
    if (confidence && !is.null(lower_col) && !is.null(upper_col)) {
      p <- p + ggplot2::geom_ribbon(
        ggplot2::aes(
          ymin = .data[[lower_col]],
          ymax = .data[[upper_col]],
          fill = .data[[var1]]
        ),
        alpha = 0.2,
        linetype = 0
      )
    }
  } else if (var1_type == "numeric" && var2_type == "numeric") {
    # Contour plot or heat map
    p <- ggplot2::ggplot(grid, ggplot2::aes(
      x = .data[[var1]], y = .data[[var2]],
      z = .data[[y_col]]
    )) +
      ggplot2::geom_contour_filled() +
      ggplot2::labs(
        title = paste("Interaction between", var1, "and", var2),
        x = var1,
        y = var2,
        fill = "Predicted Value"
      )
  } else {
    # Categorical x Categorical: Faceted bar plot
    p <- ggplot2::ggplot(grid, ggplot2::aes(
      x = .data[[var1]], y = .data[[y_col]],
      fill = .data[[var2]]
    )) +
      ggplot2::geom_col(position = "dodge") +
      ggplot2::labs(
        title = paste("Interaction between", var1, "and", var2),
        x = var1,
        y = "Predicted Value",
        fill = var2
      )
  }

  # Apply minimal theme
  p <- p + ggplot2::theme_minimal()

  p
}

#' Find important interactions automatically
#'
#' @param data A data frame containing the data
#' @param formula A formula specifying the base model without interactions
#' @param top_n Number of top interactions to return
#' @param min_r2_change Minimum change in R-squared to consider
#' @param max_p_value Maximum p-value for significance
#' @param exclude_vars Character vector of variables to exclude
#'   from interaction testing
#' @return A tidylearn model object (class \code{"tidylearn_model"}) fitted
#'   with the top significant interaction terms added to the formula.
#'   The interaction test results and selected interactions are stored as
#'   attributes \code{"interaction_tests"} and
#'   \code{"selected_interactions"}.
#' @examples
#' \donttest{
#' model <- tl_auto_interactions(mtcars, mpg ~ wt + hp + cyl, top_n = 2)
#' }
#' @export
tl_auto_interactions <- function(data, formula, top_n = 3, min_r2_change = 0.01,
                                 max_p_value = 0.05, exclude_vars = NULL) {
  # Extract predictor variables
  predictors <- all.vars(formula)[-1]

  # Remove excluded variables
  if (!is.null(exclude_vars)) {
    predictors <- setdiff(predictors, exclude_vars)
  }

  # Test all interactions
  test_results <- tl_test_interactions(
    data, formula, all_pairs = TRUE,
    alpha = max_p_value
  )

  # Filter significant interactions
  significant <- test_results %>%
    dplyr::filter(.data$p_value < max_p_value, .data$delta_r2 >= min_r2_change)

  # Select top interactions
  if (nrow(significant) > top_n) {
    top_interactions <- significant[1:top_n, ]
  } else {
    top_interactions <- significant
  }

  if (nrow(top_interactions) == 0) {
    message("No significant interactions found based on criteria")
    # Return original model
    return(tl_model(data, formula, method = "linear"))
  }

  # Build formula with interactions
  interaction_terms <- apply(top_interactions, 1, function(row) {
    paste0(row["var1"], ":", row["var2"])
  })

  int_str <- paste(interaction_terms, collapse = " + ")
  new_formula <- update(
    formula, paste0(". ~ . +", int_str)
  )

  # Fit model with interactions
  interaction_model <- tl_model(data, new_formula, method = "linear")

  # Add interaction test results to model
  attr(interaction_model, "interaction_tests") <- test_results
  attr(interaction_model, "selected_interactions") <- top_interactions

  interaction_model
}

#' Calculate partial effects based on a model with interactions
#'
#' @param model A tidylearn model object
#' @param var Variable to calculate effects for
#' @param by_var Variable to calculate effects by (interaction variable)
#' @param at_values Named list of values at which to hold other variables
#' @param intervals Logical; whether to include confidence intervals
#' @return For numeric \code{var}: a list with \code{effects} (data frame of
#'   predicted values across the variable range for each level of
#'   \code{by_var}) and \code{slopes} (data frame with estimated slopes and
#'   standard errors per level). For categorical \code{var}: a data frame of
#'   predicted values at each factor level for each level of \code{by_var}.
#' @export
tl_interaction_effects <- function(model, var, by_var,
                                   at_values = NULL,
                                   intervals = TRUE) {
  # Extract data
  data <- model$data
  formula <- model$spec$formula

  # Check if variables exist in the model
  all_vars <- all.vars(formula)[-1] # Remove response variable
  if (!var %in% all_vars || !by_var %in% all_vars) {
    stop("Variables not found in model formula", call. = FALSE)
  }

  # Check if the interaction term exists in the model
  int_term <- paste0(var, ":", by_var)
  rev_int_term <- paste0(by_var, ":", var)

  has_interaction <- int_term %in% attr(terms(formula), "term.labels") ||
    rev_int_term %in% attr(terms(formula), "term.labels")

  if (!has_interaction) {
    warning(paste("Interaction term", int_term, "not found in model formula"),
            call. = FALSE)
  }

  # Determine by_var values
  if (is.factor(data[[by_var]]) || is.character(data[[by_var]])) {
    if (is.factor(data[[by_var]])) {
      by_values <- levels(data[[by_var]])
    } else {
      by_values <- unique(data[[by_var]])
    }
  } else {
    # For continuous by_var, use quantiles
    by_values <- stats::quantile(
      data[[by_var]], probs = seq(0, 1, 0.25),
      na.rm = TRUE
    )
    names(by_values) <- paste0("Q", seq(0, 100, 25))
  }

  # Set up variable values
  if (is.factor(data[[var]]) || is.character(data[[var]])) {
    if (is.factor(data[[var]])) {
      var_values <- levels(data[[var]])
    } else {
      var_values <- unique(data[[var]])
    }
    var_is_numeric <- FALSE
  } else {
    var_range <- range(data[[var]], na.rm = TRUE)
    var_values <- seq(var_range[1], var_range[2], length.out = 100)
    var_is_numeric <- TRUE
  }

  # Create grid for predictions
  grid_list <- list()

  for (bv in by_values) {
    # Create data frame for this by_value
    grid <- data.frame(var_values)
    names(grid) <- var
    grid[[by_var]] <- bv

    # Add values for other variables
    other_vars <- setdiff(all_vars, c(var, by_var))
    for (v in other_vars) {
      if (!is.null(at_values) && v %in% names(at_values)) {
        grid[[v]] <- at_values[[v]]
      } else {
        # Use median/mode
        if (is.numeric(data[[v]])) {
          grid[[v]] <- median(data[[v]], na.rm = TRUE)
        } else if (is.factor(data[[v]])) {
          grid[[v]] <- levels(data[[v]])[1]
        } else {
          grid[[v]] <- unique(data[[v]])[1]
        }
      }
    }

    # Make predictions -- use raw model for se.fit (tidylearn predict
    # doesn't support it), and extract .pred for the plain case
    if (intervals) {
      raw_preds <- stats::predict(model$fit, newdata = grid, se.fit = TRUE)
      grid$fit <- as.vector(raw_preds$fit)
      grid$se <- as.vector(raw_preds$se.fit)
      grid$lower <- grid$fit - 1.96 * grid$se
      grid$upper <- grid$fit + 1.96 * grid$se
    } else {
      preds <- predict(model, grid)
      grid$fit <- if (is.data.frame(preds)) preds$.pred else preds
    }

    # Add by_value label
    grid$by_value <- bv
    if (!is.null(names(by_values)) && !is.na(match(bv, by_values))) {
      grid$by_label <- names(by_values)[match(bv, by_values)]
    } else {
      grid$by_label <- as.character(bv)
    }

    # Store in list
    grid_list[[length(grid_list) + 1]] <- grid
  }

  # Combine all grids
  final_grid <- do.call(rbind, grid_list)

  # Calculate slopes for numeric variables
  if (var_is_numeric) {
    # Group by by_value
    slopes <- lapply(by_values, function(bv) {
      # Get data for this by_value
      sub_grid <- final_grid[final_grid$by_value == bv, ]

      # If only one value, can't compute slope
      if (nrow(sub_grid) <= 1) {
        data.frame(
          by_value = bv,
          by_label = if (is.null(names(by_values))) {
            as.character(bv)
          } else {
            names(by_values)[match(bv, by_values)]
          },
          slope = NA,
          slope_se = NA
        )
      }

      # Fit linear model to get slope
      slope_formula <- stats::as.formula(paste("fit ~", var))
      slope_model <- lm(slope_formula, data = sub_grid)
      slope_coef <- coef(summary(slope_model))

      data.frame(
        by_value = bv,
        by_label = if (is.null(names(by_values))) {
          as.character(bv)
        } else {
          names(by_values)[match(bv, by_values)]
        },
        slope = slope_coef[2, 1],
        slope_se = slope_coef[2, 2]
      )
    })

    # Combine all slopes
    final_slopes <- do.call(rbind, slopes)

    # Return both grid and slopes
    list(
      effects = final_grid,
      slopes = final_slopes
    )
  } else {
    # For categorical variables, just return the effects
    final_grid
  }
}
