R/plot_methods.R

In FPScausal: Functional Propensity Score for Causal Inference

#' Plot diagnostics for fps_weighting objects
#'
#' @param x An object of class \code{"fps_weighting"}.
#' @param type Character. One of:
#'   \describe{
#'     \item{\code{"balance"}}{Absolute Pearson correlations between each
#'       treatment FPC score and each confounder, before (red) and after
#'       (blue) weighting, displayed as connected point-line chart.}
#'     \item{\code{"fpca_treatment"}}{Scree plot (eigenvalues + cumulative PVE)
#'       and panel of leading eigenfunctions for the treatment FPCA.}
#'     \item{\code{"fpca_covariates"}}{Same as \code{"fpca_treatment"} for each
#'       functional covariate (returns a list of \pkg{ggplot2} objects).}
#'     \item{\code{"weights"}}{Boxplot of the estimated weights with a
#'       horizontal reference line at the uniform weight 1/n.}
#'   }
#' @param max_efn Integer. Maximum number of eigenfunctions to show in FPCA
#'   panels (default 4).
#' @param ... Ignored.
#' @return A \pkg{ggplot2} object (or a list of them for
#'   \code{type = "fpca_covariates"}).
#' @export
plot.fps_weighting <- function(x, type = "balance", max_efn = 4, ...) {
  type <- match.arg(type,
                    c("balance", "fpca_treatment", "fpca_covariates", "weights"))
  switch(type,
    balance        = .plot_balance(x),
    fpca_treatment = .plot_fpca_panel(x$fpca_treatment,
                                       title      = "Treatment FPCA",
                                       max_efn    = max_efn,
                                       domain_name = x$domain_name),
    fpca_covariates = {
      if (is.null(x$fpca_covariates)) {
        message("No functional covariates in this fps_weighting object.")
        return(invisible(NULL))
      }
      lapply(seq_along(x$fpca_covariates), function(k) {
        .plot_fpca_panel(x$fpca_covariates[[k]],
                         title      = paste0("Functional Covariate ", k,
                                             " FPCA"),
                         max_efn    = max_efn,
                         domain_name = "t")
      })
    },
    weights = .plot_weights(x)
  )
}

#' Plot diagnostics for fps_effect_estimation objects
#'
#' @param x An object of class \code{"fps_effect_estimation"}.
#' @param type Character. One of:
#'   \describe{
#'     \item{\code{"effect"}}{Weighted estimate of \eqn{\mu(t)} (or
#'       \eqn{\mu(s,t)}) with CI ribbon and optional true-mu overlay.
#'       Uses bootstrap CI if available, otherwise analytical SE-based CI.}
#'     \item{\code{"comparison"}}{Weighted vs unweighted side by side, both
#'       with CI ribbons.}
#'     \item{\code{"fpca_treatment"}}{Scree + eigenfunctions of the treatment
#'       FPCA used in estimation.}
#'     \item{\code{"fpca_outcome"}}{Scree + eigenfunctions of the outcome FPCA
#'       (functional outcome only).}
#'     \item{\code{"bootstrap_slice"}}{1-D slice(s) of the causal effect
#'       surface with bootstrap CI band.  Pass a scalar or vector to
#'       \code{point}; multiple points yield a patchwork panel.  Functional
#'       outcome only.}
#'     \item{\code{"significance"}}{Regions / points where the CI excludes 0.}
#'   }
#' @param point Numeric scalar or vector. Time point(s) at which to slice the
#'   effect surface (for \code{type = "bootstrap_slice"}).
#' @param which_domain Character. Either \code{"treatment"} or
#'   \code{"outcome"} (for \code{type = "bootstrap_slice"}).
#' @param alpha Numeric. Significance level; defaults to \code{x$alpha}.
#' @param max_efn Integer. Maximum eigenfunctions shown in FPCA panels.
#' @param ... Ignored.
#' @return A \pkg{ggplot2} object.
#' @export
plot.fps_effect_estimation <- function(x, type = "effect",
                                        point        = NULL,
                                        which_domain = "treatment",
                                        alpha        = NULL,
                                        max_efn      = 4,
                                        ...) {
  type <- match.arg(type, c("effect", "comparison", "fpca_treatment",
                              "fpca_outcome", "bootstrap_slice", "significance"))
  if (is.null(alpha)) alpha <- x$alpha

  switch(type,
    effect         = .plot_effect(x),
    comparison     = .plot_comparison(x),
    fpca_treatment = .plot_fpca_panel(x$fpca_treatment,
                                       title      = "Treatment FPCA",
                                       max_efn    = max_efn,
                                       domain_name = x$domain_name),
    fpca_outcome   = {
      if (is.null(x$fpca_outcome)) {
        stop("No outcome FPCA available (only for functional outcomes).")
      }
      .plot_fpca_panel(x$fpca_outcome,
                       title      = "Outcome FPCA",
                       max_efn    = max_efn,
                       domain_name = x$outcome_domain_name)
    },
    bootstrap_slice = {
      if (is.null(x$ci_lower)) {
        stop(paste0("Bootstrap CIs not available. Re-run ",
                    "fps_effect_estimation with bootstrap = TRUE."))
      }
      if (x$outcome_type != "functional") {
        stop("'bootstrap_slice' is only available for functional outcomes.")
      }
      if (is.null(point)) stop("'point' must be specified for bootstrap_slice.")
      if (length(point) == 1) {
        .plot_bootstrap_slice(x, point, which_domain)
      } else {
        plots <- lapply(point,
                        function(pt) .plot_bootstrap_slice(x, pt, which_domain))
        patchwork::wrap_plots(plots, ncol = length(point))
      }
    },
    significance = .plot_significance(x, alpha)
  )
}

# ============================================================
# Internal plot helpers
# ============================================================

#' Balance plot: point-line chart of absolute correlations
#' @keywords internal
.plot_balance <- function(x) {
  w    <- x$weights
  A    <- x$fpca_treatment$scr
  conf <- x$conf_matrix
  L    <- ncol(A)
  p    <- ncol(conf)

  conf_names <- colnames(conf)
  if (is.null(conf_names)) {
    conf_names <- paste0("Conf", seq_len(p))
  } else {
    blank <- !nzchar(conf_names)
    if (any(blank)) {
      counter <- 0L
      conf_names[blank] <- vapply(which(blank), function(i) {
        counter <<- counter + 1L
        paste0("Conf", counter)
      }, character(1))
    }
  }

  rows <- expand.grid(conf = seq_len(p), fpc = seq_len(L))

  cor_uw <- vapply(seq_len(nrow(rows)), function(i) {
    abs(stats::cor(A[, rows$fpc[i]], conf[, rows$conf[i]]))
  }, numeric(1))

  cor_w <- vapply(seq_len(nrow(rows)), function(i) {
    abs(wCorr::weightedCorr(A[, rows$fpc[i]], conf[, rows$conf[i]],
                             method = "Pearson", weights = w))
  }, numeric(1))

  lab <- paste0(conf_names[rows$conf], " | FPC", rows$fpc)

  df <- data.frame(
    label      = factor(lab, levels = lab),
    Unweighted = cor_uw,
    Weighted   = cor_w
  )

  df_long <- tidyr::pivot_longer(df, cols = c("Unweighted", "Weighted"),
                                  names_to  = "type",
                                  values_to = "correlation")
  df_long$type <- factor(df_long$type, levels = c("Unweighted", "Weighted"))

  ggplot2::ggplot(df_long,
    ggplot2::aes(x = label, y = correlation,
                  colour = type, group = type)) +
    ggplot2::geom_line(linewidth = 0.7) +
    ggplot2::geom_point(size = 2.5) +
    ggplot2::scale_colour_manual(
      values = c("Unweighted" = "#8B0000", "Weighted" = "#00008B")
    ) +
    ggplot2::geom_hline(yintercept = 0.1, linetype = "dashed",
                        colour = "grey40", linewidth = 0.4) +
    ggplot2::labs(x = NULL, y = "|Pearson correlation|",
                  colour = NULL, title = "Covariate Balance") +
    ggplot2::theme_bw() +
    ggplot2::theme(
      axis.text.x        = ggplot2::element_text(
                             angle = 90, hjust = 1, vjust = 0.5, size = 8),
      legend.position    = "top",
      plot.title         = ggplot2::element_text(hjust = 0.5, face = "bold"),
      panel.grid.major.x = ggplot2::element_blank()
    )
}

#' FPCA panel: scree + eigenfunctions
#' @keywords internal
.plot_fpca_panel <- function(fpca_res, title = "FPCA", max_efn = 4,
                              domain_name = "t") {
  tgrid  <- fpca_res$t_grid
  L      <- fpca_res$L
  evals  <- fpca_res$eval
  perc   <- fpca_res$perc
  efn    <- fpca_res$efn
  n_show <- min(L, max_efn)

  df_scree <- data.frame(
    Component     = seq_along(evals),
    Eigenvalue    = evals,
    CumulativePVE = perc
  )
  df_scree_show <-
    df_scree[seq_len(min(nrow(df_scree), max(n_show + 2, 6))), ]

  p_scree <- ggplot2::ggplot(df_scree_show,
    ggplot2::aes(x = Component, y = Eigenvalue)) +
    ggplot2::geom_col(fill = "#4393C3", alpha = 0.8) +
    ggplot2::geom_line(ggplot2::aes(y = CumulativePVE * max(Eigenvalue)),
                       colour = "#D62728", linewidth = 0.8) +
    ggplot2::geom_point(ggplot2::aes(y = CumulativePVE * max(Eigenvalue)),
                        colour = "#D62728", size = 2) +
    ggplot2::scale_y_continuous(
      name     = "Eigenvalue",
      sec.axis = ggplot2::sec_axis(
        ~ . / max(df_scree_show$Eigenvalue),
        name = "Cumulative PVE"
      )
    ) +
    ggplot2::scale_x_continuous(breaks = df_scree_show$Component) +
    ggplot2::labs(x = "Component", title = paste(title, "- Scree")) +
    ggplot2::theme_bw() +
    ggplot2::theme(plot.title = ggplot2::element_text(hjust = 0.5))

  df_efn <- do.call(rbind, lapply(seq_len(n_show), function(k) {
    data.frame(
      time      = tgrid,
      value     = efn[, k],
      Component = factor(paste0("FPC ", k,
                                " (", round(fpca_res$varprop[k] * 100, 1),
                                "%)"))
    )
  }))

  p_efn <- ggplot2::ggplot(df_efn,
    ggplot2::aes(x = time, y = value, colour = Component)) +
    ggplot2::geom_line(linewidth = 0.9) +
    ggplot2::geom_hline(yintercept = 0, linetype = "dashed",
                        colour = "grey50") +
    ggplot2::labs(x = domain_name, y = "Eigenfunction value",
                  colour = NULL,
                  title  = paste(title, "- Eigenfunctions")) +
    ggplot2::theme_bw() +
    ggplot2::theme(legend.position = "top",
                   plot.title = ggplot2::element_text(hjust = 0.5))

  patchwork::wrap_plots(p_scree, p_efn, ncol = 2)
}

#' Weights boxplot
#' @keywords internal
.plot_weights <- function(x) {
  n  <- length(x$weights)
  df <- data.frame(group = "Weights", weight = x$weights)

  ggplot2::ggplot(df, ggplot2::aes(x = group, y = weight)) +
    ggplot2::geom_boxplot(fill = "#4393C3", alpha = 0.8,
                           outlier.colour = "#D62728",
                           outlier.size   = 1.5) +
    ggplot2::geom_hline(yintercept = 1 / n, linetype = "dashed",
                        colour = "#D62728", linewidth = 0.8) +
    ggplot2::annotate("text", x = 1, y = 1 / n, label = "  1/n",
                      vjust = -0.5, colour = "#D62728", size = 3.5) +
    ggplot2::labs(x = NULL, y = "Weight",
                  title = "Distribution of FPS Weights") +
    ggplot2::theme_bw() +
    ggplot2::theme(
      axis.text.x  = ggplot2::element_blank(),
      axis.ticks.x = ggplot2::element_blank(),
      plot.title   = ggplot2::element_text(hjust = 0.5)
    )
}

#' Effect plot: mu(t) or mu(s,t) with CI
#' @keywords internal
.plot_effect <- function(x) {
  if (x$outcome_type == "functional") {
    .plot_effect_surface(x, beta = x$beta,
                          ci_lower = x$ci_lower, ci_upper = x$ci_upper,
                          title = "Weighted Causal Effect Surface")
  } else {
    .plot_effect_1d(x, use_weighted = TRUE, show_true = TRUE)
  }
}

#' Comparison plot: weighted vs unweighted, both with CI
#' @keywords internal
.plot_comparison <- function(x) {
  if (x$outcome_type == "functional") {
    p1 <- .plot_effect_surface(x, beta = x$beta_unweighted,
                                ci_lower = NULL, ci_upper = NULL,
                                title = "Unweighted")
    p2 <- .plot_effect_surface(x, beta = x$beta,
                                ci_lower = x$ci_lower, ci_upper = x$ci_upper,
                                title = "Weighted")
    patchwork::wrap_plots(p1, p2, ncol = 2)
  } else {
    tgrid <- x$t_grid
    alpha <- x$alpha
    dn    <- x$domain_name
    z     <- stats::qnorm(1 - alpha / 2)

    # Determine CI type for title
    ci_type_lbl <- if (!is.null(x$ci_lower)) {
      "bootstrap CI"
    } else if (!is.null(x$analytical_se)) {
      "analytical CI"
    } else {
      "no CI"
    }

    # CIs: bootstrap preferred, otherwise analytical SE
    if (!is.null(x$ci_lower)) {
      lo_w <- x$ci_lower; hi_w <- x$ci_upper
    } else if (!is.null(x$analytical_se)) {
      lo_w <- x$beta - z * x$analytical_se
      hi_w <- x$beta + z * x$analytical_se
    } else {
      lo_w <- hi_w <- rep(NA_real_, length(tgrid))
    }

    if (!is.null(x$analytical_se)) {
      lo_u <- x$beta_unweighted - z * x$analytical_se
      hi_u <- x$beta_unweighted + z * x$analytical_se
    } else {
      lo_u <- hi_u <- rep(NA_real_, length(tgrid))
    }

    # line data
    df_lines <- data.frame(
      time      = rep(tgrid, 2),
      effect    = c(x$beta_unweighted, x$beta),
      beta_type = rep(c("Unweighted", "Weighted"), each = length(tgrid))
    )
    if (!is.null(x$true_beta)) {
      df_lines <- rbind(df_lines,
        data.frame(time = tgrid, effect = x$true_beta, beta_type = "True"))
    }
    lvls <- intersect(c("True", "Unweighted", "Weighted"),
                      unique(df_lines$beta_type))
    df_lines$beta_type <- factor(df_lines$beta_type, levels = lvls)

    cols <- c("True" = "black", "Unweighted" = "#4393C3", "Weighted" = "#D62728")
    ltys <- c("True" = "dashed", "Unweighted" = "solid", "Weighted" = "solid")

    # ribbon data (separate, no colour scale conflict)
    df_rib_w <- data.frame(time = tgrid, lower = lo_w, upper = hi_w)
    df_rib_u <- data.frame(time = tgrid, lower = lo_u, upper = hi_u)

    y_label <- bquote(hat(mu)(.(dn)))

    p <- ggplot2::ggplot(df_lines,
                          ggplot2::aes(x = time, y = effect,
                                       colour   = beta_type,
                                       linetype = beta_type)) +
      ggplot2::geom_ribbon(data = df_rib_w,
                            ggplot2::aes(x = time, ymin = lower, ymax = upper),
                            inherit.aes = FALSE,
                            fill = "#D62728", alpha = 0.15, na.rm = TRUE) +
      ggplot2::geom_ribbon(data = df_rib_u,
                            ggplot2::aes(x = time, ymin = lower, ymax = upper),
                            inherit.aes = FALSE,
                            fill = "#4393C3", alpha = 0.15, na.rm = TRUE) +
      ggplot2::geom_line(linewidth = 0.9) +
      ggplot2::geom_hline(yintercept = 0, linetype = "dotted",
                          colour = "grey50") +
      ggplot2::scale_colour_manual(values  = cols[lvls], name = NULL) +
      ggplot2::scale_linetype_manual(values = ltys[lvls], name = NULL) +
      ggplot2::labs(
        x     = dn,
        y     = y_label,
        title = paste0("Weighted vs Unweighted Causal Effect (",
                       ci_type_lbl, ")")
      ) +
      ggplot2::theme_bw() +
      ggplot2::theme(
        legend.position  = "top",
        legend.key.width = ggplot2::unit(2, "cm"),
        plot.title       = ggplot2::element_text(hjust = 0.5)
      )
    p
  }
}

#' 1-D effect plot with CI ribbon and legend
#' @keywords internal
.plot_effect_1d <- function(x, use_weighted = TRUE, show_true = TRUE) {
  tgrid <- x$t_grid
  alpha <- x$alpha
  dn    <- x$domain_name
  beta  <- if (use_weighted) x$beta else x$beta_unweighted

  # Pick CIs: bootstrap preferred, else analytical
  if (use_weighted && !is.null(x$ci_lower)) {
    lo <- x$ci_lower
    hi <- x$ci_upper
    ci_label <- paste0(round((1 - alpha) * 100), "% bootstrap CI")
  } else if (!is.null(x$analytical_se)) {
    z  <- stats::qnorm(1 - alpha / 2)
    lo <- beta - z * x$analytical_se
    hi <- beta + z * x$analytical_se
    ci_label <- paste0(round((1 - alpha) * 100), "% analytical CI")
  } else {
    lo <- hi <- NULL
    ci_label <- NULL
  }

  # Build line data frame for legend
  df_lines <- data.frame(time = tgrid, effect = beta,
                          curve_type = "Estimated")
  if (show_true && !is.null(x$true_beta)) {
    df_lines <- rbind(df_lines,
      data.frame(time = tgrid, effect = x$true_beta, curve_type = "True"))
  }
  lvls <- unique(df_lines$curve_type)
  lvls <- c("Estimated", "True")[c("Estimated", "True") %in% lvls]
  df_lines$curve_type <- factor(df_lines$curve_type, levels = lvls)

  cols <- c("Estimated" = "#D62728", "True" = "black")
  ltys <- c("Estimated" = "solid",   "True" = "dashed")

  # CI ribbon (separate data frame, not part of colour scale)
  df_ci <- data.frame(
    time  = tgrid,
    lower = if (!is.null(lo)) lo else NA_real_,
    upper = if (!is.null(hi)) hi else NA_real_
  )

  y_label <- bquote(hat(mu)(.(dn)))

  p <- ggplot2::ggplot(df_lines,
                        ggplot2::aes(x = time, y = effect,
                                     colour   = curve_type,
                                     linetype = curve_type)) +
    ggplot2::geom_ribbon(data = df_ci,
                          ggplot2::aes(x = time, ymin = lower, ymax = upper),
                          inherit.aes = FALSE,
                          fill = "#D62728", alpha = 0.2, na.rm = TRUE) +
    ggplot2::geom_line(linewidth = 1) +
    ggplot2::scale_colour_manual(values = cols[lvls], name = NULL) +
    ggplot2::scale_linetype_manual(values = ltys[lvls], name = NULL) +
    ggplot2::geom_hline(yintercept = 0, linetype = "dotted",
                        colour = "grey50") +
    ggplot2::labs(
      x     = dn,
      y     = y_label,
      title = if (!is.null(ci_label))
        paste0("Causal Effect Function (", ci_label, ")")
      else
        "Causal Effect Function"
    ) +
    ggplot2::theme_bw() +
    ggplot2::theme(
      legend.position  = "top",
      legend.key.width = ggplot2::unit(2, "cm"),
      plot.title       = ggplot2::element_text(hjust = 0.5)
    )
  p
}

#' Surface heatmap for functional outcome
#' @keywords internal
.plot_effect_surface <- function(x, beta, ci_lower = NULL, ci_upper = NULL,
                                  title = "Causal Effect Surface") {
  tgrid <- x$t_grid
  sgrid <- x$outcome_t_grid
  df <- expand.grid(s = tgrid, t = sgrid)
  df$effect <- as.vector(beta)

  p <- ggplot2::ggplot(df, ggplot2::aes(x = s, y = t, fill = effect)) +
    ggplot2::geom_tile() +
    ggplot2::scale_fill_gradient2(
      low      = "#2166AC", mid = "white", high = "#D62728",
      midpoint = 0,
      name     = expression(hat(mu)(s, t))
    ) +
    ggplot2::labs(x = x$domain_name, y = x$outcome_domain_name,
                  title = title) +
    ggplot2::theme_bw() +
    ggplot2::theme(
      plot.title      = ggplot2::element_text(hjust = 0.5, face = "bold"),
      legend.position = "right"
    )

  if (!is.null(x$true_beta)) {
    df_true        <- expand.grid(s = tgrid, t = sgrid)
    df_true$effect <- as.vector(x$true_beta)
    p <- p + ggplot2::geom_contour(
      data        = df_true,
      ggplot2::aes(x = s, y = t, z = effect),
      inherit.aes = FALSE,
      colour      = "black",
      linewidth   = 0.3,
      linetype    = "dashed"
    )
  }
  p
}

#' Bootstrap slice plot for functional outcome (single point) with legend
#' @keywords internal
.plot_bootstrap_slice <- function(x, point, which_domain) {
  which_domain <- match.arg(which_domain, c("treatment", "outcome"))
  tgrid <- x$t_grid
  sgrid <- x$outcome_t_grid

  if (which_domain == "treatment") {
    idx        <- which.min(abs(tgrid - point))
    fixed      <- tgrid[idx]
    slice_beta <- x$beta[idx, ]
    slice_lo   <- x$ci_lower[idx, ]
    slice_hi   <- x$ci_upper[idx, ]
    free_grid  <- sgrid
    xlab       <- x$outcome_domain_name
    ttl        <- sprintf("%s = %.2f", x$domain_name, fixed)
    true_slice <- if (!is.null(x$true_beta)) x$true_beta[idx, ] else NULL
  } else {
    idx        <- which.min(abs(sgrid - point))
    fixed      <- sgrid[idx]
    slice_beta <- x$beta[, idx]
    slice_lo   <- x$ci_lower[, idx]
    slice_hi   <- x$ci_upper[, idx]
    free_grid  <- tgrid
    xlab       <- x$domain_name
    ttl        <- sprintf("%s = %.2f", x$outcome_domain_name, fixed)
    true_slice <- if (!is.null(x$true_beta)) x$true_beta[, idx] else NULL
  }

  # Lines data frame for legend
  df_est <- data.frame(domain_val = free_grid, effect = slice_beta,
                        curve_type = "Estimated")
  df_all <- df_est
  if (!is.null(true_slice)) {
    df_all <- rbind(df_all,
      data.frame(domain_val = free_grid, effect = true_slice,
                 curve_type = "True"))
  }
  lvls <- unique(df_all$curve_type)
  lvls <- c("Estimated", "True")[c("Estimated", "True") %in% lvls]
  df_all$curve_type <- factor(df_all$curve_type, levels = lvls)

  # CI ribbon (separate)
  df_ci <- data.frame(domain_val = free_grid,
                       lower = slice_lo, upper = slice_hi)

  cols <- c("Estimated" = "#D62728", "True" = "black")
  ltys <- c("Estimated" = "solid",   "True" = "dashed")

  ggplot2::ggplot(df_all,
                   ggplot2::aes(x = domain_val, y = effect,
                                colour = curve_type, linetype = curve_type)) +
    ggplot2::geom_ribbon(data = df_ci,
                          ggplot2::aes(x = domain_val, ymin = lower,
                                       ymax = upper),
                          inherit.aes = FALSE,
                          fill = "#4393C3", alpha = 0.25, na.rm = TRUE) +
    ggplot2::geom_line(linewidth = 1) +
    ggplot2::scale_colour_manual(values = cols[lvls], name = NULL) +
    ggplot2::scale_linetype_manual(values = ltys[lvls], name = NULL) +
    ggplot2::geom_hline(yintercept = 0, linetype = "dotted",
                        colour = "grey50") +
    ggplot2::labs(x = xlab, y = expression(hat(mu)), title = ttl) +
    ggplot2::theme_bw() +
    ggplot2::theme(
      legend.position  = "top",
      legend.key.width = ggplot2::unit(2, "cm"),
      plot.title       = ggplot2::element_text(hjust = 0.5)
    )
}

#' Significance plot
#' @keywords internal
.plot_significance <- function(x, alpha) {
  if (is.null(x$ci_lower)) {
    stop("Bootstrap CIs not available. Re-run with bootstrap = TRUE.")
  }

  if (x$outcome_type == "functional") {
    tgrid <- x$t_grid
    sgrid <- x$outcome_t_grid
    sig   <- ((x$ci_lower > 0) | (x$ci_upper < 0))
    sig[is.na(sig)] <- FALSE
    df <- expand.grid(s = tgrid, t = sgrid)
    df$significant <- as.logical(as.vector(sig))

    ggplot2::ggplot(df, ggplot2::aes(x = s, y = t, fill = significant)) +
      ggplot2::geom_tile() +
      ggplot2::scale_fill_manual(
        values = c("TRUE" = "#D62728", "FALSE" = "#DEEBF7"),
        labels = c("TRUE"  = paste0("Significant (alpha=", alpha, ")"),
                   "FALSE" = "Not significant"),
        name   = NULL
      ) +
      ggplot2::labs(x = x$domain_name, y = x$outcome_domain_name,
                    title = paste0("Significance Map (alpha = ", alpha, ")")) +
      ggplot2::theme_bw() +
      ggplot2::theme(
        legend.position = "top",
        plot.title      = ggplot2::element_text(hjust = 0.5, face = "bold")
      )
  } else {
    tgrid   <- x$t_grid
    dn      <- x$domain_name
    sig     <- !is.na(x$ci_lower) & !is.na(x$ci_upper) &
               ((x$ci_lower > 0) | (x$ci_upper < 0))
    half_dt <- (tgrid[2] - tgrid[1]) / 2

    nonsig_ivs <- .contiguous_intervals(tgrid, sig)
    if (nrow(nonsig_ivs) > 0) {
      nonsig_ivs$start <- nonsig_ivs$start - half_dt
      nonsig_ivs$end   <- nonsig_ivs$end   + half_dt
    }

    df <- data.frame(
      time        = tgrid,
      effect      = x$beta,
      lower       = x$ci_lower,
      upper       = x$ci_upper,
      effect_uw   = x$beta_unweighted
    )

    # Add true_beta column for legend integration
    has_true <- !is.null(x$true_beta)
    if (has_true) df$effect_true <- x$true_beta

    p <- ggplot2::ggplot(df, ggplot2::aes(x = time)) +
      # 1. Pink CI ribbon (weighted estimate)
      ggplot2::geom_ribbon(ggplot2::aes(ymin = lower, ymax = upper),
                            fill = "#D62728", alpha = 0.25, na.rm = TRUE) +
      # 2. Grey rectangles for NON-significant regions
      {
        if (nrow(nonsig_ivs) > 0) {
          ggplot2::geom_rect(
            data        = nonsig_ivs,
            ggplot2::aes(xmin = start, xmax = end,
                          ymin = -Inf,  ymax = Inf),
            inherit.aes = FALSE,
            fill        = "grey80",
            alpha       = 0.35
          )
        } else {
          list()
        }
      } +
      # 3. Unweighted line
      ggplot2::geom_line(ggplot2::aes(y = effect_uw,
                                       colour = "Unweighted"),
                          linewidth = 0.9) +
      # 4. Weighted line
      ggplot2::geom_line(ggplot2::aes(y = effect,
                                       colour = "Weighted"),
                          linewidth = 0.9)

    # 5. True mu line (in legend via colour aesthetic)
    if (has_true) {
      p <- p + ggplot2::geom_line(
        ggplot2::aes(y = effect_true,
                     colour = "True \u03bc"),
        linetype  = "dashed",
        linewidth = 0.8
      )
    }

    col_vals <- c("Unweighted" = "#4393C3", "Weighted" = "#D62728")
    if (has_true) col_vals["True \u03bc"] <- "black"

    lty_vals <- c("Unweighted" = "solid", "Weighted" = "solid")
    if (has_true) lty_vals["True \u03bc"] <- "dashed"

    p +
      ggplot2::geom_hline(yintercept = 0, linetype = "dashed",
                           colour = "grey40", linewidth = 0.5) +
      ggplot2::scale_colour_manual(values = col_vals, name = NULL) +
      ggplot2::scale_linetype_manual(values = lty_vals, name = NULL) +
      ggplot2::labs(
        x     = dn,
        y     = bquote(hat(mu)(.(dn))),
        title = paste0(
          "Pointwise Confidence Intervals with Significant Regions",
          " at Level  \u03b1 = ", alpha
        )
      ) +
      ggplot2::theme_bw() +
      ggplot2::theme(
        legend.position  = "right",
        legend.key.width = ggplot2::unit(2, "cm"),
        plot.title       = ggplot2::element_text(size = 10)
      )
  }
}