R/plot_mov.cor.R

In dendRoAnalyst: A Tool for Processing and Analyzing Dendrometer Data

#' @title Plot method for moving dendrometer-climate correlation
#'
#' @description
#' S3 plotting method for output of \code{mov.cor.dm()}.
#' Supports heatmaps, line plots, faceted line plots, peak-correlation summaries,
#' and comparison across multiple \code{mov.cor.dm()} objects.
#'
#' @param x Object returned by \code{mov.cor.dm()}.
#' @param y Unused.
#' @param sig.only Logical. If \code{TRUE}, only significant windows are shown in
#'   heatmaps and significance is highlighted in other plot types.
#' @param ci Numeric confidence level between 0 and 1. For non-bootstrap
#'   results, significance is based on \code{p_val < 1 - ci} or
#'   \code{p_adj < 1 - ci}. For bootstrap results, the stored logical
#'   \code{significance} column is used.
#' @param clim_vars Character vector of climate variables to plot, or
#'   \code{"all"} for all variables.
#' @param type Plot type. One of \code{"heatmap"}, \code{"line"},
#'   \code{"facet"}, \code{"peak"}, or \code{"compare"}.
#' @param x_axis X-axis style. One of \code{"time"} or \code{"doy"}.
#' @param use_adjusted Logical. For non-bootstrap objects, if \code{TRUE},
#'   significance uses \code{p_adj}; otherwise \code{p_val}.
#' @param show_na Logical. If \code{TRUE}, missing values are shown in heatmaps.
#' @param show_ci Logical. If \code{TRUE}, bootstrap confidence intervals are
#'   shown in line/facet/compare plots when available.
#' @param sig_mode One of \code{"outline"}, \code{"point"}, \code{"filter"},
#'   or \code{"none"} describing how significance is displayed in non-heatmap plots.
#' @param annotate_peak Logical. If \code{TRUE}, the strongest absolute
#'   correlation is marked for each climate variable in line/facet/compare plots.
#' @param show_window_label Logical. If \code{TRUE}, \code{type = "peak"} adds
#'   the climate transformation settings above each bar.
#' @param compare_with Optional list of additional \code{mov_cor_dm} objects for
#'   comparison. Used only when \code{type = "compare"}.
#' @param compare_labels Optional character vector of labels for the comparison
#'   objects. If \code{NULL}, defaults are generated automatically.
#' @param low_col Colour for negative correlations.
#' @param mid_col Colour for zero correlations.
#' @param high_col Colour for positive correlations.
#' @param na_col Fill colour for missing values in heatmaps.
#' @param line_size Numeric line width for line plots.
#' @param point_size Numeric point size for significance markers and peak markers.
#' @param alpha_sig Numeric transparency used for significance highlighting.
#' @param ... Unused.
#'
#' @return A \code{ggplot2} object, returned invisibly.
#'
#' @method plot mov_cor_dm
#' @importFrom dplyr bind_rows mutate filter arrange %>% left_join
#' @importFrom tibble as_tibble
#' @importFrom lubridate yday
#' @importFrom ggplot2 ggplot geom_line geom_point geom_raster geom_tile geom_col
#'   geom_ribbon geom_text aes theme_minimal labs facet_wrap theme element_text
#'   scale_fill_gradient2 scale_x_date theme_bw
#' @export
plot.mov_cor_dm <- function(x,
                            y = NULL,
                            sig.only = TRUE,
                            ci = 0.95,
                            clim_vars = "all",
                            type = c("heatmap", "line", "facet", "peak", "compare"),
                            x_axis = c("time", "doy"),
                            use_adjusted = TRUE,
                            show_na = TRUE,
                            show_ci = FALSE,
                            sig_mode = c("outline", "point", "filter", "none"),
                            annotate_peak = FALSE,
                            show_window_label = TRUE,
                            compare_with = NULL,
                            compare_labels = NULL,
                            low_col = "red",
                            mid_col = "white",
                            high_col = "blue",
                            na_col = "grey79",
                            line_size = 0.5,
                            point_size = 1.8,
                            alpha_sig = 0.9,
                            ...) {

  TIME <- corr <- para <- clim_var <- X <- source <- corr_lw_lim <- corr_up_lim <- NULL
  clim_fun <- lag_days <- accum_days <- window_label <- y_lab <- NULL

  if (!inherits(x, "mov_cor_dm")) {
    stop("The input object is not an output of 'mov.cor.dm()'.")
  }

  type <- match.arg(type)
  x_axis <- match.arg(x_axis)
  sig_mode <- match.arg(sig_mode)

  if (!is.numeric(ci) || length(ci) != 1 || is.na(ci) || ci <= 0 || ci >= 1) {
    stop("'ci' must be a single number between 0 and 1.")
  }

  .sig_note <- function(obj, ci, use_adjusted, sig.only, sig_mode) {
    if (isTRUE(obj$metadata$boot)) {
      base <- "Significance based on bootstrap CI excluding 0."
    } else {
      alpha_level <- 1 - ci
      if (isTRUE(use_adjusted)) {
        base <- paste0("Significance based on adjusted p-values < ", signif(alpha_level, 3), ".")
      } else {
        base <- paste0("Significance based on raw p-values < ", signif(alpha_level, 3), ".")
      }
    }

    mode_note <- if (type == "heatmap") {
      if (isTRUE(sig.only)) " Heatmap displays only significant windows." else " Heatmap displays all windows."
    } else {
      paste0(" sig_mode = '", sig_mode, "'.")
    }

    paste0(base, mode_note)
  }

  .extract_plot_data <- function(obj, keep_vars, label, ci, use_adjusted) {
    if (!inherits(obj, "mov_cor_dm")) {
      stop("All comparison objects must inherit from 'mov_cor_dm'.")
    }
    if (!is.list(obj$results)) {
      stop("Comparison object does not contain valid 'results'.")
    }

    varnames <- names(obj$results)
    miss <- setdiff(keep_vars, varnames)
    if (length(miss) > 0) {
      stop(sprintf(
        "The following climate variables are missing from comparison object '%s': %s",
        label, paste(miss, collapse = ", ")
      ))
    }

    dat <- dplyr::bind_rows(obj$results[keep_vars])

    if (!"clim_var" %in% names(dat)) {
      dat$clim_var <- rep(keep_vars, lengths(obj$results[keep_vars]))
    }

    dat$TIME <- as.Date(dat$TIME)
    dat$para <- toupper(dat$clim_var)
    dat$source <- label

    if (isTRUE(obj$metadata$boot)) {
      dat$sig_plot <- dat$significance
    } else {
      alpha_level <- 1 - ci
      if (use_adjusted && "p_adj" %in% names(dat)) {
        dat$sig_plot <- !is.na(dat$p_adj) & dat$p_adj < alpha_level
      } else {
        dat$sig_plot <- !is.na(dat$p_val) & dat$p_val < alpha_level
      }
    }

    dat
  }

  .peak_df <- function(din, by_source = FALSE) {
    if (by_source) {
      split_list <- split(din, interaction(din$clim_var, din$source, drop = TRUE))
    } else {
      split_list <- split(din, din$clim_var)
    }

    pk <- lapply(split_list, function(tab) {
      tab <- tab[!is.na(tab$corr), , drop = FALSE]
      if (nrow(tab) == 0) return(NULL)
      tab[which.max(abs(tab$corr)), , drop = FALSE]
    })
    dplyr::bind_rows(pk)
  }

  varnames <- names(x$results)
  if (length(varnames) == 0) {
    stop("No climate-variable results found in object.")
  }

  if (length(clim_vars) == 1 && identical(clim_vars, "all")) {
    keep_vars <- varnames
  } else {
    if (!is.character(clim_vars)) {
      stop("'clim_vars' must be 'all' or a character vector of climate variable names.")
    }
    miss <- setdiff(clim_vars, varnames)
    if (length(miss) > 0) {
      stop(sprintf(
        "The following climate variables are not available in the object: %s",
        paste(miss, collapse = ", ")
      ))
    }
    keep_vars <- clim_vars
  }

  dat <- .extract_plot_data(x, keep_vars, label = "x1", ci = ci, use_adjusted = use_adjusted)

  if (type == "compare") {
    objs <- c(list(x), compare_with)
    if (length(objs) < 2) {
      stop("For type = 'compare', provide at least one object in 'compare_with'.")
    }
    if (is.null(compare_labels)) {
      compare_labels <- paste0("obj", seq_along(objs))
    }
    if (length(compare_labels) != length(objs)) {
      stop("'compare_labels' must have the same length as c(list(x), compare_with).")
    }

    dat <- dplyr::bind_rows(
      lapply(seq_along(objs), function(i) {
        .extract_plot_data(objs[[i]], keep_vars, label = compare_labels[i], ci = ci, use_adjusted = use_adjusted)
      })
    )
  }

  if (nrow(dat) == 0) {
    stop("No data available for plotting.")
  }

  if (x_axis == "doy") {
    dat$X <- lubridate::yday(dat$TIME)
    x_lab <- "Day of year"
  } else {
    dat$X <- dat$TIME
    x_lab <- "Time"
  }

  sig_caption <- .sig_note(x, ci = ci, use_adjusted = use_adjusted, sig.only = sig.only, sig_mode = sig_mode)

  if (type != "heatmap" && sig_mode == "filter") {
    dat <- dat[dat$sig_plot %in% TRUE, , drop = FALSE]
    if (nrow(dat) == 0) {
      stop("No significant windows available after filtering.")
    }
  }

  if (type == "heatmap") {
    dat_hm <- dat

    if (isTRUE(sig.only)) {
      dat_hm <- dat_hm[dat_hm$sig_plot %in% TRUE, , drop = FALSE]
      if (nrow(dat_hm) == 0) {
        stop("No significant windows available for the heatmap.")
      }
    }

    p <- ggplot2::ggplot(dat_hm, ggplot2::aes(x = X, y = para, fill = corr)) +
      ggplot2::geom_tile() +
      ggplot2::scale_fill_gradient2(
        low = low_col,
        mid = mid_col,
        high = high_col,
        midpoint = 0,
        limits = c(-1, 1),
        na.value = if (show_na) na_col else "transparent"
      ) +
      ggplot2::labs(
        title = "Running correlation with climate",
        x = x_lab,
        y = "Climate variable",
        fill = "r",
        caption = sig_caption
      ) +
      ggplot2::theme_bw() +
      ggplot2::theme(
        axis.text.x = ggplot2::element_text(angle = 45, hjust = 1),
        axis.title = ggplot2::element_text(size = 14),
        axis.text = ggplot2::element_text(size = 11),
        plot.caption = ggplot2::element_text(hjust = 0, size = 9)
      )

    if (x_axis == "time") {
      p <- p + ggplot2::scale_x_date(date_labels = "%b-%y", date_breaks = "1 month")
    }

    print(p)
    return(invisible(p))
  }

  if (type == "line") {
    if (length(unique(dat$para)) > 1) {
      warning("Multiple climate variables selected; lines are overlaid. Use type = 'facet' for separate panels.")
    }

    p <- ggplot2::ggplot(dat, ggplot2::aes(x = X, y = corr, color = para)) +
      ggplot2::geom_line(linewidth = line_size) +
      ggplot2::theme_bw() +
      ggplot2::labs(
        title = "Running correlation with climate",
        x = x_lab,
        y = "Correlation",
        color = "Climate variable",
        caption = sig_caption
      ) +
      ggplot2::theme(
        axis.text.x = ggplot2::element_text(angle = 45, hjust = 1),
        axis.title = ggplot2::element_text(size = 14),
        axis.text = ggplot2::element_text(size = 11),
        plot.caption = ggplot2::element_text(hjust = 0, size = 9)
      )

    if (show_ci && all(c("corr_lw_lim", "corr_up_lim") %in% names(dat))) {
      p <- p +
        ggplot2::geom_ribbon(
          ggplot2::aes(ymin = corr_lw_lim, ymax = corr_up_lim, fill = para),
          alpha = 0.15,
          color = NA,
          inherit.aes = TRUE
        )
    }

    if (sig_mode == "point" && sig.only) {
      dat2 <- dat[dat$sig_plot %in% TRUE, , drop = FALSE]
      if (nrow(dat2) > 0) {
        p <- p + ggplot2::geom_point(
          data = dat2,
          ggplot2::aes(x = X, y = corr, color = para),
          size = point_size,
          alpha = alpha_sig
        )
      }
    }

    if (annotate_peak) {
      pk <- .peak_df(dat, by_source = FALSE)
      if (nrow(pk) > 0) {
        p <- p + ggplot2::geom_point(
          data = pk,
          ggplot2::aes(x = X, y = corr, color = para),
          size = point_size + 1,
          shape = 21,
          fill = "yellow",
          stroke = 1.1
        )
      }
    }

    print(p)
    return(invisible(p))
  }

  if (type == "facet") {
    p <- ggplot2::ggplot(dat, ggplot2::aes(x = X, y = corr)) +
      ggplot2::geom_line(linewidth = line_size, color = "black") +
      ggplot2::facet_wrap(~ para, scales = "free_y", ncol = 1) +
      ggplot2::theme_bw() +
      ggplot2::labs(
        title = "Running correlation with climate",
        x = x_lab,
        y = "Correlation",
        caption = sig_caption
      ) +
      ggplot2::theme(
        axis.text.x = ggplot2::element_text(angle = 45, hjust = 1),
        axis.title = ggplot2::element_text(size = 14),
        axis.text = ggplot2::element_text(size = 11),
        plot.caption = ggplot2::element_text(hjust = 0, size = 9)
      )

    if (show_ci && all(c("corr_lw_lim", "corr_up_lim") %in% names(dat))) {
      p <- p +
        ggplot2::geom_ribbon(
          ggplot2::aes(ymin = corr_lw_lim, ymax = corr_up_lim),
          alpha = 0.15,
          fill = "grey70",
          color = NA
        )
    }

    if (sig_mode %in% c("point", "outline") && sig.only) {
      dat2 <- dat[dat$sig_plot %in% TRUE, , drop = FALSE]
      if (nrow(dat2) > 0) {
        p <- p + ggplot2::geom_point(
          data = dat2,
          ggplot2::aes(x = X, y = corr),
          size = point_size,
          color = "red",
          alpha = alpha_sig
        )
      }
    }

    if (annotate_peak) {
      pk <- .peak_df(dat, by_source = FALSE)
      if (nrow(pk) > 0) {
        p <- p + ggplot2::geom_point(
          data = pk,
          ggplot2::aes(x = X, y = corr),
          size = point_size + 1,
          shape = 21,
          fill = "yellow",
          color = "black",
          stroke = 1.1
        )
      }
    }

    print(p)
    return(invisible(p))
  }

  if (type == "compare") {
    p <- ggplot2::ggplot(dat, ggplot2::aes(x = X, y = corr, color = source)) +
      ggplot2::geom_line(linewidth = line_size) +
      ggplot2::facet_wrap(~ para, scales = "free_y", ncol = 1) +
      ggplot2::theme_bw() +
      ggplot2::labs(
        title = "Comparison of running correlations",
        x = x_lab,
        y = "Correlation",
        color = "Object",
        caption = sig_caption
      ) +
      ggplot2::theme(
        axis.text.x = ggplot2::element_text(angle = 45, hjust = 1),
        axis.title = ggplot2::element_text(size = 14),
        axis.text = ggplot2::element_text(size = 11),
        plot.caption = ggplot2::element_text(hjust = 0, size = 9)
      )

    if (show_ci && all(c("corr_lw_lim", "corr_up_lim") %in% names(dat))) {
      p <- p +
        ggplot2::geom_ribbon(
          ggplot2::aes(ymin = corr_lw_lim, ymax = corr_up_lim, fill = source),
          alpha = 0.12,
          color = NA
        )
    }

    if (sig_mode == "point" && sig.only) {
      dat2 <- dat[dat$sig_plot %in% TRUE, , drop = FALSE]
      if (nrow(dat2) > 0) {
        p <- p + ggplot2::geom_point(
          data = dat2,
          ggplot2::aes(x = X, y = corr, color = source),
          size = point_size,
          alpha = alpha_sig
        )
      }
    }

    if (annotate_peak) {
      pk <- .peak_df(dat, by_source = TRUE)
      if (nrow(pk) > 0) {
        p <- p + ggplot2::geom_point(
          data = pk,
          ggplot2::aes(x = X, y = corr, color = source),
          size = point_size + 1,
          shape = 21,
          fill = "yellow",
          stroke = 1.1
        )
      }
    }

    print(p)
    return(invisible(p))
  }

  if (type == "peak") {
    peak_df <- .peak_df(dat, by_source = FALSE)

    if (nrow(peak_df) == 0) {
      stop("No non-missing correlation values available.")
    }

    if (!is.null(x$metadata$climate_settings) && nrow(x$metadata$climate_settings) > 0) {
      peak_df <- dplyr::left_join(peak_df, x$metadata$climate_settings, by = "clim_var")
      peak_df$window_label <- paste0(
        peak_df$clim_fun,
        ", lag=", peak_df$lag_days,
        ", win=", peak_df$accum_days
      )
    } else {
      peak_df$window_label <- NA_character_
    }

    p <- ggplot2::ggplot(peak_df, ggplot2::aes(x = para, y = corr, fill = corr)) +
      ggplot2::geom_col() +
      ggplot2::scale_fill_gradient2(
        low = low_col,
        mid = mid_col,
        high = high_col,
        midpoint = 0,
        limits = c(-1, 1),
        na.value = na_col
      ) +
      ggplot2::labs(
        title = "Peak absolute running correlation",
        x = "Climate variable",
        y = "Peak correlation",
        fill = "r",
        caption = sig_caption
      ) +
      ggplot2::theme_bw() +
      ggplot2::theme(
        axis.title = ggplot2::element_text(size = 14),
        axis.text = ggplot2::element_text(size = 11),
        plot.caption = ggplot2::element_text(hjust = 0, size = 9)
      )

    if (isTRUE(show_window_label) && any(!is.na(peak_df$window_label))) {
      y_off <- 0.04 * max(abs(peak_df$corr), na.rm = TRUE)
      peak_df$y_lab <- ifelse(peak_df$corr >= 0, peak_df$corr + y_off, peak_df$corr - y_off)

      p <- p + ggplot2::geom_text(
        data = peak_df,
        ggplot2::aes(x = para, y = y_lab, label = window_label),
        inherit.aes = FALSE,
        size = 3
      )
    }

    print(p)
    return(invisible(p))
  }
}


#' @title Backward-compatible wrapper for plotting moving correlation
#'
#' @description
#' Wrapper around \code{plot()} for objects returned by \code{mov.cor.dm()}.
#'
#' @param mov.cor.output Output of \code{mov.cor.dm()}.
#' @param ... Passed to \code{plot.mov_cor_dm()}.
#'
#' @return A \code{ggplot2} object.
#' @export
plot_mov.cor <- function(mov.cor.output, ...) {
  plot(mov.cor.output, ...)
}


#' @method plot mov.cor
#' @export
plot.mov.cor <- function(x, y = NULL, ...) {
  if (!inherits(x, "mov_cor_dm")) {
    x <- list(
      results = x,
      metadata = list(
        boot = FALSE,
        clim_names = names(x),
        tree_name = NA_character_,
        win_size = NA_integer_,
        cor_method = NA_character_,
        dm_stat = NA_character_,
        climate_settings = NULL
      )
    )
    class(x) <- c("mov_cor_dm", "mov.cor")
  }
  plot.mov_cor_dm(x, ...)
}


#' @method plot mov.cor.boot
#' @export
plot.mov.cor.boot <- function(x, y = NULL, ...) {
  if (!inherits(x, "mov_cor_dm")) {
    x <- list(
      results = x,
      metadata = list(
        boot = TRUE,
        clim_names = names(x),
        tree_name = NA_character_,
        win_size = NA_integer_,
        cor_method = NA_character_,
        dm_stat = NA_character_,
        climate_settings = NULL
      )
    )
    class(x) <- c("mov_cor_dm_boot", "mov_cor_dm", "mov.cor.boot")
  }
  plot.mov_cor_dm(x, ...)
}


#' @title Plot method for summaries of moving dendrometer-climate correlation
#'
#' @description
#' S3 plotting method for objects returned by \code{summary.mov_cor_dm()}.
#'
#' @param x Object of class \code{"summary_mov_cor_dm"}.
#' @param y Unused.
#' @param type Plot type. One of \code{"peak_corr"}, \code{"prop_significant"},
#'   or \code{"peak_time"}.
#' @param x_axis X-axis style for \code{type = "peak_time"}. One of
#'   \code{"time"} or \code{"doy"}.
#' @param low_col Colour for negative values.
#' @param mid_col Colour for zero.
#' @param high_col Colour for positive values.
#' @param show_window_label Logical. If \code{TRUE}, climate settings are shown
#'   as labels in \code{type = "peak_corr"}.
#' @param ... Unused.
#'
#' @return A \code{ggplot2} object, returned invisibly.
#'
#' @method plot summary_mov_cor_dm
#' @importFrom ggplot2 ggplot aes geom_col geom_point geom_text scale_fill_gradient2
#' @export
plot.summary_mov_cor_dm <- function(x,
                                    y = NULL,
                                    type = c("peak_corr", "prop_significant", "peak_time"),
                                    x_axis = c("time", "doy"),
                                    low_col = "red",
                                    mid_col = "white",
                                    high_col = "blue",
                                    show_window_label = TRUE,
                                    ...) {
  clim_var <- peak_corr <- y_lab <- window_label <- prop_significant <- X <- NULL

  if (!inherits(x, "summary_mov_cor_dm")) {
    stop("Input must be an object of class 'summary_mov_cor_dm'.")
  }

  type <- match.arg(type)
  x_axis <- match.arg(x_axis)

  tab <- tibble::as_tibble(x$table)

  if (nrow(tab) == 0) {
    stop("No summary data available for plotting.")
  }

  if (type == "peak_corr") {
    p <- ggplot2::ggplot(tab, ggplot2::aes(x = clim_var, y = peak_corr, fill = peak_corr)) +
      ggplot2::geom_col() +
      ggplot2::scale_fill_gradient2(
        low = low_col,
        mid = mid_col,
        high = high_col,
        midpoint = 0,
        limits = c(-1, 1)
      ) +
      ggplot2::labs(
        x = "Climate variable",
        y = "Peak correlation",
        fill = "r",
        title = "Peak running correlation by climate variable"
      ) +
      ggplot2::theme_bw() +
      ggplot2::theme(
        axis.title = ggplot2::element_text(size = 14),
        axis.text = ggplot2::element_text(size = 11)
      )

    if (isTRUE(show_window_label) &&
        all(c("clim_fun", "lag_days", "accum_days") %in% names(tab))) {
      tab$window_label <- paste0(tab$clim_fun, ", lag=", tab$lag_days, ", win=", tab$accum_days)
      y_off <- 0.04 * max(abs(tab$peak_corr), na.rm = TRUE)
      tab$y_lab <- ifelse(tab$peak_corr >= 0, tab$peak_corr + y_off, tab$peak_corr - y_off)

      p <- p + ggplot2::geom_text(
        data = tab,
        ggplot2::aes(x = clim_var, y = y_lab, label = window_label),
        inherit.aes = FALSE,
        size = 3
      )
    }

    print(p)
    return(invisible(p))
  }

  if (type == "prop_significant") {
    p <- ggplot2::ggplot(tab, ggplot2::aes(x = clim_var, y = prop_significant)) +
      ggplot2::geom_col(fill = "steelblue") +
      ggplot2::labs(
        x = "Climate variable",
        y = "Proportion significant",
        title = "Proportion of significant windows"
      ) +
      ggplot2::theme_bw() +
      ggplot2::theme(
        axis.title = ggplot2::element_text(size = 14),
        axis.text = ggplot2::element_text(size = 11)
      )

    print(p)
    return(invisible(p))
  }

  if (type == "peak_time") {
    tt <- as.Date(tab$peak_time)
    if (x_axis == "doy") {
      tab$X <- lubridate::yday(tt)
      x_lab <- "Peak time (day of year)"
    } else {
      tab$X <- tt
      x_lab <- "Peak time"
    }

    p <- ggplot2::ggplot(tab, ggplot2::aes(x = X, y = clim_var)) +
      ggplot2::geom_point(size = 3, color = "black") +
      ggplot2::labs(
        x = x_lab,
        y = "Climate variable",
        title = "Timing of strongest running correlation"
      ) +
      ggplot2::theme_bw() +
      ggplot2::theme(
        axis.title = ggplot2::element_text(size = 14),
        axis.text = ggplot2::element_text(size = 11)
      )

    print(p)
    return(invisible(p))
  }
}