R/plot.R

In bayesDP: Implementation of the Bayesian Discount Prior Approach for Clinical Trials

#' @title bdpnormal Object Plot
#' @description \code{plot} method for class \code{bdpnormal}.
#'
#' @param x object of class \code{bdpnormal}. The result of a call to the
#'   \code{\link{bdpnormal}} function.
#' @param type character. Optional. Select plot type to print. Supports the
#'   following: "discount" gives the discount function; "posteriors" gives the
#'   posterior plots of the event rates; and "density" gives the augmented
#'   posterior density plot(s). Leave NULL to display all plots in sequence.
#' @param print logical. Optional. Produce a plot (\code{print = TRUE}; default)
#'   or return a ggplot object (\code{print = FALSE}). When \code{print =
#'   FALSE}, it is possible to use \code{ggplot2} syntax to modify the plot
#'   appearance.
#' @details The \code{plot} method for \code{bdpnormal} returns up to three
#'   plots: (1) posterior density curves; (2) posterior density of the effect of
#'   interest; and (3) the discount function. A call to \code{plot} that omits
#'   the \code{type} input returns one plot at a time and prompts the user to
#'   click the plot or press return to proceed to the next plot. Otherwise, the
#'   user can specify a plot \code{type} to display the requested plot.
#'
#' @import methods
#' @importFrom utils head
#' @importFrom ggplot2 aes_string ggtitle ylim guides guide_legend theme
#'   element_blank
#' @importFrom graphics par
#' @importFrom stats sd density is.empty.model median model.offset
#'   model.response pweibull quantile rbeta rgamma rnorm var vcov
#' @export
setMethod("plot", signature(x = "bdpnormal"), function(x, type = NULL, print = TRUE) {
  posterior_treatment <- x$posterior_treatment
  posterior_control <- x$posterior_control
  discount_function <- x$args1$discount_function
  N0_t <- x$args1$N0_t
  N0_c <- x$args1$N0_c
  N_t <- x$args1$N_t
  N_c <- x$args1$N_c
  arm2 <- x$args1$arm2
  method <- x$args1$method

  information_sources <- NULL

  ### Create data frames for plotting via ggplot2
  D1 <- D2 <- D3 <- D5 <- D6 <- NULL

  if (arm2) {
    dens1 <- density(posterior_control$posterior_mu,
                     adjust = 0.5
    )
    D1 <- data.frame(
      information_sources = "Posterior",
      group = "Control",
      x = dens1$x,
      y = dens1$y
    )

    if (!is.null(posterior_control$posterior_flat_mu)) {
      dens2 <- density(posterior_control$posterior_flat_mu,
                       adjust = 0.5
      )
      D2 <- data.frame(
        information_sources = "Current Data",
        group = "Control",
        x = dens2$x,
        y = dens2$y
      )
    }

    if (!is.null(posterior_control$prior_mu)) {
      dens3 <- density(posterior_control$prior_mu,
                       adjust = 0.5
      )
      D3 <- data.frame(
        information_sources = "Historical Data",
        group = "Control",
        x = dens3$x,
        y = dens3$y
      )
    }
  }

  dens4 <- density(posterior_treatment$posterior_mu,
                   adjust = 0.5
  )
  D4 <- data.frame(
    information_sources = "Posterior",
    group = "Treatment",
    x = dens4$x,
    y = dens4$y
  )

  if (!is.null(posterior_treatment$posterior_flat_mu)) {
    dens5 <- density(posterior_treatment$posterior_flat_mu,
                     adjust = 0.5
    )
    D5 <- data.frame(
      information_sources = "Current Data",
      group = "Treatment",
      x = dens5$x,
      y = dens5$y
    )
  }

  if (!is.null(posterior_treatment$prior_mu)) {
    dens6 <- density(posterior_treatment$prior_mu,
                     adjust = 0.5
    )
    D6 <- data.frame(
      information_sources = "Historical Data",
      group = "Treatment",
      x = dens6$x,
      y = dens6$y
    )
  }

  D <- rbind(D1, D2, D3, D4, D5, D6)

  D$information_sources <- factor(D$information_sources,
                                  levels = c("Posterior", "Current Data", "Historical Data")
  )
  D$group <- factor(D$group, levels = c("Treatment", "Control"))

  ##############################################################################
  ### Posterior Type Plots
  ##############################################################################

  post_typeplot <- ggplot(D, aes_string(x = "x", y = "y")) +
    geom_line(
      size = 2,
      aes_string(
        colour = "information_sources",
        lty = "information_sources"
      )
    ) +
    facet_wrap(~group, ncol = 1, scales = "free") +
    ylab("Density (PDF)") +
    xlab("Values") +
    theme_bw() +
    ggtitle("Posterior Type Plot") +
    guides(fill = guide_legend(title = NULL)) +
    theme(legend.title = element_blank())

  ##############################################################################
  ### Density Plots
  ##############################################################################

  densityplot <- ggplot(
    subset(D, information_sources == "Posterior"),
    aes_string(x = "x", y = "y")
  ) +
    geom_line(size = 2, aes_string(colour = "group")) +
    ylab("Density (PDF)") +
    xlab("Values") +
    theme_bw() +
    ggtitle("Density Plot") +
    guides(fill = guide_legend(title = NULL)) +
    theme(legend.title = element_blank())

  ##############################################################################
  ### Discount function plot
  ### - Only makes sense to plot if Current/historical treatment are present or
  ###   both current/historical control are present
  ##############################################################################

  p_hat <- seq(0, 1, length.out = 100)

  discountfun_plot <- NULL

  if (!is.null(N0_t) & !is.null(N_t)) {

    ### Create discount function based on distribution type
    if (discount_function == "weibull") {
      discount_function_treatment <- pweibull(p_hat,
                                              shape = x$args1$weibull_shape[1],
                                              scale = x$args1$weibull_scale[1]
      )
    } else if (discount_function == "scaledweibull") {
      max_p <- pweibull(1,
                        shape = x$args1$weibull_shape[1],
                        scale = x$args1$weibull_scale[1]
      )

      discount_function_treatment <- pweibull(p_hat,
                                              shape = x$args1$weibull_shape[1],
                                              scale = x$args1$weibull_scale[1]
      ) / max_p
    } else if (discount_function == "identity") {
      discount_function_treatment <- p_hat
    }

    D1 <- data.frame(
      group = "Treatment",
      y = discount_function_treatment,
      x = seq(0, 1, length.out = 100)
    )
    D2 <- data.frame(
      group = c("Treatment"),
      p_hat = median(posterior_treatment$p_hat)
    )
    D3 <- data.frame(
      group = c("Treatment"),
      p_hat = median(posterior_treatment$alpha_discount)
    )

    discountfun_plot <- ggplot() +
      geom_line(data = D1, aes_string(y = "y", x = "x", color = "group"), size = 1) +
      geom_vline(data = D2, aes_string(xintercept = "p_hat", color = "group"), lty = 2) +
      geom_hline(data = D3, aes_string(yintercept = "p_hat", color = "group"), lty = 2)
  }

  if (arm2) {
    if (!is.null(N0_c) & !is.null(N_c)) {
      if (is.null(discountfun_plot)) {
        discountfun_plot <- ggplot()
      }

      ### Create discount function based on distribution type
      if (discount_function == "weibull") {
        discount_function_control <- pweibull(p_hat,
                                              shape = x$args1$weibull_shape[2],
                                              scale = x$args1$weibull_scale[2]
        )
      } else if (discount_function == "scaledweibull") {
        max_p <- pweibull(1,
                          shape = x$args1$weibull_shape[2],
                          scale = x$args1$weibull_scale[2]
        )

        discount_function_control <- pweibull(p_hat,
                                              shape = x$args1$weibull_shape[2],
                                              scale = x$args1$weibull_scale[2]
        ) / max_p
      } else if (discount_function == "identity") {
        discount_function_control <- p_hat
      }

      D4 <- data.frame(
        group = "Control",
        y = discount_function_control,
        x = seq(0, 1, length.out = 100)
      )
      D5 <- data.frame(group = "Control", p_hat = median(posterior_control$p_hat))
      D6 <- data.frame(group = "Control", p_hat = median(posterior_control$alpha_discount))

      discountfun_plot <- discountfun_plot +
        geom_line(data = D4, aes_string(y = "y", x = "x", color = "group"), size = 1) +
        geom_vline(data = D5, aes_string(xintercept = "p_hat", color = "group"), lty = 2) +
        geom_hline(data = D6, aes_string(yintercept = "p_hat", color = "group"), lty = 2)
    }
  }

  if (!is.null(discountfun_plot)) {
    discountfun_plot <- discountfun_plot +
      facet_wrap(~group, ncol = 1) +
      theme_bw() +
      ylab("Alpha Discount Value") +
      xlab("Stochastic comparison (Current vs Historical Data)") +
      ggtitle("Discount Function") +
      ylim(0, 1) +
      guides(fill = guide_legend(title = NULL)) +
      theme(legend.title = element_blank())
  }

  if (print) {
    if (is.null(type)) {
      op <- par(ask = TRUE)
      plot(post_typeplot)
      plot(densityplot)
      if (!is.null(discountfun_plot)) {
        plot(discountfun_plot)
      }
      par(op)
    } else if (type == "discount") {
      if (!is.null(discountfun_plot)) {
        plot(discountfun_plot)
      }
    } else if (type == "posteriors") {
      plot(post_typeplot)
    } else if (type == "density") {
      plot(densityplot)
    }
  } else {
    if (is.null(type)) {
      out <- list(
        posteriors = post_typeplot,
        density = densityplot
      )
      if (!is.null(discountfun_plot)) {
        out <- c(out, discount = list(discountfun_plot))
        out
      } else {
        out
      }
    } else if (type == "discount") {
      if (!is.null(discountfun_plot)) {
        discountfun_plot
      }
    } else if (type == "posteriors") {
      post_typeplot
    } else if (type == "density") {
      densityplot
    }
  }
})


#' @title bdpbinomial Object Plot
#' @description \code{plot} method for class \code{bdpbinomial}.
#'
#' @param x object of class \code{bdpbinomial}. The result of a call to the
#'   \code{\link{bdpbinomial}} function.
#' @param type character. Optional. Select plot type to print. Supports the
#'   following: "discount" gives the discount function; "posteriors" gives the
#'   posterior plots of the event rates; and "density" gives the augmented
#'   posterior density plot(s). Leave NULL to display all plots in sequence.
#' @param print logical. Optional. Produce a plot (\code{print = TRUE}; default)
#'   or return a ggplot object (\code{print = FALSE}). When \code{print =
#'   FALSE}, it is possible to use \code{ggplot2} syntax to modify the plot
#'   appearance.
#'
#' @details The \code{plot} method for \code{bdpbinomial} returns up to three
#'   plots: (1) posterior density curves; (2) posterior density of the effect of
#'   interest; and (3) the discount function. A call to \code{plot} that omits
#'   the \code{type} input returns one plot at a time and prompts the user to
#'   click the plot or press return to proceed to the next plot. Otherwise, the
#'   user can specify a plot \code{type} to display the requested plot.
#'
#' @import methods
#' @importFrom utils head
#' @importFrom ggplot2 aes_string ggtitle ylim guides guide_legend theme
#'   element_blank
#' @importFrom graphics par
#' @importFrom stats density is.empty.model median model.offset model.response
#'   pweibull quantile rbeta rgamma rnorm var vcov
#' @export
setMethod("plot", signature(x = "bdpbinomial"), function(x, type = NULL, print = TRUE) {
  posterior_treatment <- x$posterior_treatment
  posterior_control <- x$posterior_control
  discount_function <- x$args1$discount_function
  N0_t <- x$args1$N0_t
  N0_c <- x$args1$N0_c
  N_t <- x$args1$N_t
  N_c <- x$args1$N_c
  arm2 <- x$args1$arm2
  method <- x$args1$method

  information_sources <- NULL

  ### Create data frames for plotting via ggplot2
  D1 <- D2 <- D3 <- D5 <- D6 <- NULL

  if (arm2) {
    dens1 <- density(posterior_control$posterior,
                     adjust = 0.5
    )
    D1 <- data.frame(
      information_sources = "Posterior",
      group = "Control",
      x = dens1$x,
      y = dens1$y
    )

    if (!is.null(posterior_control$posterior_flat)) {
      dens2 <- density(posterior_control$posterior_flat, adjust = 0.5)
      D2 <- data.frame(
        information_sources = "Current Data",
        group = "Control",
        x = dens2$x,
        y = dens2$y
      )
    }

    if (!is.null(posterior_control$prior)) {
      dens3 <- density(posterior_control$prior, adjust = 0.5)
      D3 <- data.frame(
        information_sources = "Historical Data",
        group = "Control",
        x = dens3$x,
        y = dens3$y
      )
    }
  }

  dens4 <- density(posterior_treatment$posterior, adjust = 0.5)
  D4 <- data.frame(
    information_sources = "Posterior",
    group = "Treatment",
    x = dens4$x,
    y = dens4$y
  )

  if (!is.null(posterior_treatment$posterior_flat)) {
    dens5 <- density(posterior_treatment$posterior_flat, adjust = 0.5)
    D5 <- data.frame(
      information_sources = "Current Data",
      group = "Treatment",
      x = dens5$x,
      y = dens5$y
    )
  }

  if (!is.null(posterior_treatment$prior)) {
    dens6 <- density(posterior_treatment$prior, adjust = 0.5)
    D6 <- data.frame(
      information_sources = "Historical Data",
      group = "Treatment",
      x = dens6$x,
      y = dens6$y
    )
  }

  D <- rbind(D1, D2, D3, D4, D5, D6)

  D$information_sources <- factor(D$information_sources,
                                  levels = (c("Posterior", "Current Data", "Historical Data"))
  )
  D$group <- factor(D$group, levels = c("Treatment", "Control"))

  ##############################################################################
  ### Posterior Type Plots
  ##############################################################################

  post_typeplot <- ggplot(D, aes_string(x = "x", y = "y")) +
    geom_line(size = 2, aes_string(color = "information_sources", lty = "information_sources")) +
    theme_bw() +
    facet_wrap(~group, ncol = 1, scales = "free") +
    ylab("Density (PDF)") +
    xlab("Values") +
    ggtitle("Posterior Type Plot") +
    guides(fill = guide_legend(title = NULL)) +
    theme(legend.title = element_blank())

  ##############################################################################
  ### Density Plots
  ##############################################################################

  densityplot <- ggplot(
    subset(D, information_sources == "Posterior"),
    aes_string(x = "x", y = "y")
  ) +
    geom_line(size = 2, aes_string(color = "group")) +
    ylab("Density (PDF)") +
    xlab("Values") +
    theme_bw() +
    ggtitle("Density Plot") +
    guides(fill = guide_legend(title = NULL)) +
    theme(legend.title = element_blank())

  ##############################################################################
  ### Discount function plot
  ### - Only makes sense to plot if Current/historical treatment are present or
  ###   both current/historical control are present
  ##############################################################################

  p_hat <- seq(0, 1, length.out = 100)

  discountfun_plot <- NULL

  if (!is.null(N0_t) & !is.null(N_t)) {
    ### Create discount function based on distribution type
    if (discount_function == "weibull") {
      discount_function_treatment <- pweibull(p_hat,
                                              shape = x$args1$weibull_shape[1],
                                              scale = x$args1$weibull_scale[1]
      )
    } else if (discount_function == "scaledweibull") {
      max_p <- pweibull(1,
                        shape = x$args1$weibull_shape[1],
                        scale = x$args1$weibull_scale[1]
      )

      discount_function_treatment <- pweibull(p_hat,
                                              shape = x$args1$weibull_shape[1],
                                              scale = x$args1$weibull_scale[1]
      ) / max_p
    } else if (discount_function == "identity") {
      discount_function_treatment <- p_hat
    }

    D1 <- data.frame(
      group = "Treatment",
      y = discount_function_treatment,
      x = seq(0, 1, length.out = 100)
    )
    D2 <- data.frame(group = "Treatment", p_hat = median(posterior_treatment$p_hat))
    D3 <- data.frame(group = "Treatment", p_hat = median(posterior_treatment$alpha_discount))

    discountfun_plot <- ggplot() +
      geom_line(data = D1, aes_string(y = "y", x = "x", color = "group"), size = 1) +
      geom_vline(data = D2, aes_string(xintercept = "p_hat", color = "group"), lty = 2) +
      geom_hline(data = D3, aes_string(yintercept = "p_hat", color = "group"), lty = 2)
  }

  if (arm2) {
    if (!is.null(N0_c) & !is.null(N_c)) {
      if (is.null(discountfun_plot)) {
        discountfun_plot <- ggplot()
      }

      ### Create discount function based on distribution type
      if (discount_function == "weibull") {
        discount_function_control <- pweibull(p_hat,
                                              shape = x$args1$weibull_shape[2],
                                              scale = x$args1$weibull_scale[2]
        )
      } else if (discount_function == "scaledweibull") {
        max_p <- pweibull(1,
                          shape = x$args1$weibull_shape[2],
                          scale = x$args1$weibull_scale[2]
        )

        discount_function_control <- pweibull(p_hat,
                                              shape = x$args1$weibull_shape[2],
                                              scale = x$args1$weibull_scale[2]
        ) / max_p
      } else if (discount_function == "identity") {
        discount_function_control <- p_hat
      }

      D4 <- data.frame(
        group = "Control",
        y = discount_function_control,
        x = seq(0, 1, length.out = 100)
      )
      D5 <- data.frame(group = "Control", p_hat = median(posterior_control$p_hat))
      D6 <- data.frame(group = "Control", p_hat = median(posterior_control$alpha_discount))

      discountfun_plot <- discountfun_plot +
        geom_line(data = D4, aes_string(y = "y", x = "x", color = "group"), size = 1) +
        geom_vline(data = D5, aes_string(xintercept = "p_hat", color = "group"), lty = 2) +
        geom_hline(data = D6, aes_string(yintercept = "p_hat", color = "group"), lty = 2)
    }
  }

  if (!is.null(discountfun_plot)) {
    discountfun_plot <- discountfun_plot +
      facet_wrap(~group, ncol = 1) +
      theme_bw() +
      ylab("Alpha Discount Value") +
      xlab("Stochastic comparison (Current vs Historical Data)") +
      ggtitle("Discount Function") +
      ylim(0, 1) +
      guides(fill = guide_legend(title = NULL)) +
      theme(legend.title = element_blank())
  }

  if (print) {
    if (is.null(type)) {
      op <- par(ask = TRUE)
      plot(post_typeplot)
      plot(densityplot)
      if (!is.null(discountfun_plot)) {
        plot(discountfun_plot)
      }
      par(op)
    } else if (type == "discount") {
      if (!is.null(discountfun_plot)) {
        plot(discountfun_plot)
      }
    } else if (type == "posteriors") {
      plot(post_typeplot)
    } else if (type == "density") {
      plot(densityplot)
    }
  } else {
    if (is.null(type)) {
      out <- list(
        posteriors = post_typeplot,
        density = densityplot
      )
      if (!is.null(discountfun_plot)) {
        out <- c(out, discount = list(discountfun_plot))
        out
      } else {
        out
      }
    } else if (type == "discount") {
      if (!is.null(discountfun_plot)) {
        discountfun_plot
      }
    } else if (type == "posteriors") {
      post_typeplot
    } else if (type == "density") {
      densityplot
    }
  }
})


#' @title bdpsurvival Object Plot
#' @description \code{plot} method for class \code{bdpsurvival}.
#'
#' @param x object of class \code{bdpsurvival}. The result of a call to the
#'   \code{\link{bdpsurvival}} function.
#' @param type character. Optional. Select plot type to print. Supports the
#'   following: "discount" gives the discount function and "survival" gives the
#'   survival curves. Leave NULL to display all plots in sequence.
#' @param print logical. Optional. Produce a plot (\code{print = TRUE}; default)
#'   or return a ggplot object (\code{print = FALSE}). When \code{print =
#'   FALSE}, it is possible to use \code{ggplot2} syntax to modify the plot
#'   appearance.
#'
#' @details The \code{plot} method for \code{bdpsurvival} returns up to two
#'   plots: (1) posterior survival curves and (2) the discount function. A call
#'   to \code{plot} that omits the \code{type} input returns one plot at a time
#'   and prompts the user to click the plot or press return to proceed to the
#'   next plot. Otherwise, the user can specify a plot \code{type} to display
#'   the requested plot.
#'
#' @import methods
#' @importFrom utils head
#' @importFrom ggplot2 aes_string ggtitle ylim guides guide_legend theme
#'   element_blank
#' @importFrom graphics par
#' @importFrom stats density is.empty.model median model.offset model.response
#'   pweibull quantile rbeta rgamma rnorm var vcov
#' @export
setMethod("plot", signature(x = "bdpsurvival"), function(x, type = NULL, print = TRUE) {
  args1 <- x$args1
  discount_function <- x$args1$discount_function
  posterior_treatment <- x$posterior_treatment
  posterior_control <- x$posterior_control
  data <- args1$data
  breaks <- args1$breaks
  arm2 <- args1$arm2
  method <- args1$method

  ##############################################################################
  ### Survival curve(s)
  ### - Only computed for one-arm trial
  ##############################################################################

  ### Organize data for current treatment
  time_t <- sort(unique(args1$S_t$time))
  survival_times_posterior_flat <- lapply(time_t, ppexp,
                                          posterior_treatment$posterior_flat_hazard,
                                          cuts = c(0, breaks)
  )
  survival_median_posterior_flat <- 1 - sapply(survival_times_posterior_flat, median)

  D1 <- data.frame(
    source = "Current Data",
    group = "Treatment",
    x = time_t,
    y = survival_median_posterior_flat
  )

  ### Organize data for historical treatment
  if (!is.null(args1$S0_t)) {
    time0_t <- sort(unique(args1$S0_t$time))
    survival_times_prior <- lapply(time0_t, ppexp,
                                   posterior_treatment$prior_hazard,
                                   cuts = c(0, breaks)
    )
    survival_median_prior <- 1 - sapply(survival_times_prior, median)

    D2 <- data.frame(
      source = "Historical Data",
      group = "Treatment",
      x = time0_t,
      y = survival_median_prior
    )
  } else {
    D2 <- NULL
  }

  ### Organize data for treatment posterior
  survival_times_posterior <- lapply(time_t, ppexp, posterior_treatment$posterior_hazard, cuts = c(0, breaks))
  survival_median_posterior <- 1 - sapply(survival_times_posterior, median)

  D3 <- data.frame(
    source = "Posterior",
    group = "Treatment",
    x = time_t,
    y = survival_median_posterior
  )

  ### Organize data for current control
  if (!is.null(args1$S_c)) {
    time_c <- sort(unique(args1$S_c$time))
    survival_times_posterior_flat <- lapply(time_c, ppexp,
                                            posterior_control$posterior_flat_hazard,
                                            cuts = c(0, breaks)
    )
    survival_median_posterior_flat <- 1 - sapply(survival_times_posterior_flat, median)

    D4 <- data.frame(
      source = "Current Data",
      group = "Control",
      x = time_c,
      y = survival_median_posterior_flat
    )
  } else {
    D4 <- NULL
  }

  ### Organize data for historical control
  if (!is.null(args1$S0_c)) {
    time0_c <- sort(unique(args1$S0_c$time))
    survival_times_prior <- lapply(time0_c, ppexp,
                                   posterior_control$prior_hazard,
                                   cuts = c(0, breaks)
    )
    survival_median_prior <- 1 - sapply(survival_times_prior, median)

    D5 <- data.frame(
      source = "Historical Data",
      group = "Control",
      x = time0_c,
      y = survival_median_prior
    )
  } else {
    D5 <- NULL
  }

  ### Organize data for control posterior
  if (!is.null(args1$S_c) & !is.null(args1$S0_c)) {
    survival_times_posterior <- lapply(time_c, ppexp, posterior_control$posterior_hazard, cuts = c(0, breaks))
    survival_median_posterior <- 1 - sapply(survival_times_posterior, median)

    D6 <- data.frame(
      source = "Posterior",
      group = "Control",
      x = time_c,
      y = survival_median_posterior
    )
  } else if (is.null(args1$S_c) & !is.null(args1$S0_c)) {
    survival_times_posterior <- lapply(time0_c, ppexp, posterior_control$posterior_hazard, cuts = c(0, breaks))
    survival_median_posterior <- 1 - sapply(survival_times_posterior, median)

    D6 <- data.frame(
      source = "Posterior",
      group = "Control",
      x = time0_c,
      y = survival_median_posterior
    )
  } else if (!is.null(args1$S_c) & is.null(args1$S0_c)) {
    survival_times_posterior <- lapply(time_c, ppexp, posterior_control$posterior_hazard, cuts = c(0, breaks))
    survival_median_posterior <- 1 - sapply(survival_times_posterior, median)

    D6 <- data.frame(
      source = "Posterior",
      group = "Control",
      x = time_c,
      y = survival_median_posterior
    )
  } else {
    D6 <- NULL
  }

  D <- rbind(D4, D5, D6, D1, D2, D3)
  D$group <- factor(D$group, levels = c("Treatment", "Control"))

  ### Plot survival curve
  survival_curves <- ggplot(D, aes_string(x = "x", y = "y")) +
    geom_line(size = 1.4, aes_string(color = "source", lty = "source")) +
    facet_wrap(~group, ncol = 1, scales = "free") +
    ylab("Survival probability") +
    xlab("Time") +
    theme_bw() +
    ggtitle("Survival Curve(s)") +
    guides(fill = guide_legend(title = NULL)) +
    theme(legend.title = element_blank())

  ##############################################################################
  ### Discount function plot
  ### - Only makes sense to plot if Current/historical treatment are present or
  ###   both current/historical control are present
  ##############################################################################

  p_hat <- seq(0, 1, length.out = 100)

  discountfun_plot <- NULL

  if (!is.null(args1$S_t) & !is.null(args1$S0_t)) {

    ### Create discount function based on distribution type
    if (discount_function == "weibull") {
      discount_function_treatment <- pweibull(p_hat,
                                              shape = x$args1$weibull_shape[1],
                                              scale = x$args1$weibull_scale[1]
      )
    } else if (discount_function == "scaledweibull") {
      max_p <- pweibull(1,
                        shape = x$args1$weibull_shape[1],
                        scale = x$args1$weibull_scale[1]
      )

      discount_function_treatment <- pweibull(p_hat,
                                              shape = x$args1$weibull_shape[1],
                                              scale = x$args1$weibull_scale[1]
      ) / max_p
    } else if (discount_function == "identity") {
      discount_function_treatment <- p_hat
    }

    D1 <- data.frame(
      group = "Treatment",
      y = discount_function_treatment,
      x = seq(0, 1, length.out = 100)
    )
    D2 <- data.frame(group = "Treatment", p_hat = c(median(posterior_treatment$p_hat)))
    D3 <- data.frame(group = "Treatment", p_hat = c(median(posterior_treatment$alpha_discount)))

    discountfun_plot <- ggplot() +
      geom_line(data = D1, aes_string(y = "y", x = "x", color = "group"), size = 1) +
      geom_vline(data = D2, aes_string(xintercept = "p_hat", color = "group"), lty = 2) +
      geom_hline(data = D3, aes_string(yintercept = "p_hat", color = "group"), lty = 2)
  }

  if (arm2) {
    if (!is.null(args1$S_c) & !is.null(args1$S0_c)) {
      if (is.null(discountfun_plot)) {
        discountfun_plot <- ggplot()
      }

      ### Create discount function based on distribution type
      if (discount_function == "weibull") {
        discount_function_control <- pweibull(p_hat,
                                              shape = x$args1$weibull_shape[2],
                                              scale = x$args1$weibull_scale[2]
        )
      } else if (discount_function == "scaledweibull") {
        max_p <- pweibull(1,
                          shape = x$args1$weibull_shape[2],
                          scale = x$args1$weibull_scale[2]
        )

        discount_function_control <- pweibull(p_hat,
                                              shape = x$args1$weibull_shape[2],
                                              scale = x$args1$weibull_scale[2]
        ) / max_p
      } else if (discount_function == "identity") {
        discount_function_control <- p_hat
      }

      D4 <- data.frame(
        group = "Control",
        y = discount_function_control,
        x = seq(0, 1, length.out = 100)
      )
      D5 <- data.frame(group = "Control", p_hat = c(median(posterior_control$p_hat)))
      D6 <- data.frame(group = "Control", p_hat = c(median(posterior_control$alpha_discount)))

      discountfun_plot <- discountfun_plot +
        geom_line(data = D4, aes_string(y = "y", x = "x", color = "group"), size = 1) +
        geom_vline(data = D5, aes_string(xintercept = "p_hat", color = "group"), lty = 2) +
        geom_hline(data = D6, aes_string(yintercept = "p_hat", color = "group"), lty = 2)
    }
  }

  if (!is.null(discountfun_plot)) {
    discountfun_plot <- discountfun_plot +
      facet_wrap(~group, ncol = 1) +
      theme_bw() +
      ylab("Alpha Discount Value") +
      xlab("Stochastic comparison (Current vs Historical Data)") +
      ggtitle("Discount Function") +
      ylim(0, 1) +
      guides(fill = guide_legend(title = NULL)) +
      theme(legend.title = element_blank())
  }

  if (print) {
    if (is.null(type)) {
      op <- par(ask = TRUE)
      plot(survival_curves)
      if (!is.null(args1$S0_t) | (!is.null(args1$S_c) & !is.null(args1$S0_c))) {
        plot(discountfun_plot)
      }
      par(op)
    } else if (type == "discount") {
      if (!is.null(args1$S0_t) | (!is.null(args1$S_c) & !is.null(args1$S0_c))) {
        plot(discountfun_plot)
      }
    } else if (type == "survival") {
      plot(survival_curves)
    }
  } else {
    if (is.null(type)) {
      out <- list(survival = survival_curves)
      if (!is.null(args1$S0_t) | (!is.null(args1$S_c) & !is.null(args1$S0_c))) {
        out <- c(out, discount = list(discountfun_plot))
        out
      } else {
        out
      }
    } else if (type == "discount") {
      if (!is.null(args1$S0_t) | (!is.null(args1$S_c) & !is.null(args1$S0_c))) {
        discountfun_plot
      }
    } else if (type == "survival") {
      survival_curves
    }
  }
})