#' Plot Methods for \code{fastml} Objects
#'
#' \code{plot.fastml} produces visual diagnostics for a trained \code{fastml} object.
#'
#' @param x A \code{fastml} object (output of \code{\link{fastml}()}).
#' @param algorithm Character vector specifying which algorithm(s) to include when
#'   generating certain plots (e.g., ROC curves). Defaults to \code{"best"}.
#' @param type Character vector indicating which plot(s) to produce. Options are:
#'   \describe{
#'     \item{\code{"bar"}}{Bar plot of performance metrics across all models/engines.}
#'     \item{\code{"roc"}}{ROC curve(s) for binary classification models.}
#'     \item{\code{"calibration"}}{Calibration plot for the best model(s).}
#'     \item{\code{"residual"}}{Residual diagnostics for the best model.}
#'     \item{\code{"learning_curve"}}{Learning-curve plot if recorded during training.}
#'     \item{\code{"all"}}{Produce all available plots.}
#'   }
#' @param ... Additional arguments (currently unused).
#'
#' @examples
#' \donttest{
#'   ## Create a binary classification dataset from iris
#'   data(iris)
#'   iris <- iris[iris$Species != "setosa",]
#'   iris$Species <- factor(iris$Species)
#'
#'   ## Fit fastml model on binary classification task
#'   model <- fastml(data = iris, label = "Species", algorithms = c("rand_forest", "svm_rbf"))
#'
#'   ## 1. Plot all available diagnostics
#'   plot(model, type = "all")
#'
#'   ## 2. Bar plot of performance metrics
#'   plot(model, type = "bar")
#'
#'   ## 3. ROC curves (only for classification models)
#'   plot(model, type = "roc")
#'
#'   ## 4. Calibration plot (requires 'probably' package)
#'   plot(model, type = "calibration")
#'
#'   ## 5. ROC curves for specific algorithm(s) only
#'   plot(model, type = "roc", algorithm = "rand_forest")
#'
#'   ## 6. Residual diagnostics (only available for regression tasks)
#'   model <- fastml(data = mtcars, label = "mpg", algorithms = c("linear_reg", "xgboost"))
#'   plot(model, type = "residual")
#'}
#'
#' @details
#' When \code{type = "all"}, \code{plot.fastml} will produce a bar plot of metrics,
#' ROC curves (classification), calibration plot, and residual diagnostics (regression).
#' If you specify a subset of types, only those will be drawn.
#'
#' @method plot fastml
#' @importFrom graphics plot
#' @importFrom pROC multiclass.roc
#' @export
plot.fastml <- function(x,
                        algorithm = "best",
                        type = c("all", "bar", "roc", "calibration", "residual", "learning_curve"),
                        ...) {

  if (!inherits(x, "fastml")) {
    stop("The input must be a 'fastml' object.")
  }

  # Validate 'type' argument
  type <- match.arg(type, several.ok = TRUE)
  if ("all" %in% type) {
    type <- c("bar", "roc", "calibration", "residual", "learning_curve")
  }

  performance      <- x$performance
  predictions_list <- x$predictions
  task             <- x$task
  best_model_name  <- x$best_model_name
  optimized_metric <- x$metric
  positive_class   <- x$positive_class
  engine_names     <- x$engine_names
  resampling_plan  <- x$resampling_plan
  resampling_desc  <- fastml_describe_resampling(resampling_plan)

  # Rebuild performance_wide (robust to single- or multi-engine structures)
  metrics_list <- lapply(names(performance), function(model_name) {
    perf_entry <- performance[[model_name]]

    resolve_engine <- function(default_engine = NA_character_) {
      if (!is.null(engine_names) && !is.null(engine_names[[model_name]]) && !is.na(engine_names[[model_name]])) {
        return(engine_names[[model_name]])
      }
      if (!is.null(best_model_name) && model_name %in% names(best_model_name)) {
        return(as.character(best_model_name[[model_name]]))
      }
      default_engine
    }

    if (is.list(perf_entry) && !is.data.frame(perf_entry)) {
      engine_dfs <- lapply(names(perf_entry), function(engine_name) {
        df <- as.data.frame(perf_entry[[engine_name]])
        df$Engine <- engine_name
        df
      })
      combined_engines <- do.call(rbind, engine_dfs)
    } else {
      combined_engines <- as.data.frame(perf_entry)
      combined_engines$Engine <- resolve_engine()
    }

    combined_engines$Model <- model_name
    combined_engines
  })
  performance_df <- do.call(rbind, metrics_list)

  all_metric_names <- unique(performance_df$.metric)
  if (optimized_metric %in% all_metric_names) {
    main_metric <- optimized_metric
  } else {
    main_metric <- all_metric_names[1]
    warning("Optimized metric not available; using first metric.")
  }

  if (task == "classification") {
    desired_metrics <- c("accuracy", "f_meas", "kap", "precision", "sens", "spec", "roc_auc", "logloss", "brier_score")
  } else if (task == "regression") {
    desired_metrics <- c("rmse", "rsq", "mae")
  } else if (task == "survival") {
    brier_metrics <- sort(grep("^brier_t", all_metric_names, value = TRUE))
    desired_metrics <- c("c_index", "uno_c", "ibs", "rmst_diff", brier_metrics)
  } else {
    desired_metrics <- unique(performance_df$.metric)
  }
  desired_metrics <- intersect(desired_metrics, all_metric_names)
  if (length(desired_metrics) == 0) {
    desired_metrics <- main_metric
  }

  performance_sub <- performance_df[
    performance_df$.metric %in% desired_metrics,
  ] %>%
    dplyr::select(-dplyr::any_of(".estimator"))

  if (length(engine_names) == 1 && "LiblineaR" %in% engine_names) {
    performance_wide <- tidyr::pivot_wider(
      performance_sub,
      names_from  = ".metric",
      values_from = ".estimate"
    ) %>%
      dplyr::select(Model, Engine, accuracy, kap, sens, spec, precision, f_meas)
  } else {
    keep_metrics <- desired_metrics
    if (length(engine_names) == 1 && "LiblineaR" %in% engine_names) {
      keep_metrics <- intersect(keep_metrics, c("accuracy", "kap", "sens", "spec", "precision", "f_meas"))
    }
    performance_wide <- tidyr::pivot_wider(
      performance_sub,
      names_from  = ".metric",
      values_from = ".estimate"
    )
    select_cols <- c("Model", "Engine", keep_metrics)
    select_cols <- intersect(select_cols, colnames(performance_wide))
    performance_wide <- dplyr::select(performance_wide, dplyr::all_of(select_cols))
  }

  brier_cols <- grep("^brier_t", colnames(performance_wide), value = TRUE)
  ascending_metrics <- unique(c("rmse", "mae", "ibs", "logloss", "brier_score", "mse", brier_cols))
  if (main_metric %in% ascending_metrics) {
    performance_wide <- performance_wide[order(performance_wide[[main_metric]], na.last = TRUE), ]
  } else {
    performance_wide <- performance_wide[order(-performance_wide[[main_metric]], na.last = TRUE), ]
  }

  display_names <- c(
    accuracy  = "Accuracy",
    f_meas    = "F1 Score",
    kap       = "Kappa",
    precision = "Precision",
    roc_auc   = "ROC AUC",
    sens      = "Sensitivity",
    spec      = "Specificity",
    rsq       = "R-squared",
    mae       = "MAE",
    rmse      = "RMSE",
    c_index   = "Harrell C-index",
    uno_c     = "Uno's C-index",
    ibs       = "Integrated Brier Score",
    rmst_diff = "RMST diff (t_max)"
  )
  t_max_val <- x$survival_t_max
  if (!is.null(t_max_val) && length(t_max_val) == 1 && is.finite(t_max_val) && t_max_val > 0) {
    display_names[["rmst_diff"]] <- sprintf(
      "RMST diff (t<=%s)",
      format(t_max_val, trim = TRUE, digits = 4)
    )
  }
  if (!is.null(x$survival_brier_times)) {
    for (nm in names(x$survival_brier_times)) {
      if (!is.null(x$survival_brier_times[[nm]]) && is.finite(x$survival_brier_times[[nm]])) {
        time_label <- format(x$survival_brier_times[[nm]], trim = TRUE, digits = 4)
      } else {
        time_label <- nm
      }
      display_names[[nm]] <- sprintf("Brier(t=%s)", time_label)
    }
  }
  auto_brier_cols <- setdiff(grep("^brier_t", colnames(performance_wide), value = TRUE), names(display_names))
  if (length(auto_brier_cols) > 0) {
    display_names[auto_brier_cols] <- auto_brier_cols
  }

  # ============================
  # 1. Bar plot of metrics
  # ============================
  if ("bar" %in% type) {
    performance_melt <- performance_wide %>%
      tidyr::pivot_longer(
        cols      = -c(Model, Engine),
        names_to  = "Metric",
        values_to = "Value"
      )
    performance_melt$Value <- as.numeric(performance_melt$Value)
    performance_melt$Metric <- ifelse(
      performance_melt$Metric %in% names(display_names),
      display_names[performance_melt$Metric],
      performance_melt$Metric
    )

    if (task == "classification") {
      class_order <- c("Accuracy", "F1 Score", "Kappa", "Precision", "Sensitivity", "Specificity", "ROC AUC")
      present_class_metrics <- intersect(class_order, unique(performance_melt$Metric))
      if (length(present_class_metrics) > 0) {
        performance_melt$Metric <- factor(performance_melt$Metric, levels = present_class_metrics)
      }
    } else if (task == "regression") {
      reg_order <- c("RMSE", "R-squared", "MAE")
      present_reg_metrics <- intersect(reg_order, unique(performance_melt$Metric))
      if (length(present_reg_metrics) > 0) {
        performance_melt$Metric <- factor(performance_melt$Metric, levels = present_reg_metrics)
      }
    } else if (task == "survival") {
      rmst_label <- display_names[["rmst_diff"]]
      if (is.null(rmst_label)) {
        rmst_label <- "RMST diff (t_max)"
      }
      surv_order <- c("Harrell C-index", "Uno's C-index", "Integrated Brier Score", rmst_label)
      brier_labels <- display_names[grep("^brier_t", names(display_names))]
      surv_order <- c(surv_order, brier_labels)
      present_surv_metrics <- intersect(surv_order, unique(performance_melt$Metric))
      if (length(present_surv_metrics) > 0) {
        performance_melt$Metric <- factor(performance_melt$Metric, levels = present_surv_metrics)
      }
    }

    p_bar <- ggplot2::ggplot(
      performance_melt,
      ggplot2::aes(x = .data$Model, y = .data$Value, fill = .data$Engine)
    ) +
      ggplot2::geom_bar(stat = "identity", position = ggplot2::position_dodge()) +
      ggplot2::facet_wrap(~ Metric, scales = "free_y") +
      ggplot2::theme_bw() +
      ggplot2::theme(axis.text.x = ggplot2::element_text(angle = 45, hjust = 1)) +
      ggplot2::labs(
        title = "Model Performance Comparison",
        x     = "Model",
        y     = "Metric Value"
      )

    if (task == "classification") {
      p_bar <- p_bar + ggplot2::ylim(0, 1)
    }

    print(p_bar)
  }

  # ============================
  # 2. ROC curves (classification)
  # ============================
  if ("roc" %in% type) {
    if (task == "classification" && !is.null(predictions_list) && length(predictions_list) > 0) {

      # Create a list to collect data frames from all models/engines
      dfs <- list()

      # Loop over each algorithm (e.g., "rand_forest", "logistic_reg")

      if(all(algorithm == "best")){
        pred_list = predictions_list
      }else{

        pred_list = predictions_list[algorithm]

      }

      for (algo in names(pred_list)) {
        # Check if predictions are nested by engine or flat (single engine)
        algo_preds <- predictions_list[[algo]]

        if (is.data.frame(algo_preds)) {
          # Flat structure: single engine, get engine name from best_model_name or engine_names
          eng <- if (!is.null(best_model_name) && algo %in% names(best_model_name)) {
            as.character(best_model_name[[algo]])
          } else if (!is.null(engine_names) && algo %in% names(engine_names)) {
            as.character(engine_names[[algo]])
          } else {
            "default"
          }
          engines_to_process <- stats::setNames(list(algo_preds), eng)
        } else {
          # Nested structure: multiple engines
          engines_to_process <- algo_preds
        }

        # Loop over each engine in the current algorithm group
        for (eng in names(engines_to_process)) {
          if(eng == "LiblineaR"){
            warning("Engine 'LiblineaR' does not provide probability predictions; no ROC curve (roc_auc) will be computed.")

          }else{
            df <- engines_to_process[[eng]]
            if (!is.null(df) && "truth" %in% names(df)) {
              # Get probability columns matching ".pred_"

              if(eng == "h2o"){
                prob_cols <- grep("^\\.pred_p", names(df), value = TRUE)
              } else {
                prob_cols <- grep("^\\.pred_", names(df), value = TRUE)
              }

              if (length(prob_cols) >= 2) {  # binary classification should have two probability columns
                # Determine the predictor column for the positive class
                if(eng == "h2o"){
                  pred_col <- paste0(".pred_p", positive_class)

                }else{
                  pred_col <- paste0(".pred_", positive_class)
                }

                if (pred_col %in% prob_cols) {
                  # Add columns for algorithm and engine
                  df$Model  <- algo
                  df$Engine <- eng
                  # Use a compound key to store in the list
                  key <- paste(algo, eng, sep = "_")
                  dfs[[key]] <- df
                }
              }
            }
          }
        }
      }


      if (length(dfs) > 0) {
        # Combine all prediction data frames

        dfs_standardized <- lapply(dfs, function(df) {
          # Check if any column starts with ".pred_p"
          if (all(grepl("^\\.pred_p", names(df)))) {
            df <- df %>%
              rename_with(~ sub("^\\.pred_p", ".pred_", .x), starts_with(".pred_p"))
          }
          df
        })

        dfs_roc <-  bind_rows(dfs_standardized)

        # Ensure 'truth', 'Model', and 'Engine' are factors
        dfs_roc$truth  <- factor(dfs_roc$truth)
        dfs_roc$Model  <- as.factor(dfs_roc$Model)
        dfs_roc$Engine <- as.factor(dfs_roc$Engine)

        # We'll compute ROC curves for each unique combination of Model and Engine.
        pred_col <- paste0(".pred_", positive_class)

        if (length(levels(dfs_roc$truth)) != 2) {
          cat("\nROC curves are only generated for binary classification tasks.\n\n")
        } else if (!(pred_col %in% colnames(dfs_roc))) {
          cat("\nProbability column for the positive class not found; cannot compute ROC curves.\n\n")
        } else {
          roc_list <- list()

          # Get the unique combinations by splitting the compound key back into Model and Engine
          groups <- unique(dfs_roc[, c("Model", "Engine")])
          for (i in seq_len(nrow(groups))) {
            mod <- as.character(groups$Model[i])
            eng <- as.character(groups$Engine[i])
            data_model <- subset(dfs_roc, Model == mod & Engine == eng)
            data_model <- data_model[!is.na(data_model[[pred_col]]), ]

            roc_obj <- pROC::roc(response = data_model$truth,
                                predictor = data_model[[pred_col]],
                                direction = "auto",
                                quiet = TRUE)

            key <- paste(mod, eng, sep = " - ")
            roc_list[[key]] <- roc_obj
          }

          # Compute AUC values for each model/engine combination
          auc_values <- sapply(roc_list, function(x) pROC::auc(x))

          # Sort keys by AUC in alphabetical order of names
          sorted_keys <- auc_values[order(names(auc_values))]

          # Create a data frame for ROC curves by combining the ROC objects
          roc_data <- data.frame()
          for (key in names(sorted_keys)) {
            roc_obj <- roc_list[[key]]
            roc_df <- data.frame(
              FalsePositiveRate = rev(roc_obj$specificities),
              TruePositiveRate  = rev(roc_obj$sensitivities),
              ModelEngine       = key
            )
            roc_data <- rbind(roc_data, roc_df)
          }


          # Create the ROC curve plot, using the compound ModelEngine label for color
          roc_curve_plot <- ggplot(
            data = roc_data,
            aes(
              x = 1 - .data$FalsePositiveRate,
              y = .data$TruePositiveRate,
              color = .data$ModelEngine
            )
          ) +
            geom_line(linewidth = 1) +
            geom_abline(intercept = 0, slope = 1, linetype = "dashed", color = "grey") +
            theme_minimal() +
            labs(title = "ROC Curves for Models",
                 x = "1 - Specificity",
                 y = "Sensitivity") +
            theme(plot.title = element_text(hjust = 0.5)) +
            # Use sorted keys to label the legend with AUC values
            scale_color_manual(values = 1:length(sorted_keys),
                               labels = paste0(names(sorted_keys), " (AUC = ",
                                               sprintf("%.3f", auc_values[names(sorted_keys)]), ")")) +
            theme(legend.title = element_blank())

          print(roc_curve_plot)
        }

      } else {
        cat("\nNo suitable probability predictions for ROC curves.\n\n")
      }

    } else {
      if (is.null(predictions_list) || length(predictions_list) == 0) {
        cat("\nNo predictions available to generate ROC curves.\n\n")
      }
    }

  }

  # ============================
  # 4. Calibration Plot (classification)
  # ============================
  if ("calibration" %in% type) {
    if (task == "classification") {
      if (requireNamespace("probably", quietly = TRUE)) {
        df_best <- list()
        for (model in names(best_model_name)) {
          engine <- best_model_name[model]
          name_combined <- paste(model, "(", engine, ")", sep = "")

          model_preds <- predictions_list[[model]]
          if (!is.null(model_preds)) {
            if (is.data.frame(model_preds)) {
              # Flat structure: predictions stored directly as data.frame
              df_best[[name_combined]] <- model_preds
            } else if (!is.null(model_preds[[engine]])) {
              # Nested structure: predictions stored by engine
              df_best[[name_combined]] <- model_preds[[engine]]
            } else {
              cat("\nNo predictions found for", model, "with engine", engine, "\n")
            }
          } else {
            cat("\nNo predictions found for", model, "with engine", engine, "\n")
          }
        }

        names_df_best <- unique(unlist(lapply(df_best, names)))

        prob_cols <- grep("^\\.pred", names_df_best, value = TRUE)
        if (length(prob_cols) > 1) {
          for (model_name in names(df_best)) {
            model_predictions <- df_best[[model_name]]
            if (grepl("h2o", model_name)) {
              names(model_predictions) <- sub("^\\.pred_p", ".pred_", names(model_predictions))
            }
            pred_col <- paste0(".pred_", positive_class)
            if (!pred_col %in% colnames(model_predictions)) {
              cat("Error: Column", pred_col, "not found in model predictions for", model_name, "\n")
              next
            }
            prob_cols_metric <- setdiff(names(model_predictions), c("truth", "estimate"))
            prob_cols_metric <- prob_cols_metric[
              vapply(model_predictions[prob_cols_metric], is.numeric, logical(1))
            ]
            prob_info <- fastml_prepare_prob_matrix(model_predictions, prob_cols_metric)
            calib_label <- NULL
            if (!is.null(prob_info)) {
              calib <- fastml_class_calibration_metrics(
                truth_vec = model_predictions$truth,
                prob_mat = prob_info$mat,
                prob_cols = prob_info$cols,
                event_class = x$event_class
              )
              calib_vals <- c(calib$logloss, calib$brier_score, calib$ece)
              if (any(is.finite(calib_vals))) {
                calib_label <- paste0(
                  "Logloss: ",
                  ifelse(is.finite(calib$logloss), sprintf("%.3f", calib$logloss), "NA"),
                  ", Brier: ",
                  ifelse(is.finite(calib$brier_score), sprintf("%.3f", calib$brier_score), "NA"),
                  ", ECE: ",
                  ifelse(is.finite(calib$ece), sprintf("%.3f", calib$ece), "NA")
                )
              }
            }
            p_cal <- probably::cal_plot_breaks(
              model_predictions,
              truth    = truth,
              estimate = !!rlang::sym(pred_col),
              event_level = x$event_class
            ) +
              ggplot2::labs(
                title = paste("Calibration Plot for", model_name),
                subtitle = calib_label
              )
            print(p_cal)
          }
        } else {
          cat("\nNot enough probability columns for calibration plots.\n\n")
        }
      } else {
        cat("\nInstall the 'probably' package to generate calibration plots.\n")
      }
    } else {
      cat("\nCalibration plots are only available for classification tasks.\n\n")
    }
  }

  # ============================
  # 5. Residual Diagnostics (regression)
  # ============================
  if ("residual" %in% type) {
    if (task == "regression") {
      df_best <- list()
      # Determine the single best model engine
      for (model in names(best_model_name)) {
        engine <- best_model_name[model]
        name_combined <- paste(model, "(", engine, ")", sep = "")

        model_preds <- predictions_list[[model]]
        if (!is.null(model_preds)) {
          if (is.data.frame(model_preds)) {
            # Flat structure: predictions stored directly as data.frame
            df_best[[name_combined]] <- model_preds
          } else if (!is.null(model_preds[[engine]])) {
            # Nested structure: predictions stored by engine
            df_best[[name_combined]] <- model_preds[[engine]]
          } else {
            cat("\nNo predictions found for", model, "with engine", engine, "\n")
          }
        } else {
          cat("\nNo predictions found for", model, "with engine", engine, "\n")
        }
      }


      names_df_best <- unique(unlist(lapply(df_best, names)))
      if (!is.null(df_best) && length(df_best) > 0 && "truth" %in% names_df_best && "estimate" %in% names_df_best) {
        df_best_all <- dplyr::bind_rows(df_best, .id = "ModelEngine") %>%
          dplyr::mutate(residual = .data$truth - .data$estimate)

        cat("\nResidual Diagnostics for Best Model:\n")

        p_truth_pred <- ggplot2::ggplot(
          df_best_all,
          ggplot2::aes(x = .data$estimate, y = .data$truth)
        ) +
          ggplot2::geom_point(alpha = 0.6) +
          ggplot2::geom_abline(linetype = "dashed", color = "red") +
          ggplot2::labs(
            title = "Truth vs Predicted",
            x = "Predicted",
            y = "Truth"
          ) +
          ggplot2::theme_bw()
        print(p_truth_pred)

        p_resid_hist <- ggplot2::ggplot(
          df_best_all,
          ggplot2::aes(x = .data$residual)
        ) +
          ggplot2::geom_histogram(bins = 30, fill = "steelblue", color = "white", alpha = 0.7) +
          ggplot2::labs(
            title = "Residual Distribution",
            x = "Residual",
            y = "Count"
          ) +
          ggplot2::theme_bw()
        print(p_resid_hist)
      } else {
        cat("\nNo valid predictions to compute residual diagnostics.\n\n")
      }
    } else {
      cat("\nResidual diagnostics are only available for regression tasks.\n\n")
    }
  }

  # 5. Learning curve (optional)
  if ("learning_curve" %in% type) {
    lc <- x$learning_curve
    if (!is.null(lc) && !is.null(lc$plot)) {
      print(lc$plot)
    } else if (!is.null(lc) && !is.null(lc$data)) {
      lc_plot <- ggplot2::ggplot(
        lc$data,
        ggplot2::aes(x = .data$Fraction, y = .data$Performance)
      ) +
        ggplot2::geom_line(color = "blue") +
        ggplot2::geom_point(color = "blue") +
        ggplot2::labs(
          title = "Learning Curve",
          x = "Training Set Size (fraction)",
          y = paste("Mean", optimized_metric, "across models")
        ) +
        ggplot2::theme_minimal()
      print(lc_plot)
    } else {
      cat("\nNo learning curve data available. Set `learning_curve = TRUE` when fitting to record it.\n\n")
    }
  }

  invisible(x)
}
