R/autoplot_pfun.R

In confMeta: Confidence Curves and P-Value Functions for Meta-Analysis

#' @title Visualizations of \code{confMeta} objects
#'
#' @description Plots one or more \code{confMeta} objects. This function can
#'     create two types of plots: the \emph{p}-value function plot (also known
#'     as drapery plot) and the forest plot. This allows for direct visual
#'     comparison of different \emph{p}-value functions.
#'
#' Optionally, a Bayesian meta-analysis object created with the
#' \code{bayesmeta::bayesmeta()} function can be supplied via the \code{bayesmeta} argument.
#' When provided, its posterior summary is displayed as an additional diamond
#' at the bottom of the forest plot for comparison with the frequentist methods
#' 
#' \strong{Important Note:} If supplying a Bayesian meta-analysis object 
#' via the \code{bayesmeta} argument, this function explicitly extracts the 
#' 95% credible intervals. If the Bayesian model was fit using a different 
#' credible level, the function will crash.
#' 
#' @param ... One or more objects of class \code{confMeta}.
#' @param type A character vector of length 1 or 2. Indicates what type of
#'     plot should be returned. Accepted values are \code{"p"}, \code{"forest"},
#'     or both. Defaults to \code{c("p", "forest")}.
#' @param diamond_height Numeric scalar. Indicates the maximal
#'     possible height of the diamonds in the forest plot. Defaults to 0.5.
#'     This argument is only relevant if \code{type} contains \code{"forest"} and will
#'     be ignored otherwise.
#' @param v_space Numeric scalar. Indicates the vertical space
#'     between two diamonds in the forest plot. Defaults to 1.5.
#'     This argument is only relevant if \code{type} contains \code{"forest"} and will
#'     be ignored otherwise.
#' @param scale_diamonds Logical. If \code{TRUE} (default), the diamond is rescaled to the
#'      interval \[0, 1\] in cases where the maximum of the p-value
#'     function is not equal to 1. This argument is only relevant if \code{type}
#'     contains \code{"forest"} and will be ignored otherwise.
#' @param show_studies Logical. If \code{TRUE} (default), the forest plot shows the
#'     confidence intervals for the individual effect estimates. Otherwise,
#'     the intervals are suppressed. This argument is only relevant if \code{type}
#'     contains \code{"forest"} and will be ignored otherwise.
#' @param drapery Logical. If \code{TRUE} (default), individual
#'     study effects are represented as drapery plots in the \emph{p}-value function plot. 
#'     If \code{FALSE} the studies are represented by a simple vertical line at their effect estimates.
#'      This argument is only relevant if \code{type}
#'     contains \code{"p"} and will be ignored otherwise.
#' @param reference_methods_p Character vector controlling which reference
#'     meta-analysis methods are shown as baseline curves in the
#'     \emph{p}-value function plot. Valid options are \code{"fe"},
#'     \code{"re"}, or \code{c("fe", "re")} for both. Defaults to \code{"re"}.
#' @param reference_methods_forest Character vector of length 1 to 4.
#'     Specifies which reference meta-analysis methods should be shown as
#'     diamonds in the forest plot. Valid options are any subset of
#'     \code{c("fe", "re", "hk", "hc")}, which correspond to:
#'     \itemize{
#'       \item \code{"fe"}: fixed-effect meta-analysis
#'       \item \code{"re"}: random-effects meta-analysis
#'       \item \code{"hk"}: Hartung-Knapp adjustment
#'       \item \code{"hc"}: Henmi-Copas adjustment
#'     }
#'     Defaults to \code{c("re", "hk", "hc")}.
#' @param ref_labels A named character vector to customize the labels of the reference 
#'     methods in the plot legends and axes (e.g., \code{c("fe" = "Fixed-effect (IV)", 
#'     "re" = "Random-effects (DL)")}). Defaults to \code{NULL}.
#' @param xlim Numeric vector of length 2. Global limits for the x-axis. 
#'   If \code{NULL}, limits are calculated automatically.
#' @param xlim_p Numeric vector of length 2. Specific x-axis limits for the 
#'   \emph{p}-value plot. Overrides \code{xlim}.
#' @param xlim_forest Numeric vector of length 2. Specific x-axis limits for 
#'   the forest plot. Overrides \code{xlim}.
#' @param same_xlim Logical. If \code{TRUE}, forces the \emph{p}-value plot to use the 
#'   same x-axis limits as the forest plot. Defaults to \code{TRUE}.
#' @param xlab Character string. Label for the x-axis. Defaults to \code{NULL} 
#'   (renders as \eqn{\mu}).
#' @param n_breaks Numeric. Approximate number of tick marks to display 
#'   on the x-axis. Defaults to 7.
#' @param bayesmeta An object of class \code{bayesmeta}, typically created using 
#'     \code{bayesmeta::bayesmeta()}. When provided, the posterior median (or mean) and
#'     95% credible interval are displayed as an additional diamond at the bottom 
#'     of the forest plot. Has no effect on the \emph{p}-value plot. Defaults to \code{NULL}.
#' @param n_points Numeric. Number of points used to create the \emph{p}-value plot. 
#'     Higher values (e.g., 10000) yield higher resolution but take longer to render. 
#'     Defaults to 1000.
#' @return An object of class \code{ggplot} containing the specified plot(s).
#'
#'
#'
#'
#' @importFrom patchwork wrap_plots
#'
#' @export
autoplot.confMeta <- function(
    ...,
    type = c("p", "forest"),
    diamond_height = 0.5,
    v_space = 1.5,
    scale_diamonds = TRUE,
    show_studies = TRUE,
    drapery = TRUE,
    # drapery_ma = c("fe"),
    reference_methods_p      = c("re"),
    reference_methods_forest = c("re", "hk", "hc"),
    ref_labels = NULL,
    xlim = NULL,
    xlim_p = NULL,
    xlim_forest = NULL,
    same_xlim = TRUE,
    xlab = NULL,
    n_breaks = 7,
    bayesmeta = NULL,
    n_points = 1000
) {

    # Check validity of reference methods
  check_ref_methods(reference_methods = reference_methods_forest)
  if (!all(reference_methods_p %in% c("fe", "re"))) {
    stop("`reference_methods_p` can only contain \"fe\", \"re\", or both.")
  }
  type <- match.arg(type, several.ok = TRUE)
  reference_methods_forest <- match.arg(
    reference_methods_forest,
    several.ok = TRUE,
    choices = c("fe", "re", "hk", "hc")
  )
  reference_methods_p <- match.arg(
    reference_methods_p,
    several.ok = TRUE,
    choices = c("fe", "re")
  )
  
  
  # Check for ref_labels
  valid_method_keys <- c("fe", "re", "hk", "hc")
  
  if (!is.null(ref_labels)) { 
    if (!is.character(ref_labels) || is.null (names(ref_labels))) {
      stop("'ref_labels' must be a named character vector, e.g. c(fe = 'Fixed-effect (IV)', re = 'Random-effects (REML)')")
    }
    invalid_keys <- setdiff(names(ref_labels), valid_method_keys)
    if (length(invalid_keys) > 0) {
      stop(paste0("Invalid names in `ref_labels`: ", paste(invalid_keys, collapse = ", "),
                  ". Must be one of: ", paste(valid_method_keys, collapse = ", "), "."))
    }
  }
  

    # Check all the confMeta objects
    cms <- list(...)
    check_cms(cms = cms)
    check_unique_names(cms = cms)

    # Check that the levels, estimates, SEs are equal
    # in all confMeta objects
    ests <- lapply(cms, "[[", i = "estimates")
    ses <- lapply(cms, "[[", i = "SEs")
    lvl <- lapply(cms, "[[", i = "conf_level")
    nms <- lapply(cms, "[[", i = "study_names")
    ok <- c(
        check_equal(ests),
        check_equal(ses),
        check_equal(lvl),
        check_equal(nms)
    )
    if (any(!ok)) {
        stop(
            paste0(
                "For plotting, all confMeta objects must have the same ",
                "'estimates', 'SEs', 'study_names', and 'conf_level' elements."
            )
        )
    }

    # Add some more checks for other graphics parameters
    check_TF(x = scale_diamonds)
    check_TF(x = show_studies)
    check_TF(x = drapery)
    check_length_1(x = diamond_height)
    check_class(x = diamond_height, "numeric")
    check_length_1(x = v_space)
    check_class(x = v_space, "numeric")

    # if xlim not specified, builds a range that covers all CIs and add 5% of margin
    # I added the possibility to handle xlim_p and xlim_forest divvertently!
    
    # COLOR HANDLING
    fun_names_shared  <- vapply(cms, "[[", i = "fun_name", character(1L)) #take the fun_name from every confMeta object (Edgington, Edgington_w ...)
    ref_base_p        <- intersect(c("fe", "re"), reference_methods_p)
    ref_forest_only   <- setdiff(
      map_ref_methods(reference_methods_forest, labels = ref_labels),
      map_ref_methods(ref_base_p, labels = ref_labels)
    )
    fac_levels_all    <- c(map_ref_methods(ref_base_p, labels = ref_labels), ref_forest_only, fun_names_shared)
    shared_colors     <- c("#000000", scales::hue_pal()(length(fac_levels_all) - 1))
    names(shared_colors) <- fac_levels_all #attach each method name to the color vector
    
    #This is the complete list of everything that needs a color 
    #we want one color x method -> first is black, then hue_pal() generates evenly sapced color
    
    
    # xlim for p plot:
    if (!is.null(xlim_p)) {
        check_xlim(x = xlim_p)
    } else if (!is.null(xlim)) { #goes to the global xlim if specified
      xlim_p <- xlim
      check_xlim(xlim_p)
    } else {
      candidates_p <- unname(
            do.call(
                "c",
                lapply(cms, function(x) {
                  
                  comp <- x$comparison_cis
                  
                  # For the p value plot, we only care about FE or RE. So we don't consider the other comparison method
                  keep <- rownames(comp) %in% c("Fixed-effect", "Random-effects", "fe", "re")
                  comp_filtered <- comp[keep, , drop = FALSE]
                  
                  with(x, c(individual_cis, joint_cis, comp_filtered))
                    }
                )))
        
        candidates_p <- candidates_p[is.finite(candidates_p)]
        
        ext_perc <- 5
        lower_p <- min(candidates_p)
        upper_p <- max(candidates_p)
        margin_p <- (upper_p - lower_p) * ext_perc / 100
        xlim_p <- c(lower_p - margin_p, upper_p + margin_p)
    }
    
    
    # xlim for f plot:
    if (!is.null(xlim_forest)) {
      
      check_xlim(xlim_forest) 
    } else if (!is.null(xlim)) {
      xlim_forest <- xlim     
      check_xlim(xlim_forest)
      
    } else {
      candidates_f <- unname(do.call("c", lapply(cms, function(x) {
        
        comp <- x$comparison_cis
        
        # create x axis only on the actual ref methods
        allowed_methods <- map_ref_methods (reference_methods_forest)
        
        keep <- rownames(comp)%in%allowed_methods
        comp_filtered <- comp[keep, , drop = FALSE]
        
        
        with(x, c(individual_cis, joint_cis, comp_filtered))
      })))
      
      candidates_f <- candidates_f[is.finite(candidates_f)]
      
      if (length(candidates_f) == 0) {
        xlim_forest <- c(-1, 1)
      } else {
        diff_f <- max(candidates_f) - min(candidates_f)
        margin_f <- if(diff_f == 0) 0.5 else diff_f * 0.05
        xlim_forest <- c(min(candidates_f) - margin_f, max(candidates_f) + margin_f)
      }
    }
    
    
    # if you want to have same x lim, set the x lim_p (which can be different since we excluded some methods) = lim_f
    if (same_xlim) {
      xlim_p = xlim_forest
    }

    plots <- list()
    if ("p" %in% type) {
      plots[["p"]] <- ggPvalueFunction(
        cms=cms, 
        drapery = drapery, 
        xlim = xlim_p,
        xlab = xlab,
        reference_methods = reference_methods_p, 
        n_breaks =  n_breaks,
        n_points = n_points,
        shared_colors = shared_colors,
        ref_labels = ref_labels
        
        
      )
    }
    
    if ("forest" %in% type) {
      plots[["forest"]] <- ForestPlot(
        cms = cms, 
        diamond_height = diamond_height, 
        v_space = v_space, 
        xlim = xlim_forest, 
        show_studies = show_studies, 
        scale_diamonds = scale_diamonds, 
        reference_methods = reference_methods_forest, 
        xlab = xlab, 
        n_breaks =  n_breaks,
        shared_colors = shared_colors,
        ref_labels = ref_labels
        
        
      )
    }
    
    
    # Add Bayesmeta if requested
    if (!is.null(bayesmeta)) {
      if (!inherits(bayesmeta, "bayesmeta")) stop("bayesmeta argument MUST be of class 'bayesmeta'")
           if ("forest" %in% type) {
             plots[["forest"]] <- add_bayes_forest(
               p = plots[["forest"]], 
               bm = bayesmeta, 
               diamond_height = diamond_height,
               v_space = v_space,               
               color = "black", 
               label = "Bayesmeta"
             )
                }
    }
    
    # OUTPUT
    if (length(plots) == 0) {
      return(NULL)
    } else if (length(plots) == 1L) {
      return(plots[[1L]])
    } else {
      return(patchwork::wrap_plots(plots, ncol = 1L)) #put them vertically if there are 2 plots
    }
}
    

# ==============================================================================
# Reexport autoplot function
# ==============================================================================

#' @importFrom ggplot2 autoplot
#' @export
ggplot2::autoplot

# ==============================================================================
# Input checks
# ==============================================================================

check_cms <- function(cms) {

    l_cms <- length(cms)
    counter <- 0L
    message <- NA_character_

    while (is.na(message) && counter < l_cms) {
        counter <- counter + 1L
        y <- tryCatch(
            {
                validate_confMeta(cms[[counter]])
            },
            error = function(e) conditionMessage(e)
        )
        if (!is.null(y)) message <- y
    }

    if (!is.na(message)) {
        stop(
            paste0(
                "Problem found in confMeta object number ",
                counter,
                ": ",
                message
            ),
            call. = FALSE
        )
    }
    invisible(NULL)
}

check_unique_names <- function(cms) {
    nms <- vapply(cms, "[[", i = "fun_name", character(1L))
    if (any(duplicated(nms))) {
        stop(
            "confMeta objects must have unique `fun_name`.",
            call. = FALSE
        )
    }
    invisible(NULL)
}

check_equal <- function(x) {

    # Order-agnostic
    x <- lapply(x, sort, decreasing = FALSE)
    # Get the first for comparison
    comp <- x[[1L]]
    all(
        do.call(
            "c",
            lapply(
                x,
                function(x, comp) {
                    all(x == comp)
                },
                comp = comp
            )
        )
    )
}

check_TF <- function(x) {
    obj <- deparse1(substitute(x))
    if (!is_TF(x)) {
        stop(
            paste0(
                "Object ", obj, " must be either `TRUE` or `FALSE`."
            ),
            call. = FALSE
        )
    }
    invisible(FALSE)
}

is_TF <- function(x) {
    if (length(x) == 1L && is.logical(x) && !is.na(x)) {
        TRUE
    } else {
        FALSE
    }
}

check_xlim <- function(x) {
    ok <- length(x) == 2L && is.numeric(x) && x[1L] < x[2L]
    if (!ok) {
        stop(
            paste0(
                "Argument `xlim` must be a numeric vector of length 2 where ",
                "the first element is smaller than the second."
            ),
            call. = FALSE
        )
    }
    invisible(NULL)
}


check_ref_methods <- function(reference_methods) {
    # Check for invalid reference_methods
    valid_methods <- c("fe", "re", "hk", "hc")
    is_invalid <- !(reference_methods %in% valid_methods)
    if (any(is_invalid)) {
        msg <- paste0(
            "Detected invalid reference_methods: ",
            paste0(reference_methods[is_invalid], collapse = ", "),
            "."
        )
        stop(msg)
    }
}

# ==============================================================================
# p-value function plot
# ==============================================================================

#' @importFrom stats qnorm
#' @importFrom ggplot2 ggplot aes xlim geom_line geom_hline geom_vline
#' @importFrom ggplot2 geom_point scale_y_continuous sec_axis geom_segment
#' @importFrom ggplot2 theme theme_minimal labs element_text element_blank
#' @importFrom ggplot2 xlim scale_color_manual
#' @importFrom scales hue_pal
#' @importFrom stats setNames
#' @noRd
ggPvalueFunction <- function(
    cms,
    xlim,
    drapery,
    xlab,
    reference_methods,
     n_breaks, 
    n_points,
    shared_colors,
    ref_labels
) {
 
  # Set some constants that are equal for all grid rows
  #use const as a list for containing stuff, muSeq is for creating the grid of mu values,

  const <- list(
    estimates  = cms[[1L]]$estimates,
    SEs        = cms[[1L]]$SEs,
    conf_level = cms[[1L]]$conf_level,
    eps = 0.0025,
    eb_height = 0.025,
    muSeq = seq(xlim[1], xlim[2], length.out = n_points)
  )
  
  

    # Get the function names (for legend)
    fun_names <- vapply(cms, "[[", i = "fun_name", character(1L))
    
    ref_base <- intersect(c("fe", "re"), reference_methods)
    fac_levels <- c(map_ref_methods(ref_base, labels = ref_labels), fun_names)
    
    # Calculate the p-values and CIs
    data <- lapply(seq_along(cms), function(x) {
        
        cm <- cms[[x]]
        fun <- cm$p_fun
        fun_name <- cm$fun_name
        alpha <- 1 - const$conf_level
        w <- if (!is.null(cm$w)) cm$w else NULL
        
        
        # use the new call for the function 
        pval <- call_pfun(
          fun = fun,
          estimates = cm$estimates,
          SEs = cm$SEs,
          mu = const$muSeq,
          w = w
        )
        
        
        ################################################################################################################
        #         ######## NOTE FOR WHO CARES 
        # (!!) here we call just those arguments, cause the others are already saved inside p_fun as default parameters 
        #thanks to make_p_fun() (e.g. heterogeneity, approx, tau2, neff...)
        #
        #just remember that w in this closure is by default (1,...1) so this is extremely necessary
        #############################################################################################à##################
        
        
        
        #now we have a vector p val of length 10.000 (one for each mu) 
        
        
        pval <- pmin(pmax(pval, 0), 1) #sometimes due to approx error are a tiny bit smaller than 0, fix to not have any error messages
        
        CIs <- cm$joint_cis
        y0 <- cm$p_0[, 2L] #p_val at mu = 0

        # Data frame with the lines of the p-value function: (x,y coords, value
        # at x = 0)
        df1 <- data.frame(
            x = const$muSeq,
            y = pval,
            p_val_fun = cm$fun_name,
            y0 = y0, # p value at mu = 0
            group = factor(fun_name, levels = fac_levels), #to plot different lines each method
            stringsAsFactors = FALSE,
            row.names = NULL
        )
        # make a second data frame for the display of confidence intervals
        df2 <- data.frame(
            xmin = CIs[, 1L],
            xmax = CIs[, 2L],
            p_val_fun = fun_name,
            # y = rep(1 - conf_level, nrow(CIs$CI)) + factor * const$eps,
            y = rep(alpha, nrow(CIs)), #horiz segment at confidence level
            group = factor(fun_name, levels = fac_levels),
            stringsAsFactors = FALSE,
            row.names = NULL
        )
        #define the height of the small vertical tiks at the extreme of CI:
        #  |
        #--|---------|--
        #  |
        df2$ymax <- df2$y + const$eb_height 
        df2$ymin <- df2$y - const$eb_height
        list(df1, df2)
    })
    
    #now we have the data object (e.g. data[[1]] = list(df1, df2) is the first ConfMeta object
    # data[[2]] = list(df1, df2) second ConfMeta obj etc...)
    
    # Extract the data frame for the lines with p-value functions
    # as well as the data frame for the error bars
    plot_data <- lapply(
        list(lines = 1L, errorbars = 2L),
        function(z, data) do.call("rbind", lapply(data, "[[", i = z)),
        data = data
    )
    lines <- plot_data[["lines"]]
    errorbars <- plot_data[["errorbars"]]

    
    
    #if you print errorbars:
    #xmin     xmax   y   ymax   ymin
    # 1   -2.4671   1.9013 0.05 0.075 0.025
    # means that the graph will be drawed: an horiz segm from mu=-2.47 to mu=1.9, positioned at level 
    # p = 0.05, with two vertical ticks from 0.025 to 0.075

    # Calculate the drapery lines if TRUE (by default)
    if (drapery) {
        dp <- get_drapery_df(
            estimates = const$estimates,
            SEs = const$SEs,
            mu = const$muSeq
        )
        ## also compute the RMA p-value function
        ### Object containing CIs for the reference methods
        ci_obj <- cms[[1L]]$comparison_cis
        cl <- cms[[1L]]$conf_level
        
        ## What reference method should we calculate drapery plot for:
        ## "fe" or "re"
        ref_indices <- which(rownames(ci_obj) %in% map_ref_methods(intersect(c("fe", "re"), reference_methods)))
        
        default_to_custom_p <- setNames(
          map_ref_methods(c("fe", "re", "hk", "hc"), labels = ref_labels),
          map_ref_methods(c("fe", "re", "hk", "hc"))
        )
        
        rmadf <- do.call("rbind", lapply(ref_indices, function(idx) {
          rmaestimate <- rowMeans(ci_obj[idx, , drop = FALSE])
          rmase <- (ci_obj[idx, 2L] - ci_obj[idx, 1L]) /
            (2 * stats::qnorm(p = (1 + cl) / 2))
          df <- get_drapery_df(estimates = rmaestimate, SEs = rmase, mu = const$muSeq)
          df$group <- default_to_custom_p[rownames(ci_obj)[idx]]  # label by method name, maybe I should wrap it in unname? Not sure, future development
          df
        }))
        rmadf$group <- factor(rmadf$group, levels = fac_levels)
    }
    
    
    #now dp and rmdaf are two similar df, one for my method, the other for the basline method  

    # Define function to convert breaks from primary y-axis to
    # breaks for secondary y-axis
    trans <- function(x) abs(x - 1) * 100 #trasform p value p=0.05 -> 95%.
    
    # Define breaks for the primary y-axis (the one on the right)
    b_points <- c(1 - const$conf_level, pretty(c(lines$y, 1)))
    o <- order(b_points, decreasing = FALSE)
    breaks_y1 <- b_points[o]
    
    # Compute breaks for the secondary y-axis
    # breaks_y2 <- trans(b_points[o])
    # Set transparency
    transparency <- 1

    #add plots for single studies
    p <- ggplot2::ggplot(
      data = lines,
      ggplot2::aes(x = x, y = y, color = group)
    ) +
      ggplot2::geom_hline(yintercept = 1 - const$conf_level, linetype = "dashed") + #line at CI
      ggplot2::geom_vline(xintercept = 0, linetype = "dashed") #line at mu = 0
    
    
    # DEFINE THE axis scale!s
    # Changed completely the logic since sometimes the x axis was problematic in scaling, now pretty_breaks handles everything!
    
    p <- p +
      ggplot2::scale_x_continuous(
        # limits = xlim, old version
        breaks = scales::pretty_breaks(n = 7),  #7 number is just arbitrary
        minor_breaks = NULL
      ) +
      ggplot2::coord_cartesian(xlim = xlim, ylim = c(0, 1), expand = FALSE) # this allows us to ZOOM in the plot, not cut it 
    
    #now add the plot for baseline method
    if (!drapery) { #only vertical lines if drapery FALSE
        p <- p + ggplot2::geom_vline(
            xintercept = estimates, #or const$estimates
            linetype = "dashed"
        )
    } else { #otwise add the drapery
        p <- p +
            ggplot2::geom_line(
                data = dp,
                mapping = ggplot2::aes(x = x, y = y, group = study),
                linetype = "dashed",
                color = "lightgrey",
                show.legend = FALSE
            ) +
            ggplot2::geom_line(
                data = rmadf,
                mapping = ggplot2::aes(x = x, y = y, color = group),
                # color = "#00000099",
                show.legend = TRUE
            )


    }
    if (0 > xlim[1L] && 0 < xlim[2L]) {
        # Vertical line at 0
        p <- p + ggplot2::geom_vline(xintercept = 0, linetype = "dashed") +
            ggplot2::geom_point(
                # Points at (x = 0, y = p_0) (it is the H0, put a point there)
                data = lines[!is.na(lines$y0), ],
                ggplot2::aes(x = 0, y = y0, color = group),
                alpha = transparency
            )
    }
    p <- p +
        ggplot2::geom_hline(yintercept = 0, linetype = "dashed") + #horiz line at y=0
        ggplot2::geom_line(alpha = transparency) +
        ggplot2::scale_y_continuous(
            name = "p-value", #left axis 
            breaks = breaks_y1,
            #limits = c(0, 1), now there should be coord_cartesian
            #expand = c(0, 0),
            sec.axis = ggplot2::sec_axis(
                transform = trans,
                name = "Confidence level [%]",
                breaks = trans(breaks_y1)
            )
        ) +
        # Draw intervals on x-axis
        ggplot2::geom_segment(
            data = errorbars[!is.na(errorbars$xmin), ],
            ggplot2::aes(x = xmin, xend = xmax, y = y, yend = y)
        ) +
        ggplot2::geom_segment(
            data = errorbars[!is.na(errorbars$xmin), ],
            ggplot2::aes(x = xmin, xend = xmin, y = ymin, yend = ymax)
        ) +
        ggplot2::geom_segment(
            data = errorbars[!is.na(errorbars$xmin), ],
            ggplot2::aes(x = xmax, xend = xmax, y = ymin, yend = ymax)
        ) +
        # Set x-axis window, labels
        # ggplot2::labs(
        #     x = bquote(mu),
        #     color = "Configuration"
        # ) +
        # Set theme
        ggplot2::theme_minimal() +
        ggplot2::theme(
            axis.title.y.right = ggplot2::element_text(angle = 90),
            legend.position = "bottom"
        )

    # Add axis/legend label
    xlab <- if (is.null(xlab)) {
        bquote(mu)
    } else {
        xlab
    }
    p <- p +
        ggplot2::labs(
          x = xlab,
          color = NULL
        ) +
        ggplot2::theme(
          legend.title = ggplot2::element_blank()
        )
    

    # Add SHARED coloring for curves
    p <- p + ggplot2::scale_color_manual(values = shared_colors)
    

    
    
    
    # return
    p
}

# Calculate the drapery lines, i.e. p-value functions of single studies, that are plotted as 
# grey dashed lines
#' @importFrom stats pnorm
#' @noRd
get_drapery_df <- function(estimates, SEs, mu) {
  
    # get lenghts
    l_t <- length(estimates) #n studies
    l_m <- length(mu) 
    l_tot <- l_t * l_m #total n points to generate
    
    # Initialize vectors
    x_dp <- rep(mu, times = l_t)  #create the grids [mu1,...,mu10000, mu1, ..., mu10000, ..] repetead for each study
    y_dp <- study <- numeric(l_tot)
    
    # Indices to loop over
    idx <- seq_len(l_m)
    nms <- names(estimates)
    
    # Loop on each study
    for (i in seq_along(estimates)) {
        y_dp[idx] <- 2 *
            (1 - stats::pnorm(abs(estimates[i] - mu) / SEs[i])) #compute bilateral p value function for study i
        study[idx] <- if (is.null(nms)) rep(i, l_m) else rep(nms[i], l_m) # assing identificator, if the studies have names uses them, otwise use index
        idx <- idx + l_m
    }
    #combine studies in a single df
    data.frame( 
        x = x_dp,
        y = y_dp,
        study = study,
        stringsAsFactors = FALSE
    )
}

map_ref_methods <- function(abbrevs, labels = NULL) {
  defaults <- c(
    re = "Random-effects",
    hk = "Hartung & Knapp",
    hc = "Henmi & Copas",
    fe = "Fixed-effect"
  )
  
  # UPDATED FUNCTION: drop the switch and set the default dictionary, then if we have labels provided
  # replace them with the names we have
  if (!is.null(labels)) {
    defaults[names(labels)] <- labels
  }
  unname(defaults[abbrevs])
  
}

# ==============================================================================
# forest plot
# ==============================================================================

#' @importFrom ggplot2 ggplot xlim geom_vline geom_errorbar aes geom_point
#' @importFrom ggplot2 geom_polygon theme_minimal theme element_blank
#' @importFrom ggplot2 element_line scale_y_continuous labs scale_fill_discrete
#' @importFrom ggplot2 annotate
#' @importFrom scales hue_pal
#' @importFrom rlang .data
#' @noRd

ForestPlot <- function(
    cms,
    diamond_height,
    v_space,
    show_studies,
    scale_diamonds,
    reference_methods,
    xlim,
    xlab,
    n_breaks,
    shared_colors,
    ref_labels
) {

      # Make a data frame for the single studies
    cm <- cms[[1L]]
    studyCIs <- data.frame(
        lower = cm$individual_cis[, 1L],
        upper = cm$individual_cis[, 2L],
        estimate = cm$estimates,
        name = cm$study_names,
        plottype = 0L,
        color = 0L,
        stringsAsFactors = FALSE,
        row.names = NULL
    )

    # Get the dataframe for the polygons of the new methods
    new_method_cis <- get_CI_new_methods(
        cms = cms,
        diamond_height = diamond_height,
        scale_diamonds = scale_diamonds
    )

    # RETRIEVE (not recompute) the dataframe for the polygons of the old methods
    old_methods_cis <- get_CI_old_methods(
      cm = cm,
      diamond_height = diamond_height
    )



    # Rename the default names to ref_labels!
    # first creates named vector (like Fixed effect --> Fixed effect (IV))
    default_to_custom <- setNames(
      map_ref_methods(c("fe", "re", "hk", "hc"), labels = ref_labels),
      map_ref_methods(c("fe", "re", "hk", "hc"))
    )
    
    #replace
    old_methods_cis$CIs$name <- default_to_custom[old_methods_cis$CIs$name] # ? maybe wrap those i unname? 
    old_methods_cis$p_0$name <- default_to_custom[old_methods_cis$p_0$name]
    
    
    reference_methods <- map_ref_methods(abbrevs = reference_methods, labels = ref_labels)

    keep_cis <- with(old_methods_cis, CIs$name %in% reference_methods)
    keep_p0 <- with(old_methods_cis, p_0$name %in% reference_methods)

    old_methods_cis$CIs <- old_methods_cis$CIs[keep_cis, ]
    old_methods_cis$p_0 <- old_methods_cis$p_0[keep_p0, ]

    ###########################################################################
    # Quick fix, remove the above at some point                               #
    ###########################################################################

    # Assemble p_0
    p_0 <- rbind(old_methods_cis$p_0, new_method_cis$p_0)

    # Assemble the dataset
    polygons <- rbind(old_methods_cis$CIs, new_method_cis$CIs)

    # Some cosmetics
    ## If any of the CIs does not exist, replace the row
    na_cis <- anyNA(polygons$x)
    if (na_cis) {
        # Which methods have an NA CI
        na_methods <- with(polygons, unique(name[is.na(x)]))
        # Mark it with a new variable
        polygons$ci_exists <- ifelse(polygons$name %in% na_methods, FALSE, TRUE)
        # Set y coordinate to 0
        idx <- with(polygons, name %in% na_methods)
        polygons$y[idx] <- 0
        polygons$x[idx] <- 0
    } else {
        polygons$ci_exists <- TRUE
    }
    ## Assign correct y-coordinate
    methods <- unique(polygons$name)
    n_methods <- length(methods)
    spacing <- seq(
        (nrow(studyCIs) + n_methods + 1L) * v_space,
        v_space,
        by = -v_space
    )
    studyCIs$y <- spacing[seq_len(nrow(studyCIs))]
    method_spacing <- spacing[(nrow(studyCIs) + 2L):length(spacing)]
    for (i in seq_along(methods)) {
        idx <- with(polygons, name == methods[i])
        polygons$y[idx] <- polygons$y[idx] + method_spacing[i]
    }
    ## Manage colors
    polygons$color_hex <- ifelse(
      polygons$name %in% names(shared_colors),
      shared_colors[polygons$name],
      "gray20"
    )
    
    # Make the plot
    p <- ggplot2::ggplot()
    
    
    # DEFINE THE axis scale! smartly 
    
    if (!is.null(xlim)) {
      p <- p +
        ggplot2::scale_x_continuous(
          # limits = xlim,
          breaks = scales::pretty_breaks(n = n_breaks),
          minor_breaks = NULL
        ) + 
        ggplot2::coord_cartesian(xlim = xlim, clip = "on", expand = FALSE) # 
      
      # vertical line at 0
      if (0 > xlim[1L] && 0 < xlim[2L]) {
        p <- p + ggplot2::geom_vline(xintercept = 0, linetype = "dashed")
      }
    } else {
      p <- p +
        ggplot2::scale_x_continuous(
          breaks = scales::pretty_breaks(n = n_breaks),
          minor_breaks = NULL
        )
    }
    
    if (show_studies) {
        p <- p +
            ggplot2::geom_errorbar(
                orientation = "y",
                data = studyCIs,
                ggplot2::aes(
                    y = y,
                    xmin = lower,
                    xmax = upper,
                ),
                width = diamond_height, #this parameter was inside for older ggplot (called height), now deprecated
                show.legend = FALSE
            ) +
            ggplot2::geom_point(
                data = studyCIs,
                ggplot2::aes(x = estimate, y = y)
            )
    }
    p <- p +
      ggplot2::geom_polygon(
        data = polygons[polygons$ci_exists == TRUE, ],
        ggplot2::aes(
          x = x,
          y = y,
          group = paste0(name, ".", id),
          fill = .data$color_hex
        ),
        show.legend = FALSE
      ) +
        ggplot2::theme_minimal() +
        ggplot2::theme(
            axis.title.y = ggplot2::element_blank(),
            axis.line.x = ggplot2::element_line(linetype = "solid"),
            axis.ticks.x = ggplot2::element_line(linetype = "solid"),
            panel.border = ggplot2::element_blank(),
            axis.ticks = ggplot2::element_blank(),
            panel.grid.minor = ggplot2::element_blank(),
            panel.grid.major.y = ggplot2::element_blank()
        ) +
        ggplot2::scale_y_continuous(
            breaks = spacing,
            labels = c(studyCIs$name, "", unique(polygons$name)),
            limits = c(min(spacing) - v_space/2, max(spacing) + v_space/2) # this fixes the problem where that the last polygon was always too close to the x axis! 
            
        ) +
        # ggplot2::labs(
        #     x = bquote(mu)
        # ) +
      ggplot2::scale_fill_identity()
    
    if (is.null(xlab)) {
        p <- p +
            ggplot2::labs(
                x = bquote(mu)
            )
    } else {
        p <- p +
            ggplot2::labs(
                x = xlab
            )
    }

    if (na_cis) {
        na_rows <- polygons[polygons$ci_exists == FALSE, ]
        for (i in seq_len(nrow(na_rows))) {
            p <- p + ggplot2::annotate(
                geom = "text",
                x = na_rows$x[i],
                y = na_rows$y[i],
                label = "CI does not exist"
            )
        }
    }

    p
}


get_CI_new_methods <- function(
    cms,
    diamond_height,
    scale_diamonds
) {

    out <- lapply(seq_along(cms), function(r) {

        # Get one of the cms
        cm <- cms[[r]]

        # Calculate polygons
        ci_exists <- if (all(is.na(cm$joint_cis))) FALSE else TRUE

        if (ci_exists) {
            # Calculate the polygons for the diamond
            polygons <- calculate_polygon(
                cm = cm,
                diamond_height = diamond_height,
                scale_diamonds = scale_diamonds
            )
            polygons$name <- cm$fun_name
            polygons$color <- r

        } else {
            polygons <- data.frame(
                x = NA_real_,
                y = NA_real_,
                id = NA_real_,
                name = cm$fun_name,
                color = r
            )
        }

        # Calculate the p-value at mu = 0
        p_0 <- data.frame(
            name = cm$fun_name,
            p_0 = cm$p_0[, 2L],
            stringsAsFactors = FALSE,
            row.names = NULL
        )

        p_max <- data.frame(
            name = cm$fun_name,
            mu = cm$p_max[, 1L],
            p_max = cm$p_max[, 2L],
            stringsAsFactors = FALSE,
            row.names = NULL
        )

        # Return
        list(
            CIs = polygons,
            p_0 = p_0,
            p_max = p_max
        )
    })

    # Reorganize list
    CIs <- do.call("rbind", lapply(out, "[[", i = 1L))
    p_0 <- do.call("rbind", lapply(out, "[[", i = 2L))
    p_max <- do.call("rbind", lapply(out, "[[", i = 3L))

    # Return
    list(CIs = CIs, p_0 = p_0, p_max = p_max)
}


get_CI_old_methods <- function(
    cm,              
    diamond_height) 
  {


    # in previous version, there was a call to get_stats_others, and the CI comparison were recomputed. 
    #Now i just extract them 
  
   cis <- cm$comparison_cis
   p_0_vals <- cm$comparison_p_0
   
   
   #DIAMOND IS:
   # Left Corner: (Lower, 0)
   # Bottom Corner: (Est, -H/2) 
   # Right Corner: (Upper, 0)
   # Top Corner: (Est, +H/2)
   
  # get the estimates (center of CI)
   ests <- rowMeans(cis)
   
    # Create df
    t_ci <- t(cis) #row1 = contains all Lower bounds, row2 all lower bounds
    
    
    t_ci <- rbind(t_ci[1L, ], ests, t_ci[2L, ], ests) 
    
    # t_ci has:
    #Row 1: Lower bounds (Left corner x)
    #Row 2: Estimates (Bottom corner x)
    #Row 3: Upper bounds (Right corner x)
    #Row 4: Estimates (Top corner x)
    
    x <- c(t_ci) #flatten columnwise st we obtain:: Lower1, Est1, Upper1, Est1, Lower2, Est2, Upper2, Est2 ...
    
    #create y coordinates 
    y <- rep(
        c(0, -diamond_height / 2, 0, diamond_height / 2),
        times = length(ests)
    )
    
    method_names <- rownames(cis)

    CIs <- data.frame(
        x = x,
        y = y,
        id = 1L,
        name = rep(method_names, each = 4L),
        color = 0L,
        row.names = NULL,
        stringsAsFactors = FALSE
    )

    p_0 <- data.frame(
        name = method_names,
        p_0 = p_0_vals[, 2L],
        row.names = NULL,
        stringsAsFactors = FALSE
    )

    list(
        CIs = CIs,
        p_0 = p_0
    )
}

# Make a df that contains the data for the diamonds
calculate_polygon <- function(
    cm,
    diamond_height,
    scale_diamonds
) {

    # If gamma == NA, there is either one or no CI
    no_ci <- if (all(is.na(cm$joint_cis))) TRUE else FALSE
    no_gamma <- if (all(is.na(cm$gamma))) TRUE else FALSE

    # Which points to evaluate for the diamonds (all maxima, minima, estimates)
    # type of point
    # 0 = lower ci
    # 1 = minimum
    # 2 = estimate or maximum
    # 3 = upper ci

    
    pt_eval <- with(
        cm,
        {
          f_estimates <- call_pfun(
            fun = p_fun,
            estimates = estimates,
            SEs = SEs,
            mu = estimates,
            w = w
          )

            nrep <- nrow(joint_cis)
            m <- rbind(
                cbind(p_max, 2),  #max --> type 2
                matrix(
                    c(estimates, f_estimates, rep(2, length(cm$estimates))), #estimates --> type 2
                    ncol = 3L
                ),
                matrix(c(joint_cis[, 1L], rep(0, 2 * nrep)), ncol = 3L), #lower CI --> type 0
                matrix(
                    c(joint_cis[, 2L], rep(0, nrep), rep(3, nrep)), #upper CI --> type 3
                    ncol = 3L
                )
            )
            if (!no_gamma) m <- rbind(m, cbind(gamma, 1)) #gamma --> type 1
            m
        }
    )

    # Remove duplicates
    pt_eval <- pt_eval[!duplicated(pt_eval), ]

    # Remove those not in CI
    idx <- vapply(
        pt_eval[, 1L],
        function(x, cis) any(x >= cis[, 1L] & x <= cis[, 2L]),
        logical(1L),
        cis = cm$joint_cis
    )
    pt_eval <- pt_eval[idx, ]

    # Sort these by x-values
    o <- order(pt_eval[, 1L], decreasing = FALSE)
    pt_eval <- pt_eval[o, ]

    if (no_ci) {
        return(NULL)
    } else {
        x <- pt_eval[, 1L]
        y <- pt_eval[, 2L]
        type <- pt_eval[, 3L]
        # rescale the diamonds if needed (such that the max height is always 1)
        if (scale_diamonds) {
            scale_factor <- 1 / max(y)
            y <- y * scale_factor
        }
        # Assign polygon id
        is_upper <- which(type == 3L)
        is_lower <- which(type == 0L)
        id <- vector("numeric", length(x))
        counter <- 1L
        for (i in seq_along(is_upper)) {
            id[is_lower[i]:is_upper[i]] <- counter
            counter <- counter + 1L
        }
        # Arrange this such that it gives back a data.frame
        # that can be used with ggplot
        as.data.frame(
            do.call(
                "rbind",
                lapply(
                    unique(id),
                    function(z) {
                        idx <- id == z
                        x_sub <- x[idx]
                        y_sub <- y[idx] * diamond_height / 2
                        l <- length(x_sub)
                        mat <- matrix(
                            NA_real_,
                            ncol = 3,
                            nrow = 2L * l - 2L
                        )
                        colnames(mat) <- c("x", "y", "id")
                        mat[, 1L] <- c(x_sub, rev(x_sub[-c(1L, l)]))
                        mat[, 2L] <- c(-y_sub, rev(y_sub[-c(1L, l)]))
                        mat[, 3L] <- z
                        mat
                    }
                )
            )
        )
    }
}





# ------------------------------------------------------------------------------
# Helper to call p_val function since we added weights
# ------------------------------------------------------------------------------

call_pfun <- function(fun, estimates, SEs, mu, w = NULL) {
  args <- names(formals(fun))
  
  # if 'w' exists in the arguments and we have a valide w--> pass 
  if ("w" %in% args && !is.null(w)) {
    fun(estimates = estimates, SEs = SEs, mu = mu, w = w)
  } else {
    # otwise ignore weigths
    fun(estimates = estimates, SEs = SEs, mu = mu)
  }
}





#  ---------------------------------------------------------------------------- 
#   Add the bayesmeta object to the forest plot. 
# --------------------------------------------------------------------------------

add_bayes_forest <- function(p, bm, diamond_height, v_space, 
                             color = "#000000", label = "Bayesmeta",
                             mu_estimate = "median") {
  
  # Extract summary from the bayesmeta object
  s <- bm$summary
  mu_est <- s[mu_estimate, "mu"]
  lower  <- s["95% lower", "mu"]
  upper  <- s["95% upper", "mu"]
  
  # Get the current breaks and labels from the ggplot object's scale
  y_scale <- p$scales$get_scales("y")
  y_breaks <- y_scale$breaks
  y_labels <- y_scale$labels
  
  # Calculate new position (one step down using v_space)
  y_bayes <- min(y_breaks) - v_space
  height <- diamond_height / 2
  
  # Define polygon for the Bayesmeta diamond
  diamond <- data.frame(
    x = c(lower, mu_est, upper, mu_est),
    y = c(y_bayes, y_bayes + height, y_bayes, y_bayes - height)
  )
  
  # Add the diamond
  p <- p +
    ggplot2::geom_polygon(
      data = diamond,
      ggplot2::aes(x = x, y = y),
      fill = color,
      color = color
      )
  
  # Extend y-axis labels and breaks
  y_breaks <- c(y_breaks, y_bayes)
  y_labels <- c(y_labels, label)
  
  # Overwrite the scale
  suppressMessages({
    p <- p + ggplot2::scale_y_continuous(
      breaks = y_breaks,
      labels = y_labels,
      limits = c(y_bayes - v_space, max(y_breaks) + (v_space / 2))
      )
  })
  
  p
}