R/plot.R

In coresynth: Fast and Unified Synthetic Control Methods

# -- Internal style-merging helpers ------------------------------------------
#
# geom_vline()/geom_hline() cannot be removed from a ggplot object once added
# (only overplotted), so suppression and restyling of reference lines must
# happen inside the plot method, before the layer is ever added.

# Merge user overrides onto a default aesthetic list for a reference line.
# override = NULL or FALSE suppresses the line entirely (returns NULL).
# override = list(...) is merged onto default via modifyList (unset keys keep
# their default value).
.line_style <- function(default, override) {
  if (is.null(override) || isFALSE(override)) return(NULL)
  if (!is.list(override))
    stop("vline/hline must be a list of aesthetic overrides ",
         "(e.g. list(color = \"red\")), or NULL/FALSE to hide the line.",
         call. = FALSE)
  utils::modifyList(default, override)
}

# Merge a user-supplied named vector (colors, labels, ...) onto the package
# default, keeping any series name the user didn't mention at its default
# value. Keys are always the canonical series names, so colors and labels can
# be overridden independently.
.merge_named_vec <- function(default, override, what = "colors",
                             example = "c(treated = \"black\")") {
  if (is.null(override)) return(default)
  if (is.null(names(override)) || !all(nzchar(names(override))))
    stop(what, " must be a named vector, e.g. ", example, ". ",
         "Valid names: ", paste(names(default), collapse = ", "), ".",
         call. = FALSE)
  unknown <- setdiff(names(override), names(default))
  if (length(unknown) > 0)
    stop(what, " has unrecognized name(s): ", paste(unknown, collapse = ", "),
         ". Valid names: ", paste(names(default), collapse = ", "), ".",
         call. = FALSE)
  default[names(override)] <- override
  default
}

# `vline_offset` counts periods relative to the first post-treatment period
# (offset 0, the historical fixed position of the treatment line).
.check_vline_offset <- function(offset) {
  if (!is.numeric(offset) || length(offset) != 1L || !is.finite(offset))
    stop("`vline_offset` must be a single finite number of periods relative ",
         "to the first post-treatment period (e.g. -1 = one period earlier).",
         call. = FALSE)
}

# Resolve where the treatment vline is drawn: -1 is the last pre-treatment
# period, and fractional offsets interpolate linearly between adjacent
# observed times, so -0.5 lands midway on numeric, Date, and POSIXct axes
# alike. Out-of-range offsets return NA (line skipped), matching how a fit
# without post-treatment periods has always been handled.
.vline_position <- function(times, T_pre, offset) {
  idx <- T_pre + 1 + offset
  if (idx < 1 || idx > length(times)) {
    if (offset != 0)
      warning("`vline_offset = ", offset, "` places the treatment line ",
              "outside the observed time range; the line is not drawn.",
              call. = FALSE)
    return(times[NA_integer_])
  }
  lo   <- floor(idx)
  frac <- idx - lo
  if (frac == 0) return(times[lo])
  times[lo] + frac * (times[lo + 1L] - times[lo])
}

# Split a user-supplied absolute position (`xintercept`) out of the merged
# vline style list, so `vline = list(xintercept = ...)` places the line
# instead of the treatment time. Character positions are coerced to the axis
# class on Date/POSIXct axes, letting users write "1989-01-01" directly.
.vline_split <- function(style, times, offset) {
  if (is.null(style) || !("xintercept" %in% names(style)))
    return(list(style = style, at = NULL))
  if (offset != 0)
    stop("supply either `vline_offset` or `vline = list(xintercept = ...)`, ",
         "not both.", call. = FALSE)
  at <- style$xintercept
  style$xintercept <- NULL
  if (is.character(at)) {
    if (inherits(times, "Date")) at <- as.Date(at)
    else if (inherits(times, "POSIXct"))
      at <- as.POSIXct(at, tz = attr(times, "tzone") %||% "")
  }
  if (length(at) < 1L || anyNA(at))
    stop("`vline` xintercept must be one or more non-missing positions on ",
         "the time axis.", call. = FALSE)
  list(style = style, at = at)
}

.merge_named_colors <- function(default, override) {
  .merge_named_vec(default, override, what = "colors")
}

# Legend labels default to the display series names. `key_map` maps the
# user-facing one-word keys (names) to the display names (values), so users
# address series by identifier while the plot keeps its full legend text.
.merge_named_labels <- function(key_map, override) {
  .merge_named_vec(key_map, override,
                   what = "labels", example = "c(treated = \"California\")")
}

# Line types may be given by name ("solid") or by the integer codes 0:6 that
# `par(lty)` uses. Codes are translated to names before merging, because the
# merged vector is character and ggplot2 reads a bare "2" as a hex dash
# pattern rather than as "dashed".
.linetype_names <- c("blank", "solid", "dashed", "dotted", "dotdash",
                     "longdash", "twodash")

.as_linetype <- function(x, what = "linetypes") {
  if (is.null(x) || is.character(x)) return(x)
  if (!is.numeric(x) || anyNA(x) || any(x != trunc(x)) || any(x < 0) || any(x > 6))
    stop(what, " must be line type name(s) such as \"solid\" or \"dashed\", ",
         "or integer code(s) in 0:6.", call. = FALSE)
  stats::setNames(.linetype_names[as.integer(x) + 1L], names(x))
}

.merge_named_linetypes <- function(default, override) {
  .merge_named_vec(default, .as_linetype(override), what = "linetypes",
                   example = "c(synthetic = \"solid\")")
}

# Pre-treatment level offset between the treated and synthetic series.
# SDID matches trends only up to a free intercept (omega_0 is concentrated
# out of the QP), so its raw trend plot shows the two series at different
# levels; shifting by the lambda-weighted pre-period gap makes the average
# post-period gap equal the SDID estimate exactly. Other methods already
# match levels, so the plain pre-period mean gap is used.
.align_offset <- function(x, Y_treat, Y_synth) {
  pre   <- seq_len(x$T_pre)
  d_pre <- Y_treat[pre] - Y_synth[pre]
  lam   <- x$time_weights
  if (identical(x$method, "sdid") && !is.null(lam) && length(lam) == x$T_pre)
    sum(lam * d_pre)
  else
    mean(d_pre)
}

#' Plot a coresynth model
#'
#' @param x      A `coresynth` object.
#' @param type   One of `"trend"` (observed vs synthetic), `"gap"` (ATT over time),
#'               `"weights"` (donor unit weight bar chart; SDID fits get a
#'               second panel with the time weight profile), or
#'               `"pred_weights"` (the predictor/variable weight matrix
#'               \eqn{V} as a bar chart; sharp SCM only).
#' @param colors For `type = "trend"`: a named vector overriding series colors,
#'   e.g. `c(treated = "black")` (valid keys: `"treated"`, `"synthetic"`,
#'   plus `"donors"` when `show_donors > 0`).
#'   For `type = "gap"`: a single color string for the gap line. Ignored for
#'   `type = "weights"` (use `fill` instead).
#' @param labels For `type = "trend"`: a named vector overriding the legend
#'   text of individual series, e.g. `c(treated = "California")` (valid keys:
#'   `"treated"`, `"synthetic"`, plus `"donors"` when `show_donors > 0`).
#'   Series not mentioned keep their default label; `colors` and `labels`
#'   address series by the same keys, independent of the displayed legend
#'   text. Ignored for other types (no legend).
#' @param linetypes For `type = "trend"`: a named vector overriding the line
#'   type of individual series, e.g. `c(synthetic = "solid")` to draw both the
#'   treated and the synthetic series solid (valid keys: `"treated"`,
#'   `"synthetic"`, plus `"donors"` when `show_donors > 0`; defaults are
#'   `"solid"` for the treated and donor series and `"dashed"` for the
#'   synthetic one). For `type = "gap"`: a single line type for the gap line
#'   (default `"solid"`). Values are `ggplot2` line type names (`"solid"`,
#'   `"dashed"`, `"dotted"`, `"dotdash"`, `"longdash"`, `"twodash"`,
#'   `"blank"`) or the equivalent integer codes `0:6`. Series are addressed by
#'   the same keys as `colors` and `labels`, so the three compose
#'   independently. Ignored for `type = "weights"` and `"pred_weights"`
#'   (bar charts).
#' @param vline  Aesthetic overrides for the vertical treatment-time line, as a
#'   list passed to [ggplot2::geom_vline()] (e.g. `list(color = "red")`).
#'   `NULL` or `FALSE` hides the line entirely. The list may also carry an
#'   `xintercept` element giving one or more absolute positions on the time
#'   axis (e.g. `list(xintercept = 1988)`, or a date string such as
#'   `"1989-01-01"` on a `Date` axis), replacing the default treatment-time
#'   position; for placement relative to the treatment period use
#'   `vline_offset` instead. Applies to `"trend"` and `"gap"`.
#' @param vline_offset For `type = "trend"` and `"gap"`: where to draw the
#'   vertical treatment line, in periods relative to the first post-treatment
#'   period. The default `0` keeps the line at the first post-treatment
#'   period; `-1` moves it to the last pre-treatment period, and fractional
#'   values interpolate between adjacent observed times (`-0.5` is midway
#'   between the last pre- and first post-treatment period), which works on
#'   numeric, `Date`, and `POSIXct` axes alike. Cannot be combined with an
#'   `xintercept` element in `vline`.
#' @param hline  Aesthetic overrides for the horizontal zero line in `type =
#'   "gap"`, as a list passed to [ggplot2::geom_hline()]. `NULL` or `FALSE`
#'   hides the line. Ignored for other types.
#' @param fill   For `type = "weights"` and `type = "pred_weights"`: a
#'   single color string overriding the bar fill. Ignored for other types.
#' @param top_n  For `type = "weights"`: show only the `top_n` donors with the
#'   largest weights (default `Inf` keeps every donor with a non-negligible
#'   weight). For `type = "pred_weights"`: show only the `top_n`
#'   predictors with the largest \eqn{V} weights (default `Inf` shows every
#'   predictor). Ignored for other types (and for the SDID time weight panel,
#'   which is always shown in full).
#' @param align  For `type = "trend"` and `"gap"`: when `TRUE`, shift the
#'   synthetic series by its pre-treatment level gap to the treated series so
#'   that both are drawn on the same level. SDID matches trends only up to a
#'   free intercept, so its raw trend plot shows the synthetic control at a
#'   different level; with `align = TRUE` the offset is the time-weighted
#'   (\eqn{\lambda}) pre-period gap, which makes the average post-period gap
#'   in the plot equal the SDID estimate exactly. Other methods use the plain
#'   pre-period mean gap (usually a negligible shift, since they already
#'   match levels). Default `FALSE` (raw series).
#' @param show_donors For `type = "trend"`: also draw the outcome paths of the
#'   `show_donors` donor units with the largest weights as thin background
#'   lines (`Inf` shows every donor). Requires a fit that stores donor unit
#'   weights and outcomes (sharp SCM/SDID/SI). Default `0` (no donor paths).
#' @param ...    Ignored.
#' @return A `ggplot2` plot object.
#' @examples
#' set.seed(1)
#' panel <- expand.grid(unit = 1:10, year = 1:20)
#' panel$treated <- as.integer(panel$unit == 5 & panel$year > 15)
#' panel$gdp <- panel$unit + 0.5 * panel$year +
#'   rnorm(nrow(panel)) + 3 * panel$treated
#' fit <- scm_fit(gdp ~ treated | unit + year, data = panel, method = "scm")
#'
#' \donttest{
#' plot(fit, type = "trend")
#' plot(fit, type = "gap")
#' plot(fit, type = "weights")
#' plot(fit, type = "weights", top_n = 5)
#'
#' # Predictor (V) weights: which predictors the fit leans on
#' plot(fit, type = "pred_weights")
#'
#' # Overlay the five largest donors behind the treated/synthetic series
#' plot(fit, type = "trend", show_donors = 5)
#'
#' # SDID: align the synthetic series on the lambda-weighted pre-period level
#' fit_sdid <- scm_fit(gdp ~ treated | unit + year, data = panel, method = "sdid")
#' plot(fit_sdid, type = "trend", align = TRUE)
#' plot(fit_sdid, type = "gap",   align = TRUE)
#'
#' # Customize series colors, legend text, line types, and reference lines
#' plot(fit, type = "trend",
#'      colors = c(treated = "black"),
#'      labels = c(treated = "Unit 5"),
#'      vline  = list(color = "red", linetype = "dashed"))
#'
#' # Draw both series solid, or restyle the gap line
#' plot(fit, type = "trend", linetypes = c(synthetic = "solid"))
#' plot(fit, type = "gap",   linetypes = "dashed")
#'
#' # Move the treatment line one period earlier (last pre-treatment period),
#' # pin it to an absolute time, or drop it entirely
#' plot(fit, type = "trend", vline_offset = -1)
#' plot(fit, type = "trend", vline = list(xintercept = 15.5))
#' plot(fit, type = "trend", vline = FALSE)
#' }
#' @import ggplot2
#' @export
plot.coresynth <- function(x, type = c("trend", "gap", "weights",
                                       "pred_weights"),
                            colors = NULL, labels = NULL, linetypes = NULL,
                            vline = list(), vline_offset = 0, hline = list(),
                            fill = NULL, top_n = Inf,
                            align = FALSE, show_donors = 0, ...) {
  type <- match.arg(type)
  if (!isTRUE(align) && !isFALSE(align))
    stop("`align` must be TRUE or FALSE.", call. = FALSE)
  if (!is.numeric(show_donors) || length(show_donors) != 1L ||
      is.na(show_donors) || show_donors < 0)
    stop("`show_donors` must be a single number >= 0 (Inf shows all donors).",
         call. = FALSE)

  if(type %in% c("trend", "gap")) {
    if(is.null(x$times) || is.null(x$Y_treat))
      stop("fit object does not contain time series data for plotting.")

    times <- x$times
    # Coerce only types ggplot cannot place on a continuous/date axis; pass
    # Date/POSIXct through so the appropriate date scale is selected automatically.
    if (is.character(times) || is.factor(times))
      times <- as.numeric(as.character(times))
    # Multiple treated units are averaged per period by the accessors
    Y_treat <- treated_outcomes(x, na.rm = TRUE)
    Y_synth <- synthetic_outcomes(x, na.rm = TRUE)
    if(is.null(Y_synth))
      stop("fit object does not contain a synthetic/counterfactual series ",
           "to plot (staggered fits store their series per cohort).",
           call. = FALSE)
    if (isTRUE(align)) {
      if (is.null(x$T_pre) || x$T_pre < 1L)
        stop("align = TRUE requires a fit with pre-treatment periods.",
             call. = FALSE)
      Y_synth <- Y_synth + .align_offset(x, Y_treat, Y_synth)
    }
    .check_vline_offset(vline_offset)
    vline_style <- .line_style(list(color = "gray40", linetype = "dotted"), vline)
    vl          <- .vline_split(vline_style, times, vline_offset)
    vline_style <- vl$style
    treat_time  <- if (!is.null(vl$at)) vl$at
                   else if (!is.null(x$T_pre)) .vline_position(times, x$T_pre, vline_offset)
                   else NA
    is_sdid     <- identical(x$method, "sdid")

    if(type == "trend") {
      df_donors <- NULL
      if (show_donors >= 1) {
        w    <- x$unit_weights
        Y_co <- donor_outcomes(x)
        if (is.null(w) || all(is.na(w)) || is.null(Y_co))
          stop("show_donors requires donor unit weights and outcomes ",
               "(available for sharp SCM/SDID/SI fits).", call. = FALSE)
        k   <- as.integer(min(length(w), show_donors))
        sel <- order(w, decreasing = TRUE)[seq_len(k)]
        donor_names <- (names(w) %||% sprintf("Donor %d", seq_along(w)))[sel]
        df_donors <- data.frame(
          time  = rep(times, times = k),
          value = as.vector(Y_co[, sel, drop = FALSE]),
          unit  = rep(donor_names, each = length(times))
        )
      }

      key_map        <- c(treated = "Treated", synthetic = "Synthetic Control")
      color_defaults <- c(treated = "#2166ac", synthetic = "#d73027")
      ltype_defaults <- c(treated = "solid", synthetic = "dashed")
      if (!is.null(df_donors)) {
        key_map        <- c(key_map, donors = "Donors")
        color_defaults <- c(color_defaults, donors = "grey70")
        ltype_defaults <- c(ltype_defaults, donors = "solid")
      }
      series <- unname(key_map)
      # user overrides are keyed by the one-word identifiers; the ggplot scales
      # need the display names the data frame carries
      series_colors <- stats::setNames(
        .merge_named_colors(color_defaults, colors), series)
      series_ltypes <- stats::setNames(
        .merge_named_linetypes(ltype_defaults, linetypes), series)
      series_labels <- stats::setNames(
        .merge_named_labels(key_map, labels), series)
      # Without donors the default (alphabetical) legend order is kept for
      # backwards compatibility; with donors the breaks pin them last.
      series_breaks <- if (!is.null(df_donors)) series else waiver()
      df <- data.frame(
        time     = c(times, times),
        value    = c(Y_treat, Y_synth),
        series   = rep(c("Treated", "Synthetic Control"), each = length(times))
      )
      # labels must be identical on both scales or the merged legend splits in two
      p <- ggplot(df, aes(x = time, y = value, color = series, linetype = series)) +
        {if(!is.null(df_donors)) geom_line(
          data = df_donors,
          aes(x = time, y = value, group = unit,
              color = "Donors", linetype = "Donors"),
          linewidth = 0.4, alpha = 0.5, inherit.aes = FALSE)} +
        geom_line(linewidth = 0.9) +
        scale_color_manual(values = series_colors, breaks = series_breaks,
                           labels = series_labels) +
        scale_linetype_manual(values = series_ltypes, breaks = series_breaks,
                              labels = series_labels) +
        {if(!is.null(vline_style) && !anyNA(treat_time)) do.call(geom_vline, c(list(xintercept = treat_time), vline_style))} +
        theme_minimal(base_size = 13) +
        labs(title    = paste0("Synthetic Control Trend  [", toupper(x$method), "]"),
             subtitle = if (isTRUE(align)) {
               if (is_sdid)
                 "Synthetic control shifted by the lambda-weighted pre-period gap"
               else
                 "Synthetic control shifted by the mean pre-period gap"
             },
             x = "Time", y = "Outcome", color = NULL, linetype = NULL)
      return(p)
    }

    if(type == "gap") {
      gap_color   <- if (is.null(colors)) "#1a9641" else unname(colors[[1]])
      gap_ltype   <- if (is.null(linetypes)) "solid"
                     else unname(.as_linetype(linetypes)[[1]])
      hline_style <- .line_style(list(color = "gray50", linetype = "dashed"), hline)
      gap <- Y_treat - Y_synth
      df  <- data.frame(time = times, gap = gap)
      subtitle <- if (isTRUE(align)) {
        if (is_sdid)
          "Treated - synthetic control, lambda-aligned on the pre-period\n(post-period mean = SDID estimate)"
        else
          "Treated - synthetic control, aligned on the pre-period mean"
      } else {
        "Treated - synthetic control"
      }
      p <- ggplot(df, aes(x = time, y = gap)) +
        geom_line(color = gap_color, linetype = gap_ltype, linewidth = 0.9) +
        {if(!is.null(hline_style)) do.call(geom_hline, c(list(yintercept = 0), hline_style))} +
        {if(!is.null(vline_style) && !anyNA(treat_time)) do.call(geom_vline, c(list(xintercept = treat_time), vline_style))} +
        theme_minimal(base_size = 13) +
        labs(title    = paste0("Treatment Effect Gap  [", toupper(x$method), "]"),
             subtitle = subtitle,
             x = "Time", y = "Gap")
      return(p)
    }
  }

  if(type == "weights") {
    w <- x$unit_weights
    if(is.null(w) || all(is.na(w)))
      stop("No unit weights available for this method (GSC/MC/TASC use factor loadings).")

    bar_fill <- fill %||% "#4575b4"
    df <- data.frame(
      unit   = names(w) %||% paste0("Unit_", seq_along(w)),
      weight = as.numeric(w)
    )
    df <- df[df$weight > 1e-4, ]
    if(nrow(df) == 0) stop("All unit weights are negligibly small.")
    if(!is.numeric(top_n) || length(top_n) != 1L || is.na(top_n) || top_n < 1)
      stop("`top_n` must be a single number >= 1 (Inf shows all donors).")
    if(is.finite(top_n) && nrow(df) > top_n)
      df <- df[order(df$weight, decreasing = TRUE)[seq_len(top_n)], ]

    lam <- x$time_weights
    two_panel <- identical(x$method, "sdid") && !is.null(lam) &&
      !is.null(x$T_pre) && length(lam) == x$T_pre

    if (!two_panel) {
      p <- ggplot(df, aes(x = reorder(unit, weight), y = weight)) +
        geom_col(fill = bar_fill, alpha = 0.85) +
        coord_flip() +
        theme_minimal(base_size = 13) +
        labs(title = "Donor Unit Weights", x = NULL, y = "Weight")
      return(p)
    }

    # SDID: the estimator is defined by unit weights (omega) AND time weights
    # (lambda); show both, side by side, in the same bar style.
    t_lab <- as.character((x$times %||% seq_len(x$T_pre))[seq_len(x$T_pre)])
    df_t  <- data.frame(label = t_lab, weight = as.numeric(lam),
                        stringsAsFactors = FALSE)
    df_t  <- df_t[df_t$weight > 1e-4, ]
    # A pre-period label that collides with a donor name would duplicate the
    # shared factor levels below
    if (any(df_t$label %in% df$unit))
      df_t$label <- paste0("t = ", df_t$label)

    panel_u <- "Unit weights (omega)"
    panel_t <- "Time weights (lambda)"
    df_u <- data.frame(label = df$unit, weight = df$weight, panel = panel_u,
                       stringsAsFactors = FALSE)
    df_t$panel <- panel_t
    both <- rbind(df_u, df_t)
    # Units ordered by weight (largest on top), pre-periods chronologically
    # (earliest on top); the panel factor keeps unit weights on the left.
    both$label <- factor(both$label,
                         levels = c(df_u$label[order(df_u$weight)],
                                    rev(df_t$label)))
    both$panel <- factor(both$panel, levels = c(panel_u, panel_t))

    p <- ggplot(both, aes(x = weight, y = label)) +
      geom_col(fill = bar_fill, alpha = 0.85) +
      facet_wrap(~ panel, scales = "free_y") +
      theme_minimal(base_size = 13) +
      labs(title = "Donor Unit and Time Weights  [SDID]",
           x = "Weight", y = NULL)
    return(p)
  }

  if(type == "pred_weights") {
    v <- x$v_weights
    if(is.null(v) || all(is.na(v)))
      stop("No predictor (V) weights available. A V matrix is estimated only ",
           "by sharp SCM fits; staggered SCM and the other methods ",
           "(SDID/GSC/MC/TASC/SI) do not produce one.", call. = FALSE)
    if(!is.numeric(top_n) || length(top_n) != 1L || is.na(top_n) || top_n < 1)
      stop("`top_n` must be a single number >= 1 (Inf shows all predictors).")

    bar_fill <- fill %||% "#4575b4"
    df <- data.frame(
      predictor = names(v) %||% paste0("V", seq_along(v)),
      weight    = as.numeric(v)
    )
    if(is.finite(top_n) && nrow(df) > top_n)
      df <- df[order(df$weight, decreasing = TRUE)[seq_len(top_n)], ]

    p <- ggplot(df, aes(x = reorder(predictor, weight), y = weight)) +
      geom_col(fill = bar_fill, alpha = 0.85) +
      coord_flip() +
      theme_minimal(base_size = 13) +
      labs(title = "Predictor Weights", x = NULL, y = "Weight")
    return(p)
  }
}

#' Plot SCM In-Space Placebo Results
#'
#' Visualizes the placebo study returned by [mspe_ratio_pval()], following
#' Abadie, Diamond & Hainmueller (2010, Section 3.4).
#'
#' `type = "gaps"` overlays the treated unit's gap path (treated minus
#' synthetic control) on the placebo gap paths obtained by reassigning the
#' intervention to each donor unit (ADH 2010, Figure 4). Placebo units whose
#' synthetic control fits poorly before treatment carry no information about
#' the rarity of a large post-treatment gap, so ADH exclude units whose
#' pre-treatment MSPE exceeds a multiple of the treated unit's: 20, 5, and 2
#' in their Figures 5-7 (`mspe_prune`).
#'
#' `type = "ratios"` shows the post/pre-treatment MSPE ratio of every unit
#' (ADH 2010, Figure 8), the statistic behind the two-sided permutation
#' p-value; it requires no pruning cutoff by construction.
#'
#' @param x A `scm_placebo` object from [mspe_ratio_pval()].
#' @param type One of `"gaps"` (ADH 2010, Figures 4-7) or `"ratios"`
#'   (ADH 2010, Figure 8).
#' @param mspe_prune Only for `type = "gaps"`: exclude placebo units whose
#'   pre-treatment MSPE exceeds `mspe_prune` times the treated unit's.
#'   Default `Inf` (no pruning). A rule stated on the RMSPE scale, such as
#'   tidysynth's "2 times the treated unit's pre-period RMSPE", corresponds
#'   to the squared multiple (`mspe_prune = 4`).
#' @param colors A named vector overriding series colors, e.g.
#'   `c(treated = "black")`. Valid keys: `"treated"`, `"placebo"`.
#' @param labels A named vector overriding the legend text of individual
#'   series, e.g. `c(treated = "California")`. Valid keys: `"treated"`,
#'   `"placebo"`. Series not mentioned keep their default label; `colors`
#'   and `labels` address series by the same keys, independent of the
#'   displayed legend text.
#' @param linetypes Only for `type = "gaps"`: a named vector overriding the
#'   line type of individual series, e.g. `c(placebo = "dotted")`. Valid keys:
#'   `"treated"`, `"placebo"` (both default to `"solid"`). Values are
#'   `ggplot2` line type names (`"solid"`, `"dashed"`, `"dotted"`,
#'   `"dotdash"`, `"longdash"`, `"twodash"`, `"blank"`) or the equivalent
#'   integer codes `0:6`. Series are addressed by the same keys as `colors`
#'   and `labels`. Ignored for `type = "ratios"` (points, not lines).
#' @param vline Only for `type = "gaps"`: aesthetic overrides for the vertical
#'   treatment-time line, as a list passed to [ggplot2::geom_vline()].
#'   `NULL` or `FALSE` hides the line entirely. The list may also carry an
#'   `xintercept` element giving one or more absolute positions on the time
#'   axis, replacing the default treatment-time position.
#' @param vline_offset Only for `type = "gaps"`: where to draw the vertical
#'   treatment line, in periods relative to the first post-treatment period.
#'   The default `0` keeps the line at the first post-treatment period; `-1`
#'   moves it to the last pre-treatment period, and fractional values
#'   interpolate between adjacent observed times. Cannot be combined with an
#'   `xintercept` element in `vline`.
#' @param hline Only for `type = "gaps"`: aesthetic overrides for the
#'   horizontal zero line, as a list passed to [ggplot2::geom_hline()].
#'   `NULL` or `FALSE` hides the line entirely.
#' @param ... Ignored.
#' @return A `ggplot2` plot object.
#' @examples
#' set.seed(1)
#' panel <- expand.grid(unit = 1:10, year = 1:20)
#' panel$treated <- as.integer(panel$unit == 5 & panel$year > 15)
#' panel$gdp <- panel$unit + 0.5 * panel$year +
#'   rnorm(nrow(panel)) + 3 * panel$treated
#' fit <- scm_fit(gdp ~ treated | unit + year, data = panel, method = "scm")
#' placebo <- mspe_ratio_pval(fit)
#'
#' \donttest{
#' # Treated gap overlaid on the donor-pool placebo gaps (ADH 2010, Fig. 4)
#' plot(placebo, type = "gaps")
#'
#' # Prune poorly fitting placebos and relabel the legend
#' plot(placebo, type = "gaps", mspe_prune = 5,
#'      labels = c(treated = "Unit 5"))
#'
#' # Set the line type of the placebo paths off against the treated one
#' plot(placebo, type = "gaps", linetypes = c(placebo = "dotted"))
#'
#' # Move the treatment line one period earlier
#' plot(placebo, type = "gaps", vline_offset = -1)
#'
#' # Post/pre-treatment MSPE ratios (ADH 2010, Fig. 8)
#' plot(placebo, type = "ratios")
#' }
#' @seealso [mspe_ratio_pval()]
#' @export
plot.scm_placebo <- function(x, type = c("gaps", "ratios"), mspe_prune = Inf,
                              colors = NULL, labels = NULL, linetypes = NULL,
                              vline = list(), vline_offset = 0,
                              hline = list(), ...) {
  type <- match.arg(type)
  key_map <- c(treated = "Treated", placebo = "Placebo (donor pool)")
  series_labels <- stats::setNames(
    .merge_named_labels(key_map, labels), unname(key_map))
  if (!is.numeric(mspe_prune) || length(mspe_prune) != 1L || mspe_prune <= 0)
    stop("mspe_prune must be a single positive number (Inf = no pruning).")

  subtitle <- paste0(
    "Permutation p-value = ", formatC(x$p_value, digits = 3, format = "g"),
    " (", x$alternative, ", ", x$n_placebo_used, " placebo units)"
  )

  if (type == "gaps") {
    times <- x$times
    if (is.character(times) || is.factor(times))
      times <- as.numeric(as.character(times))
    .check_vline_offset(vline_offset)

    keep <- is.finite(x$mspe_pre_placebo) &
      x$mspe_pre_placebo <= mspe_prune * x$mspe_pre_treated
    n_pruned <- sum(!keep)
    if (!any(keep))
      warning("All placebo units were pruned; only the treated gap is shown. ",
              "Consider a larger mspe_prune.")

    gaps  <- x$gaps[, keep, drop = FALSE]
    # sprintf keeps zero-length input zero-length (paste0 would collapse it to "Donor ")
    units <- colnames(gaps) %||% sprintf("Donor %d", which(keep))
    df_pl <- data.frame(
      time = rep(times, times = ncol(gaps)),
      gap  = as.vector(gaps),
      unit = rep(units, each = length(times))
    )
    df_tr <- data.frame(time = times, gap = x$treated_gap)

    series_colors <- stats::setNames(.merge_named_colors(
      c(treated = "#2166ac", placebo = "grey70"), colors
    ), unname(key_map))
    # both series are solid by default, so the linetype scale is a no-op until
    # the user overrides it; it must carry the same breaks/labels as the color
    # scale or the merged legend splits in two
    series_ltypes <- stats::setNames(.merge_named_linetypes(
      c(treated = "solid", placebo = "solid"), linetypes
    ), unname(key_map))
    vline_style <- .line_style(list(color = "gray40", linetype = "dotted"), vline)
    hline_style <- .line_style(list(color = "gray50", linetype = "dashed"), hline)
    vl          <- .vline_split(vline_style, times, vline_offset)
    vline_style <- vl$style
    treat_time  <- vl$at %||% .vline_position(times, x$T_pre, vline_offset)

    p <- ggplot() +
      geom_line(data = df_pl,
                aes(x = time, y = gap, group = unit,
                    color = "Placebo (donor pool)",
                    linetype = "Placebo (donor pool)"),
                linewidth = 0.4, alpha = 0.8) +
      geom_line(data = df_tr, aes(x = time, y = gap, color = "Treated",
                                  linetype = "Treated"),
                linewidth = 1.0) +
      {if(!is.null(hline_style)) do.call(geom_hline, c(list(yintercept = 0), hline_style))} +
      {if(!is.null(vline_style) && !anyNA(treat_time)) do.call(geom_vline, c(list(xintercept = treat_time), vline_style))} +
      scale_color_manual(values = series_colors, breaks = names(series_colors),
                         labels = series_labels) +
      scale_linetype_manual(values = series_ltypes, breaks = names(series_ltypes),
                            labels = series_labels) +
      theme_minimal(base_size = 13) +
      labs(title = "Placebo Gaps in the Donor Pool  [SCM]",
           subtitle = subtitle,
           x = "Time", y = "Gap",
           color = NULL, linetype = NULL,
           caption = if (n_pruned > 0L) {
             paste0("Pruned ", n_pruned, " placebo unit(s) with pre-treatment MSPE > ",
                    mspe_prune, "x the treated unit's.")
           })
    return(p)
  }

  # type == "ratios" (ADH 2010, Figure 8)
  r     <- x$mspe_ratios_all
  units <- names(r)
  if (is.null(units)) units <- c("Treated", sprintf("Donor %d", seq_len(length(r) - 1L)))
  # keep the axis tick consistent with the (possibly relabeled) legend entry
  units[1L] <- series_labels[["Treated"]]
  blank <- !nzchar(units)
  units[blank] <- sprintf("Donor %d", which(blank) - 1L)

  df <- data.frame(unit = units, ratio = as.numeric(r),
                   series = c("Treated", rep("Placebo (donor pool)", length(r) - 1L)))
  n_dropped <- sum(!is.finite(df$ratio))
  df <- df[is.finite(df$ratio), ]
  if (nrow(df) == 0L) stop("No finite MSPE ratios to plot.")

  series_colors <- stats::setNames(.merge_named_colors(
    c(treated = "#2166ac", placebo = "grey60"), colors
  ), unname(key_map))

  p <- ggplot(df, aes(x = ratio, y = reorder(unit, ratio), color = series)) +
    geom_point(size = 2.5) +
    scale_color_manual(values = series_colors, breaks = names(series_colors),
                       labels = series_labels) +
    theme_minimal(base_size = 13) +
    labs(title = "Post/Pre-Treatment MSPE Ratios  [SCM]",
         subtitle = subtitle,
         x = "MSPE ratio (post / pre)", y = NULL, color = NULL,
         caption = if (n_dropped > 0L) {
           paste0(n_dropped, " unit(s) without a finite ratio (mspe_threshold filter) omitted.")
         })
  p
}

#' Extract the tidy data behind a coresynth plot
#'
#' Returns the tidy `data.frame` that [plot()] draws for a given `type`, so the
#' underlying series, weights, or placebo paths can be inspected, joined into a
#' table, or re-plotted directly. [plot()] stays the quick path; `plot_data()`
#' is the handle for anyone who wants to relabel simplified series names, feed
#' the numbers into their own figure, or postprocess them further.
#'
#' The frame mirrors what the matching `plot(x, type = ...)` call shows, with
#' two deliberate departures that make it a better data source:
#'
#' * Plain column names (`time`, `value`, `series`, `weight`, ...) are used
#'   instead of the dotted convention of [augment()], since this is data to
#'   manipulate rather than model-augmented observations.
#' * The cosmetic "drop donors with weight below 1e-4" filter that
#'   `plot(type = "weights")` applies is *not* used here: every donor is
#'   returned (use `top_n` to subset), so the frame is the complete set of
#'   weights.
#'
#' Only the arguments that change *which rows or values* appear are accepted
#' (`align`, `top_n`, `show_donors`, `mspe_prune`); purely cosmetic arguments
#' of [plot()] (`colors`, `labels`, `linetypes`, `vline`, `fill`, ...) have no
#' data counterpart and are not part of this interface.
#'
#' @param x A `coresynth` fit (from [scm_fit()]) or a `scm_placebo` object
#'   (from [mspe_ratio_pval()]).
#' @param type For a `coresynth` fit: one of `"trend"`, `"gap"`, `"weights"`,
#'   or `"pred_weights"`, matching [plot.coresynth()]. For a `scm_placebo`
#'   object: `"gaps"` or `"ratios"`, matching [plot.scm_placebo()].
#' @param align For `type = "trend"`/`"gap"`: when `TRUE`, shift the synthetic
#'   series by its pre-treatment level gap to the treated series, exactly as in
#'   [plot.coresynth()]. Default `FALSE` (raw series).
#' @param top_n For `type = "weights"`/`"pred_weights"`: keep only the `top_n`
#'   largest weights (default `Inf`, every row).
#' @param show_donors For `type = "trend"`: also return the outcome paths of the
#'   `show_donors` largest-weight donors as `series = "Donors"` rows, adding a
#'   `unit` column that identifies each donor (`NA` for the treated and
#'   synthetic series). Default `0` (treated and synthetic series only).
#' @param mspe_prune For a `scm_placebo` object with `type = "gaps"`: drop
#'   placebo units whose pre-treatment MSPE exceeds `mspe_prune` times the
#'   treated unit's, as in [plot.scm_placebo()]. Default `Inf` (no pruning).
#' @param ... Passed to methods (unused by the current methods).
#' @return A tidy `data.frame`. Columns by `type`:
#'   \describe{
#'     \item{`"trend"`}{`time`, `value`, `series` (`"Treated"` /
#'       `"Synthetic Control"`); with `show_donors > 0`, also `"Donors"` rows
#'       and a `unit` column.}
#'     \item{`"gap"`}{`time`, `gap` (treated minus synthetic control).}
#'     \item{`"weights"`}{`unit`, `weight`; SDID fits add a `panel` column
#'       (`"omega"` unit weights, `"lambda"` time weights), with `unit` holding
#'       the pre-period label for `"lambda"` rows.}
#'     \item{`"pred_weights"`}{`predictor`, `weight` (sharp SCM only).}
#'     \item{`"gaps"`}{`time`, `gap`, `unit` (`NA` for the treated series),
#'       `series` (`"Treated"` / `"Placebo (donor pool)"`).}
#'     \item{`"ratios"`}{`unit`, `ratio`, `series`.}
#'   }
#' @seealso [plot.coresynth()], [plot.scm_placebo()]
#' @examples
#' set.seed(1)
#' panel <- expand.grid(unit = 1:10, year = 1:20)
#' panel$treated <- as.integer(panel$unit == 5 & panel$year > 15)
#' panel$gdp <- panel$unit + 0.5 * panel$year +
#'   rnorm(nrow(panel)) + 3 * panel$treated
#' fit <- scm_fit(gdp ~ treated | unit + year, data = panel, method = "scm")
#'
#' head(plot_data(fit, type = "trend"))
#' plot_data(fit, type = "gap")
#' plot_data(fit, type = "weights")
#'
#' # Relabel the simplified series names, then plot it yourself
#' df <- plot_data(fit, type = "trend")
#' df$series <- sub("Synthetic Control", "Synthetic Unit 5", df$series)
#' \donttest{
#' ggplot2::ggplot(df, ggplot2::aes(time, value, color = series)) +
#'   ggplot2::geom_line()
#' }
#' @export
plot_data <- function(x, ...) UseMethod("plot_data")

#' @rdname plot_data
#' @export
plot_data.default <- function(x, ...) {
  stop("plot_data() has no method for an object of class ",
       paste(class(x), collapse = "/"),
       ". It supports coresynth fits (from scm_fit()) and scm_placebo ",
       "objects (from mspe_ratio_pval()).", call. = FALSE)
}

#' @rdname plot_data
#' @export
plot_data.coresynth <- function(x, type = c("trend", "gap", "weights",
                                            "pred_weights"),
                                align = FALSE, top_n = Inf,
                                show_donors = 0, ...) {
  type <- match.arg(type)

  if (type %in% c("trend", "gap")) {
    if (is.null(x$times) || is.null(x$Y_treat))
      stop("fit object does not contain time series data.", call. = FALSE)
    if (!isTRUE(align) && !isFALSE(align))
      stop("`align` must be TRUE or FALSE.", call. = FALSE)

    times <- x$times
    if (is.character(times) || is.factor(times))
      times <- as.numeric(as.character(times))
    Y_treat <- treated_outcomes(x, na.rm = TRUE)
    Y_synth <- synthetic_outcomes(x, na.rm = TRUE)
    if (is.null(Y_synth))
      stop("fit object does not contain a synthetic/counterfactual series ",
           "(staggered fits store their series per cohort; use augment()).",
           call. = FALSE)
    if (isTRUE(align)) {
      if (is.null(x$T_pre) || x$T_pre < 1L)
        stop("align = TRUE requires a fit with pre-treatment periods.",
             call. = FALSE)
      Y_synth <- Y_synth + .align_offset(x, Y_treat, Y_synth)
    }

    if (type == "gap")
      return(data.frame(time = times, gap = Y_treat - Y_synth,
                        stringsAsFactors = FALSE))

    # type == "trend"
    if (!is.numeric(show_donors) || length(show_donors) != 1L ||
        is.na(show_donors) || show_donors < 0)
      stop("`show_donors` must be a single number >= 0 (Inf shows all donors).",
           call. = FALSE)
    df <- data.frame(
      time   = c(times, times),
      value  = c(Y_treat, Y_synth),
      series = rep(c("Treated", "Synthetic Control"), each = length(times)),
      stringsAsFactors = FALSE
    )
    if (show_donors < 1) return(df)

    w    <- x$unit_weights
    Y_co <- donor_outcomes(x)
    if (is.null(w) || all(is.na(w)) || is.null(Y_co))
      stop("show_donors requires donor unit weights and outcomes ",
           "(available for sharp SCM/SDID/SI fits).", call. = FALSE)
    k   <- as.integer(min(length(w), show_donors))
    sel <- order(w, decreasing = TRUE)[seq_len(k)]
    donor_names <- (names(w) %||% sprintf("Donor %d", seq_along(w)))[sel]
    tname <- if (!is.null(names(x$Y_treat))) names(x$Y_treat)[1L] else "treated"
    df$unit <- rep(c(tname, NA_character_), each = length(times))
    df_don <- data.frame(
      time   = rep(times, times = k),
      value  = as.vector(Y_co[, sel, drop = FALSE]),
      series = "Donors",
      unit   = rep(donor_names, each = length(times)),
      stringsAsFactors = FALSE
    )
    out <- rbind(df, df_don)
    rownames(out) <- NULL
    return(out)
  }

  if (type == "weights") {
    w <- x$unit_weights
    if (is.null(w) || all(is.na(w)))
      stop("No unit weights available for this method ",
           "(GSC/MC/TASC use factor loadings).", call. = FALSE)
    if (!is.numeric(top_n) || length(top_n) != 1L || is.na(top_n) || top_n < 1)
      stop("`top_n` must be a single number >= 1 (Inf shows all donors).",
           call. = FALSE)

    df <- data.frame(
      unit   = names(w) %||% paste0("Unit_", seq_along(w)),
      weight = as.numeric(w),
      stringsAsFactors = FALSE
    )
    lam <- x$time_weights
    two_panel <- identical(x$method, "sdid") && !is.null(lam) &&
      !is.null(x$T_pre) && length(lam) == x$T_pre
    if (is.finite(top_n) && nrow(df) > top_n)
      df <- df[order(df$weight, decreasing = TRUE)[seq_len(top_n)], , drop = FALSE]
    if (!two_panel) {
      rownames(df) <- NULL
      return(df)
    }
    # SDID: the estimator is defined by unit weights (omega) AND time weights
    # (lambda); top_n subsets omega only, matching plot(type = "weights").
    df$panel <- "omega"
    t_lab <- as.character((x$times %||% seq_len(x$T_pre))[seq_len(x$T_pre)])
    df_t  <- data.frame(unit = t_lab, weight = as.numeric(lam), panel = "lambda",
                        stringsAsFactors = FALSE)
    out <- rbind(df, df_t)
    rownames(out) <- NULL
    return(out)
  }

  if (type == "pred_weights") {
    v <- x$v_weights
    if (is.null(v) || all(is.na(v)))
      stop("No predictor weights available. A V matrix is estimated only by ",
           "sharp SCM fits; staggered SCM and the other methods ",
           "(SDID/GSC/MC/TASC/SI) do not produce one.", call. = FALSE)
    if (!is.numeric(top_n) || length(top_n) != 1L || is.na(top_n) || top_n < 1)
      stop("`top_n` must be a single number >= 1 (Inf shows all predictors).",
           call. = FALSE)
    df <- data.frame(
      predictor = names(v) %||% paste0("V", seq_along(v)),
      weight    = as.numeric(v),
      stringsAsFactors = FALSE
    )
    if (is.finite(top_n) && nrow(df) > top_n)
      df <- df[order(df$weight, decreasing = TRUE)[seq_len(top_n)], , drop = FALSE]
    rownames(df) <- NULL
    return(df)
  }
}

#' @rdname plot_data
#' @export
plot_data.scm_placebo <- function(x, type = c("gaps", "ratios"),
                                  mspe_prune = Inf, ...) {
  type <- match.arg(type)
  if (!is.numeric(mspe_prune) || length(mspe_prune) != 1L || mspe_prune <= 0)
    stop("mspe_prune must be a single positive number (Inf = no pruning).",
         call. = FALSE)

  if (type == "gaps") {
    times <- x$times
    if (is.character(times) || is.factor(times))
      times <- as.numeric(as.character(times))
    keep <- is.finite(x$mspe_pre_placebo) &
      x$mspe_pre_placebo <= mspe_prune * x$mspe_pre_treated
    gaps  <- x$gaps[, keep, drop = FALSE]
    # sprintf keeps zero-length input zero-length (paste0 would collapse it)
    units <- colnames(gaps) %||% sprintf("Donor %d", which(keep))
    df_pl <- data.frame(
      time   = rep(times, times = ncol(gaps)),
      gap    = as.vector(gaps),
      unit   = rep(units, each = length(times)),
      series = "Placebo (donor pool)",
      stringsAsFactors = FALSE
    )
    df_tr <- data.frame(
      time   = times,
      gap    = x$treated_gap,
      unit   = NA_character_,
      series = "Treated",
      stringsAsFactors = FALSE
    )
    out <- rbind(df_tr, df_pl)
    rownames(out) <- NULL
    return(out)
  }

  # type == "ratios": every unit (including non-finite ratios excluded from the
  # plot by the mspe_threshold filter) is kept, so the frame is complete.
  r     <- x$mspe_ratios_all
  units <- names(r)
  if (is.null(units))
    units <- c("Treated", sprintf("Donor %d", seq_len(length(r) - 1L)))
  if (!nzchar(units[1L])) units[1L] <- "Treated"
  blank <- !nzchar(units)
  units[blank] <- sprintf("Donor %d", which(blank) - 1L)
  df <- data.frame(
    unit   = units,
    ratio  = as.numeric(r),
    series = c("Treated", rep("Placebo (donor pool)", length(r) - 1L)),
    stringsAsFactors = FALSE
  )
  rownames(df) <- NULL
  df
}