####**********************************************************************
####**********************************************************************
####
####  ----------------------------------------------------------------
####  Written by:
####  ----------------------------------------------------------------
####    John Ehrlinger, Ph.D.
####
####    email:  john.ehrlinger@gmail.com
####    URL:    https://github.com/ehrlinger/ggRandomForests
####  ----------------------------------------------------------------
####
####**********************************************************************
####**********************************************************************

# Map partial.type ("surv" / "chf" / "mort") to a human y-axis label.
# Falls back to "Predicted Survival" when the attribute is absent (e.g. an
# object built before this attribute was introduced).
partial_surv_y_label <- function(partial.type) {
  if (is.null(partial.type)) return("Predicted Survival")
  switch(partial.type,
         surv = "Predicted Survival",
         chf  = "Predicted CHF",
         mort = "Predicted Mortality",
         "Predicted Survival")
}

#' Plot a \code{\link{gg_partial}} object
#'
#' Turns a \code{\link{gg_partial}} object into a ggplot2 figure.  Each curve
#' is a partial dependence trace -- the forest's average prediction as one
#' predictor is swept across its range while the rest are marginalized over the
#' training data.  Continuous predictors appear as line plots; categorical
#' predictors appear as bar charts.  Both panels are faceted by variable name
#' so you can compare the shape and scale of each variable's effect at a
#' glance.
#'
#' When a \code{model} label was attached in \code{gg_partial()}, lines are
#' coloured by model -- handy for overlaying results from two forests (e.g.,
#' one tuned, one default) in the same figure.
#'
#' @param x A \code{\link{gg_partial}} object (output of \code{\link{gg_partial}}).
#' @param ... Not currently used; reserved for future arguments.
#'
#' @return A \code{ggplot} (or \code{patchwork}) object.  When only one
#'   variable type is present a single \code{ggplot} is returned.  When both
#'   continuous and categorical variables are present the two panels are
#'   combined vertically via \code{patchwork::wrap_plots()}, which also
#'   satisfies \code{inherits(p, "ggplot")}.
#'
#' @seealso \code{\link{gg_partial}}, \code{\link{plot.gg_variable}}
#'
#' @examples
#' set.seed(42)
#' airq <- na.omit(airquality)
#' rf <- randomForestSRC::rfsrc(Ozone ~ ., data = airq, ntree = 50)
#' pv <- randomForestSRC::plot.variable(rf, partial = TRUE, show.plots = FALSE)
#' pd <- gg_partial(pv)
#' plot(pd)
#'
#' @importFrom ggplot2 .data
#' @importFrom patchwork wrap_plots
#' @export
plot.gg_partial <- function(x, ...) {
  gg_dta <- x

  gg_cont <- NULL
  if (!is.null(gg_dta$continuous) && nrow(gg_dta$continuous) > 0) {
    cont <- gg_dta$continuous
    gg_cont <- ggplot2::ggplot(cont,
                               ggplot2::aes(x = .data$x, y = .data$yhat)) +
      ggplot2::geom_line()

    if ("model" %in% colnames(cont)) {
      gg_cont <- gg_cont +
        ggplot2::aes(color = .data$model, group = .data$model)
    }

    gg_cont <- gg_cont +
      ggplot2::facet_wrap(~name, scales = "free_x") +
      ggplot2::labs(x = NULL, y = "Partial Effect")
  }

  gg_cat <- NULL
  if (!is.null(gg_dta$categorical) && nrow(gg_dta$categorical) > 0) {
    cat_dta <- gg_dta$categorical
    gg_cat <- ggplot2::ggplot(cat_dta,
                              ggplot2::aes(x = .data$x, y = .data$yhat)) +
      ggplot2::geom_bar(stat = "identity", width = 0.5) +
      ggplot2::facet_wrap(~name, scales = "free_x") +
      ggplot2::labs(x = NULL, y = "Partial Effect")
  }

  if (!is.null(gg_cont) && !is.null(gg_cat)) {
    wrap_plots(gg_cont, gg_cat, ncol = 1)
  } else if (!is.null(gg_cont)) {
    gg_cont
  } else {
    gg_cat
  }
}

#' Plot a \code{\link{gg_partial_rfsrc}} object
#'
#' Renders the partial dependence curves from \code{\link{gg_partial_rfsrc}}
#' as a ggplot2 figure.  The layout adapts automatically to what the object
#' contains.
#'
#' For a standard regression or classification forest, continuous predictors
#' are drawn as line plots and categorical predictors as bar charts, both
#' faceted by variable name -- the same arrangement as
#' \code{\link{plot.gg_partial}}.
#'
#' For a survival forest, each call to \code{partial.rfsrc} returns a predicted
#' quantity (survival probability, cumulative hazard function, or mortality) at
#' one or more chosen time horizons.  When a \code{time} column is present in
#' the data, each horizon becomes a separate coloured curve over the predictor's
#' value, still faceted by variable.  The y-axis label (\dQuote{Predicted
#' Survival}, \dQuote{Predicted CHF}, or \dQuote{Predicted Mortality}) tracks
#' the \code{partial.type} attribute set by \code{gg_partial_rfsrc()}.
#'
#' For a two-variable interaction surface (when \code{xvar2.name} was supplied
#' to \code{gg_partial_rfsrc}), the secondary variable's levels become
#' separate coloured lines, faceted by the primary predictor.
#'
#' @param x A \code{\link{gg_partial_rfsrc}} object.
#' @param ... Not currently used.
#'
#' @return A \code{ggplot} (or \code{patchwork}) object.  When both continuous
#'   and categorical variables are present the two panels are combined
#'   vertically via \code{patchwork::wrap_plots()}.
#'
#' @seealso \code{\link{gg_partial_rfsrc}}, \code{\link{plot.gg_partial}}
#'
#' @examples
#' ## ------------------------------------------------------------
#' ## Regression forest -- one continuous curve per variable
#' ## ------------------------------------------------------------
#' set.seed(42)
#' airq <- na.omit(airquality)
#' rfsrc_airq <- randomForestSRC::rfsrc(Ozone ~ ., data = airq, ntree = 50)
#'
#' pd <- gg_partial_rfsrc(rfsrc_airq, xvar.names = c("Wind", "Temp"),
#'                        n_eval = 10)
#' plot(pd)
#'
#' \donttest{
#' ## ------------------------------------------------------------
#' ## Survival forest -- one curve per requested time horizon,
#' ## faceted by variable. Y-axis label tracks `partial.type`.
#' ## ------------------------------------------------------------
#' # randomForestSRC's formula parser requires the unqualified Surv() symbol;
#' # it Depends on `survival`, so Surv is on the search path once
#' # randomForestSRC is loaded.
#' data(veteran, package = "randomForestSRC")
#' set.seed(42)
#' rfsrc_v <- randomForestSRC::rfsrc(Surv(time, status) ~ .,
#'                                   data = veteran, ntree = 50)
#' ti  <- rfsrc_v$time.interest
#' t30 <- ti[which.min(abs(ti - 30))]
#' t90 <- ti[which.min(abs(ti - 90))]
#'
#' # Default partial.type = "surv" -> y-axis "Predicted Survival"
#' pd_s <- gg_partial_rfsrc(rfsrc_v, xvar.names = "age",
#'                          partial.time = c(t30, t90), n_eval = 8)
#' plot(pd_s)
#'
#' # partial.type = "chf" -> y-axis "Predicted CHF"
#' pd_c <- gg_partial_rfsrc(rfsrc_v, xvar.names = "age",
#'                          partial.time = c(t30, t90),
#'                          partial.type = "chf", n_eval = 8)
#' plot(pd_c)
#' }
#'
#' @importFrom ggplot2 .data
#' @importFrom patchwork wrap_plots
#' @export
plot.gg_partial_rfsrc <- function(x, ...) {
  gg_dta <- x

  gg_cont <- NULL
  if (!is.null(gg_dta$continuous) && nrow(gg_dta$continuous) > 0) {
    cont <- gg_dta$continuous

    if (!is.null(cont$time)) {
      ## Survival forest: predictor value on x-axis, one curve per time point.
      ## Group/colour by the *full-precision* time so distinct horizons that
      ## happen to round to the same value are not silently merged. The legend
      ## is relabelled with rounded values for readability.
      time_levels <- sort(unique(cont$time))
      cont$.time_factor <- factor(cont$time, levels = time_levels)
      legend_labels <- format(round(time_levels, 2), trim = TRUE)
      y_lab <- partial_surv_y_label(attr(gg_dta, "partial.type"))
      gg_cont <- ggplot2::ggplot(
        cont,
        ggplot2::aes(
          x     = .data$x,
          y     = .data$yhat,
          color = .data$.time_factor,
          group = .data$.time_factor
        )
      ) +
        ggplot2::geom_line() +
        ggplot2::facet_wrap(~name, scales = "free_x") +
        ggplot2::scale_color_discrete(labels = legend_labels) +
        ggplot2::labs(x = NULL, y = y_lab, color = "Time")

    } else if (!is.null(cont$grp)) {
      ## Two-variable surface: group is xvar2; x-axis is the primary predictor
      gg_cont <- ggplot2::ggplot(
        cont,
        ggplot2::aes(
          x     = .data$x,
          y     = .data$yhat,
          color = factor(.data$grp),
          group = factor(.data$grp)
        )
      ) +
        ggplot2::geom_line() +
        ggplot2::facet_wrap(~name, scales = "free_x") +
        ggplot2::labs(x = NULL, y = "Partial Effect", color = "Group")

    } else {
      ## Standard: one curve per variable
      gg_cont <- ggplot2::ggplot(cont,
                                 ggplot2::aes(x = .data$x, y = .data$yhat)) +
        ggplot2::geom_line() +
        ggplot2::facet_wrap(~name, scales = "free_x") +
        ggplot2::labs(x = NULL, y = "Partial Effect")
    }
  }

  gg_cat <- NULL
  if (!is.null(gg_dta$categorical) && nrow(gg_dta$categorical) > 0) {
    cat_dta <- gg_dta$categorical
    gg_cat <- ggplot2::ggplot(
      cat_dta,
      ggplot2::aes(x = factor(.data$x), y = .data$yhat)
    ) +
      ggplot2::geom_bar(stat = "identity", width = 0.5) +
      ggplot2::facet_wrap(~name, scales = "free_x") +
      ggplot2::labs(x = NULL, y = "Partial Effect")
  }

  if (!is.null(gg_cont) && !is.null(gg_cat)) {
    wrap_plots(gg_cont, gg_cat, ncol = 1)
  } else if (!is.null(gg_cont)) {
    gg_cont
  } else {
    gg_cat
  }
}

#' @rdname plot.gg_partial_varpro
#' @name plot.gg_partial_varpro
#' @export
plot.gg_partialpro <- function(x, type = c("parametric", "nonparametric",
                                            "causal"), ...) {
  ## Deprecated class shim: re-dispatch to plot.gg_partial_varpro.
  class(x) <- c("gg_partial_varpro", setdiff(class(x), "gg_partialpro"))
  plot.gg_partial_varpro(x, type = type, ...)
}
