# helpers ----------------------------------------------------------------------

#' @keywords internal
dmwc_parse_time <- function(tt) {
  if (inherits(tt, "POSIXct") || inherits(tt, "POSIXt")) return(as.POSIXct(tt))
  if (inherits(tt, "Date")) return(as.POSIXct(tt))

  out <- suppressWarnings(lubridate::ymd_hms(tt, quiet = TRUE))
  if (all(is.na(out))) {
    out <- suppressWarnings(
      lubridate::parse_date_time(
        tt,
        orders = c(
          "ymd HMS", "ymd HM", "ymd",
          "dmy HMS", "dmy HM", "dmy",
          "mdy HMS", "mdy HM", "mdy"
        ),
        quiet = TRUE
      )
    )
  }
  as.POSIXct(out)
}

#' @keywords internal
dmwc_infer_dt <- function(time_vec) {
  dsec <- as.numeric(diff(time_vec), units = "secs")
  dsec <- dsec[is.finite(dsec) & dsec > 0]

  if (length(dsec) == 0) {
    stop("Could not infer time step from TIME column.")
  }

  med <- stats::median(dsec, na.rm = TRUE)

  if (med >= 86400 && abs(med / 86400 - round(med / 86400)) < 1e-6) {
    return(list(dt = med / 86400, unit = "days", sec = med, hours = med / 3600))
  }
  if (med >= 3600 && abs(med / 3600 - round(med / 3600)) < 1e-6) {
    return(list(dt = med / 3600, unit = "hours", sec = med, hours = med / 3600))
  }
  if (med >= 60 && abs(med / 60 - round(med / 60)) < 1e-6) {
    return(list(dt = med / 60, unit = "mins", sec = med, hours = med / 3600))
  }

  list(dt = med, unit = "secs", sec = med, hours = med / 3600)
}

#' @keywords internal
dmwc_period_to_hours <- function(period, time_unit) {
  if (time_unit == "days") return(period * 24)
  if (time_unit == "hours") return(period)
  if (time_unit == "mins") return(period / 60)
  if (time_unit == "secs") return(period / 3600)
  period
}

#' @keywords internal
dmwc_fun_match <- function(fun_name) {
  fun_name <- match.arg(fun_name, c("mean", "sum", "max", "min"))
  switch(
    fun_name,
    mean = function(z) if (all(is.na(z))) NA_real_ else mean(z, na.rm = TRUE),
    sum  = function(z) if (all(is.na(z))) NA_real_ else sum(z, na.rm = TRUE),
    max  = function(z) if (all(is.na(z))) NA_real_ else max(z, na.rm = TRUE),
    min  = function(z) if (all(is.na(z))) NA_real_ else min(z, na.rm = TRUE)
  )
}

#' @keywords internal
dmwc_resample_regular <- function(dat, target_sec, fun_map) {
  tz_use <- attr(dat$TIME, "tzone")
  if (is.null(tz_use) || length(tz_use) == 0) tz_use <- "UTC"

  bucket_num <- floor(as.numeric(dat$TIME) / target_sec) * target_sec
  dat$TIME_BIN <- as.POSIXct(bucket_num, origin = "1970-01-01", tz = tz_use)

  val_cols <- setdiff(names(dat), c("TIME", "TIME_BIN"))

  agg_list <- lapply(val_cols, function(cc) {
    ff <- fun_map[[cc]]
    stats::aggregate(dat[[cc]], by = list(TIME = dat$TIME_BIN), FUN = ff)
  })

  out <- agg_list[[1]]
  names(out)[2] <- val_cols[1]

  if (length(agg_list) > 1) {
    for (i in 2:length(agg_list)) {
      tmp <- agg_list[[i]]
      names(tmp)[2] <- val_cols[i]
      out <- merge(out, tmp, by = "TIME", all = TRUE)
    }
  }

  full_time <- seq(min(out$TIME, na.rm = TRUE), max(out$TIME, na.rm = TRUE), by = target_sec)
  out <- merge(
    data.frame(TIME = as.POSIXct(full_time, origin = "1970-01-01", tz = tz_use)),
    out,
    by = "TIME",
    all.x = TRUE
  )
  out <- out[order(out$TIME), , drop = FALSE]
  tibble::as_tibble(out)
}

#' @keywords internal
dmwc_fill_na <- function(z) {
  z2 <- zoo::na.approx(z, na.rm = FALSE)
  z2 <- zoo::na.locf(z2, na.rm = FALSE)
  z2 <- zoo::na.locf(z2, fromLast = TRUE, na.rm = FALSE)
  z2
}

#' @keywords internal
dmwc_select_cols <- function(dat, selector, exclude = character(0)) {
  cols <- setdiff(names(dat), exclude)
  cols <- cols[vapply(dat[cols], is.numeric, logical(1))]

  if (length(cols) == 0) stop("No numeric series available.")

  if (is.character(selector) && length(selector) == 1 && tolower(selector) == "all") {
    return(cols)
  }

  if (is.numeric(selector)) {
    idx <- as.integer(selector)
    if (any(is.na(idx)) || any(idx < 1) || any(idx > length(cols))) {
      stop("Numeric selector is out of range.")
    }
    return(cols[idx])
  }

  if (is.character(selector)) {
    miss <- setdiff(selector, cols)
    if (length(miss) > 0) {
      warning("These requested columns were not found and were ignored: ",
              paste(miss, collapse = ", "))
    }
    keep <- intersect(selector, cols)
    if (length(keep) == 0) stop("None of the requested columns were found.")
    return(keep)
  }

  stop("Selector must be 'all', numeric indices, or character names.")
}


# main function ----------------------------------------------------------------

#' Wavelet coherence between dendrometer and climate series
#'
#' @description
#' Computes cross-wavelet power and wavelet coherence between dendrometer series
#' and climate variables using \code{WaveletComp::analyze.coherency()}.
#'
#' The function:
#' \itemize{
#'   \item crops dendrometer and climate inputs to their common time window,
#'   \item detects their temporal resolutions,
#'   \item resamples both to the coarser detected resolution,
#'   \item optionally interpolates missing values,
#'   \item optionally uses first differences of the dendrometer series,
#'   \item computes pairwise coherence for all selected tree/climate combinations.
#' }
#'
#' @param dm Dendrometer input. Either a data frame with time in the first column
#'   or an object of class \code{"dm_detrended"}.
#' @param clim Climate data frame. The first column must be time; remaining
#'   selected columns must be numeric climate variables.
#' @param TreeNum Dendrometer series selector: \code{"all"}, numeric indices,
#'   or character names.
#' @param clim_vars Climate variable selector: \code{"all"}, numeric indices,
#'   or character names.
#' @param source One of \code{"auto"}, \code{"raw"}, \code{"detrended"},
#'   or \code{"first_diff"}.
#' @param detrended_col For \code{dm_detrended} input, which table to use.
#' @param dm_fun Aggregation function for dendrometer resampling: one of
#'   \code{"mean"}, \code{"max"}, \code{"min"}, \code{"sum"}.
#' @param clim_funs Either a single aggregation function applied to all climate
#'   variables, or a named character vector such as
#'   \code{c(temp = "mean", VPD = "mean", prec = "sum")}.
#' @param na_action One of \code{"interpolate"} or \code{"fail"}.
#' @param loess_span Detrending span passed to \code{WaveletComp::analyze.coherency()}.
#'   Use \code{0} to disable internal loess detrending.
#' @param dj Frequency resolution passed to \code{WaveletComp::analyze.coherency()}.
#' @param lowerPeriod Optional lower period in native analysis units.
#' @param upperPeriod Optional upper period in native analysis units.
#' @param window.type.t Time-direction smoothing window type.
#' @param window.type.s Scale-direction smoothing window type.
#' @param window.size.t Time-direction smoothing window size.
#' @param window.size.s Scale-direction smoothing window size.
#' @param make.pval Logical. If \code{TRUE}, compute p-values.
#' @param make_pval Optional alias for \code{make.pval}.
#' @param method Surrogate generation method for p-values.
#' @param n.sim Number of simulations.
#' @param n_sim Optional alias for \code{n.sim}.
#' @param verbose Logical.
#'
#' @return
#' An object of class \code{"dm_wavelet_coherence"}.
#'
#' @details
#' \code{WaveletComp::analyze.coherency()} expects two aligned series in one data
#' frame plus a \code{dt} value in the analysis time units, and returns
#' cross-wavelet power, coherence, phase angles, p-values, Fourier periods, and
#' cone-of-influence fields.
#'
#' @examples
#' \donttest{
#' wc <- dm_wavelet_coherence(
#'   dm = gf_nepa17,
#'   clim = ktm_clim_hourly,
#'   TreeNum = 1,
#'   clim_vars = c("temp", "VPD"),
#'   source = "raw",
#'   dm_fun = "mean",
#'   clim_funs = c(temp = "mean", VPD = "mean"),
#'   loess_span = 0,
#'   make.pval = TRUE,
#'   n.sim = 10,
#'   verbose = FALSE
#' )
#'
#' plot(wc)
#' plot(wc, type = "cross_power")
#' plot(wc, type = "average_coherence")
#' plot(wc, type = "phase")
#' plot(wc, type = "series")
#' }
#'
#' @importFrom dplyr bind_rows arrange %>%
#' @importFrom tibble tibble as_tibble
#' @importFrom lubridate ymd_hms parse_date_time
#' @importFrom zoo na.approx na.locf
#' @importFrom stats setNames
#' @export
dm_wavelet_coherence <- function(dm,
                                 clim,
                                 TreeNum = "all",
                                 clim_vars = "all",
                                 source = c("auto", "raw", "detrended", "first_diff"),
                                 detrended_col = c("detrended_data"),
                                 dm_fun = c("mean", "max", "min", "sum"),
                                 clim_funs = "mean",
                                 na_action = c("interpolate", "fail"),
                                 loess_span = 0,
                                 dj = 1 / 20,
                                 lowerPeriod = NULL,
                                 upperPeriod = NULL,
                                 window.type.t = 1,
                                 window.type.s = 1,
                                 window.size.t = 5,
                                 window.size.s = 1 / 4,
                                 make.pval = TRUE,
                                 make_pval = NULL,
                                 method = "white.noise",
                                 n.sim = 10,
                                 n_sim = NULL,
                                 verbose = TRUE) {

  TIME <- NULL

  if (!requireNamespace("WaveletComp", quietly = TRUE)) {
    stop("Package 'WaveletComp' is required for dm_wavelet_coherence().")
  }

  source <- match.arg(source)
  detrended_col <- match.arg(detrended_col)
  dm_fun <- match.arg(dm_fun)
  na_action <- match.arg(na_action)

  # support snake_case aliases too
  if (!is.null(make_pval)) {
    make.pval <- make_pval
  }
  if (!is.null(n_sim)) {
    n.sim <- n_sim
  }

  cl <- match.call()

  # dendrometer input
  if (inherits(dm, "dm_detrended")) {
    dm_type <- "dm_detrended"
    if (source == "raw") {
      stop("source = 'raw' is not valid for a dm_detrended object.")
    }
    dm0 <- tibble::as_tibble(dm[[detrended_col]])
    dm_meta <- intersect(
      c("TIME", "doy", "season_label", "season_start", "season_end", "season_day"),
      names(dm0)
    )
    dm_cols <- dmwc_select_cols(dm0, TreeNum, exclude = dm_meta)
    dm_dat <- dm0[, c("TIME", dm_cols), drop = FALSE]
    source_used <- if (source == "auto") "detrended" else source
  } else if (is.data.frame(dm)) {
    dm_type <- "raw"
    if (source == "detrended") {
      stop("source = 'detrended' requires a dm_detrended object.")
    }
    dm0 <- tibble::as_tibble(dm)
    names(dm0)[1] <- "TIME"
    dm_cols <- dmwc_select_cols(dm0, TreeNum, exclude = "TIME")
    dm_dat <- dm0[, c("TIME", dm_cols), drop = FALSE]
    source_used <- if (source == "auto") "raw" else source
  } else {
    stop("'dm' must be a data.frame or an object of class 'dm_detrended'.")
  }

  # climate input
  if (!is.data.frame(clim)) {
    stop("'clim' must be a data.frame.")
  }
  clim0 <- tibble::as_tibble(clim)
  names(clim0)[1] <- "TIME"
  clim_cols <- dmwc_select_cols(clim0, clim_vars, exclude = "TIME")
  clim_dat <- clim0[, c("TIME", clim_cols), drop = FALSE]

  # parse times
  dm_dat$TIME <- dmwc_parse_time(dm_dat$TIME)
  clim_dat$TIME <- dmwc_parse_time(clim_dat$TIME)

  if (any(is.na(dm_dat$TIME))) {
    stop("Some dendrometer timestamps could not be parsed.")
  }
  if (any(is.na(clim_dat$TIME))) {
    stop("Some climate timestamps could not be parsed.")
  }

  dm_dat <- dm_dat[order(dm_dat$TIME), , drop = FALSE]
  clim_dat <- clim_dat[order(clim_dat$TIME), , drop = FALSE]

  # crop to common overlap
  t_start <- max(min(dm_dat$TIME, na.rm = TRUE), min(clim_dat$TIME, na.rm = TRUE))
  t_end   <- min(max(dm_dat$TIME, na.rm = TRUE), max(clim_dat$TIME, na.rm = TRUE))

  if (!is.finite(as.numeric(t_start)) || !is.finite(as.numeric(t_end)) || t_start >= t_end) {
    stop("No overlapping time window between dendrometer and climate data.")
  }

  dm_dat <- dm_dat[dm_dat$TIME >= t_start & dm_dat$TIME <= t_end, , drop = FALSE]
  clim_dat <- clim_dat[clim_dat$TIME >= t_start & clim_dat$TIME <= t_end, , drop = FALSE]

  if (nrow(dm_dat) < 4 || nrow(clim_dat) < 4) {
    stop("Too few overlapping observations after cropping.")
  }

  # detect resolution and choose coarser
  dm_dt <- dmwc_infer_dt(dm_dat$TIME)
  clim_dt <- dmwc_infer_dt(clim_dat$TIME)

  target_sec <- max(dm_dt$sec, clim_dt$sec)
  target_info <- dmwc_infer_dt(
    seq(as.POSIXct("2000-01-01", tz = "UTC"), by = target_sec, length.out = 5)
  )

  # function map
  dm_fun_map <- setNames(
    replicate(length(dm_cols), dmwc_fun_match(dm_fun), simplify = FALSE),
    dm_cols
  )

  if (length(clim_funs) == 1) {
    clim_fun_map <- setNames(
      replicate(length(clim_cols), dmwc_fun_match(clim_funs), simplify = FALSE),
      clim_cols
    )
  } else {
    if (is.null(names(clim_funs))) {
      stop("'clim_funs' must be length 1 or a named vector/list matching climate variables.")
    }
    missing_fun <- setdiff(clim_cols, names(clim_funs))
    if (length(missing_fun) > 0) {
      stop("Missing aggregation functions for climate variables: ",
           paste(missing_fun, collapse = ", "))
    }
    clim_fun_map <- lapply(clim_cols, function(cc) dmwc_fun_match(clim_funs[[cc]]))
    names(clim_fun_map) <- clim_cols
  }

  # resample both to coarser resolution
  dm_res <- dmwc_resample_regular(dm_dat, target_sec = target_sec, fun_map = dm_fun_map)
  clim_res <- dmwc_resample_regular(clim_dat, target_sec = target_sec, fun_map = clim_fun_map)

  # align exactly on same time grid
  common_time <- intersect(dm_res$TIME, clim_res$TIME)
  if (length(common_time) < 8) {
    stop("Too few aligned observations after resampling.")
  }

  dm_res <- dm_res[dm_res$TIME %in% common_time, , drop = FALSE]
  clim_res <- clim_res[clim_res$TIME %in% common_time, , drop = FALSE]

  dm_res <- dm_res[order(dm_res$TIME), , drop = FALSE]
  clim_res <- clim_res[order(clim_res$TIME), , drop = FALSE]

  # NA handling
  for (cc in setdiff(names(dm_res), "TIME")) {
    if (na_action == "interpolate") {
      dm_res[[cc]] <- dmwc_fill_na(dm_res[[cc]])
    } else if (any(is.na(dm_res[[cc]]))) {
      stop("Missing values found in dendrometer series after resampling: ", cc)
    }
  }

  for (cc in setdiff(names(clim_res), "TIME")) {
    if (na_action == "interpolate") {
      clim_res[[cc]] <- dmwc_fill_na(clim_res[[cc]])
    } else if (any(is.na(clim_res[[cc]]))) {
      stop("Missing values found in climate series after resampling: ", cc)
    }
  }

  # first differences on dendrometer if requested
  if (source_used == "first_diff") {
    for (cc in dm_cols) {
      dm_res[[cc]] <- c(NA_real_, diff(dm_res[[cc]]))
    }
    dm_res <- dm_res[-1, , drop = FALSE]
    clim_res <- clim_res[-1, , drop = FALSE]
  }

  # final NA check
  for (cc in dm_cols) {
    if (any(is.na(dm_res[[cc]]))) {
      stop("Dendrometer series still contains missing values: ", cc)
    }
  }
  for (cc in clim_cols) {
    if (any(is.na(clim_res[[cc]]))) {
      stop("Climate series still contains missing values: ", cc)
    }
  }

  # pairwise coherence
  pair_results <- list()
  aligned_list <- list()
  summary_rows <- list()
  idx <- 1L

  for (dm_col in dm_cols) {
    for (clim_col in clim_cols) {
      pair_name <- paste(dm_col, clim_col, sep = "__")

      pair_df <- data.frame(
        date = dm_res$TIME,
        dm = as.numeric(dm_res[[dm_col]]),
        clim = as.numeric(clim_res[[clim_col]])
      )

      coh_args <- list(
        my.data = pair_df,
        my.pair = c("dm", "clim"),
        loess.span = loess_span,
        dt = target_info$dt,
        dj = dj,
        window.type.t = window.type.t,
        window.type.s = window.type.s,
        window.size.t = window.size.t,
        window.size.s = window.size.s,
        make.pval = make.pval,
        method = method,
        n.sim = n.sim,
        verbose = FALSE
      )

      if (!is.null(lowerPeriod)) coh_args$lowerPeriod <- lowerPeriod
      if (!is.null(upperPeriod)) coh_args$upperPeriod <- upperPeriod

      wc <- do.call(WaveletComp::analyze.coherency, coh_args)

      period_hours <- dmwc_period_to_hours(wc$Period, target_info$unit)

      pair_results[[pair_name]] <- list(
        wavelet = wc,
        time = pair_df$date,
        dm_series = pair_df$dm,
        clim_series = pair_df$clim,
        dm_name = dm_col,
        clim_name = clim_col,
        period_native = wc$Period,
        period_hours = period_hours,
        avg_coherence = if (!is.null(wc$Coherence.avg)) as.numeric(wc$Coherence.avg) else NA_real_,
        avg_cross_power = if (!is.null(wc$Power.xy.avg)) as.numeric(wc$Power.xy.avg) else NA_real_
      )

      aligned_list[[idx]] <- tibble::tibble(
        TIME = pair_df$date,
        pair = pair_name,
        dm_series = pair_df$dm,
        clim_series = pair_df$clim,
        dm_name = dm_col,
        clim_name = clim_col
      )

      avg_coh <- pair_results[[pair_name]]$avg_coherence
      if (all(is.na(avg_coh))) {
        dom_period <- NA_real_
        max_avg_coh <- NA_real_
      } else {
        jj <- which.max(avg_coh)
        dom_period <- period_hours[jj]
        max_avg_coh <- avg_coh[jj]
      }

      summary_rows[[idx]] <- tibble::tibble(
        pair = pair_name,
        dm_name = dm_col,
        clim_name = clim_col,
        n_obs = nrow(pair_df),
        dominant_coherence_period_hours = dom_period,
        max_average_coherence = max_avg_coh
      )

      idx <- idx + 1L
    }
  }

  out <- list(
    call = cl,
    dm_type = dm_type,
    source = source_used,
    dm_series = dm_cols,
    clim_vars = clim_cols,
    time_unit = target_info$unit,
    dt = target_info$dt,
    dt_hours = target_info$hours,
    resolution_seconds = target_sec,
    overlap = c(start = as.character(t_start), end = as.character(t_end)),
    aligned_data = dplyr::bind_rows(aligned_list),
    results = pair_results,
    pair_summary = dplyr::bind_rows(summary_rows),
    settings = list(
      dm_fun = dm_fun,
      clim_funs = clim_funs,
      loess_span = loess_span,
      dj = dj,
      lowerPeriod = lowerPeriod,
      upperPeriod = upperPeriod,
      window.type.t = window.type.t,
      window.type.s = window.type.s,
      window.size.t = window.size.t,
      window.size.s = window.size.s,
      make.pval = make.pval,
      method = method,
      n.sim = n.sim,
      na_action = na_action
    )
  )

  class(out) <- "dm_wavelet_coherence"

  if (isTRUE(verbose)) {
    message(
      "dm_wavelet_coherence completed for ", length(pair_results),
      " pair(s). Coarser detected resolution: ",
      out$dt, " ", out$time_unit, " (", round(out$dt_hours, 4), " hours)."
    )
  }

  out
}


# plot method ------------------------------------------------------------------

#' Plot method for dm_wavelet_coherence objects
#'
#' @param x An object of class \code{"dm_wavelet_coherence"}.
#' @param y Unused.
#' @param pair Optional pair name(s) to plot. If \code{NULL}, the first pair is used.
#' @param type One of \code{"coherence"}, \code{"cross_power"},
#'   \code{"average_coherence"}, \code{"average_cross_power"}, \code{"phase"},
#'   or \code{"series"}.
#' @param facet Logical. If \code{TRUE} and multiple pairs are selected, facet them.
#' @param log_period Logical.
#' @param log_power Logical for raster intensity.
#' @param clip_quantile Optional clipping quantiles for raster intensity.
#' @param show_sig Logical.
#' @param show_coi Logical.
#' @param siglvl Significance level.
#' @param sig_color Significance contour color.
#' @param sig_size Significance contour linewidth.
#' @param coi_fill COI fill color.
#' @param coi_alpha COI alpha.
#' @param main Optional title.
#' @param ... Unused.
#'
#' @return A \code{ggplot2} object.
#'
#' @method plot dm_wavelet_coherence
#' @importFrom dplyr bind_rows
#' @importFrom tibble tibble
#' @importFrom ggplot2 ggplot aes geom_raster geom_line geom_contour geom_segment
#'   geom_polygon facet_wrap theme_bw theme element_text labs scale_y_log10
#'   scale_x_log10 scale_fill_viridis_c coord_cartesian
#' @export
plot.dm_wavelet_coherence <- function(x,
                                      y = NULL,
                                      pair = NULL,
                                      type = c("coherence", "cross_power", "average_coherence",
                                               "average_cross_power", "phase", "series"),
                                      facet = TRUE,
                                      log_period = TRUE,
                                      log_power = TRUE,
                                      clip_quantile = c(0.01, 0.99),
                                      show_sig = TRUE,
                                      show_coi = TRUE,
                                      siglvl = 0.05,
                                      sig_color = "black",
                                      sig_size = 0.4,
                                      coi_fill = "white",
                                      coi_alpha = 0.45,
                                      main = NULL,
                                      ...) {
  TIME <- value <- signal <- period_hours <- pval <- TIME2 <- period2 <- NULL


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

  if (!requireNamespace("ggplot2", quietly = TRUE)) {
    stop("Package 'ggplot2' is required for plot.dm_wavelet_coherence().")
  }

  type <- match.arg(type)

  if (is.null(pair)) {
    pair <- names(x$results)[1]
  }

  miss <- setdiff(pair, names(x$results))
  if (length(miss) > 0) {
    stop("These requested pairs were not found: ", paste(miss, collapse = ", "))
  }

  # aligned series
  if (type == "series") {
    dat <- x$aligned_data
    dat <- dat[dat$pair %in% pair, , drop = FALSE]

    long_dat <- tibble::tibble(
      TIME = c(dat$TIME, dat$TIME),
      pair = c(dat$pair, dat$pair),
      value = c(dat$dm_series, dat$clim_series),
      signal = c(rep("dendrometer", nrow(dat)), rep("climate", nrow(dat)))
    )

    p <- ggplot2::ggplot(long_dat, ggplot2::aes(x = TIME, y = value, colour = signal)) +
      ggplot2::geom_line() +
      ggplot2::theme_bw() +
      ggplot2::theme(
        legend.position = "right",
        axis.title = ggplot2::element_text(size = 14),
        axis.text = ggplot2::element_text(size = 11)
      ) +
      ggplot2::labs(
        x = "Time",
        y = "Value",
        colour = "Series",
        title = if (is.null(main)) "Aligned dendrometer and climate series" else main
      )

    if (length(pair) > 1 && isTRUE(facet)) {
      p <- p + ggplot2::facet_wrap(stats::as.formula("~ pair"), scales = "free_y", ncol = 1)
    }

    return(p)
  }

  # average lines
  if (type %in% c("average_coherence", "average_cross_power")) {
    df_list <- lapply(pair, function(pp) {
      res <- x$results[[pp]]
      tibble::tibble(
        pair = pp,
        period_hours = res$period_hours,
        value = if (type == "average_coherence") res$avg_coherence else res$avg_cross_power
      )
    })
    dat <- dplyr::bind_rows(df_list)

    p <- ggplot2::ggplot(dat, ggplot2::aes(x = period_hours, y = value, colour = pair)) +
      ggplot2::geom_line() +
      ggplot2::theme_bw() +
      ggplot2::theme(
        legend.position = if (length(pair) > 1) "right" else "none",
        axis.title = ggplot2::element_text(size = 14),
        axis.text = ggplot2::element_text(size = 11)
      ) +
      ggplot2::labs(
        x = "Period (hours)",
        y = if (type == "average_coherence") "Average coherence" else "Average cross-wavelet power",
        colour = "Pair",
        title = if (is.null(main)) {
          if (type == "average_coherence") "Average wavelet coherence" else "Average cross-wavelet power"
        } else main
      )

    if (isTRUE(log_period) && all(dat$period_hours > 0, na.rm = TRUE)) {
      p <- p + ggplot2::scale_x_log10()
    }

    if (length(pair) > 1 && isTRUE(facet)) {
      p <- p + ggplot2::facet_wrap(stats::as.formula("~ pair"), scales = "free_y", ncol = 1)
    }

    return(p)
  }

  # raster-like plots
  df_list <- lapply(pair, function(pp) {
    res <- x$results[[pp]]
    wc <- res$wavelet

    if (type == "coherence") {
      mat <- as.matrix(wc$Coherence)
      pmat <- if (!is.null(wc$Coherence.pval)) as.matrix(wc$Coherence.pval) else NULL
    } else if (type == "cross_power") {
      mat <- as.matrix(wc$Power.xy)
      pmat <- if (!is.null(wc$Power.xy.pval)) as.matrix(wc$Power.xy.pval) else NULL
    } else { # phase
      mat <- as.matrix(wc$Angle)
      pmat <- if (!is.null(wc$Coherence.pval)) as.matrix(wc$Coherence.pval) else NULL
    }

    n_p <- min(nrow(mat), length(res$period_hours))
    n_t <- min(ncol(mat), length(res$time))
    mat <- mat[seq_len(n_p), seq_len(n_t), drop = FALSE]
    per <- res$period_hours[seq_len(n_p)]
    tt <- res$time[seq_len(n_t)]

    df <- expand.grid(
      period_hours = per,
      TIME = tt,
      KEEP.OUT.ATTRS = FALSE,
      stringsAsFactors = FALSE
    )
    df$value <- c(mat)
    df$pair <- pp
    if (!is.null(pmat)) {
      pmat <- pmat[seq_len(n_p), seq_len(n_t), drop = FALSE]
      df$pval <- c(pmat)
    } else {
      df$pval <- NA_real_
    }
    df
  })

  dat <- dplyr::bind_rows(df_list)
  dat <- dat[is.finite(dat$period_hours) & dat$period_hours > 0, , drop = FALSE]
  y_rng <- range(dat$period_hours, na.rm = TRUE)

  plot_val <- dat$value
  fill_lab <- if (type == "coherence") "Coherence" else if (type == "cross_power") "Cross-power" else "Phase (rad)"

  if (type == "cross_power") {
    if (!is.null(clip_quantile)) {
      qv <- stats::quantile(plot_val, probs = clip_quantile, na.rm = TRUE, names = FALSE)
      plot_val <- pmin(pmax(plot_val, qv[1]), qv[2])
    }
    if (isTRUE(log_power)) {
      eps <- max(min(plot_val[plot_val > 0], na.rm = TRUE) / 10, 1e-12)
      if (!is.finite(eps)) eps <- 1e-12
      plot_val <- log10(plot_val + eps)
      fill_lab <- "log10(Cross-power)"
    }
  }

  dat$plot_val <- plot_val

  p <- ggplot2::ggplot(dat, ggplot2::aes(x = TIME, y = period_hours, fill = plot_val)) +
    ggplot2::geom_raster() +
    ggplot2::scale_fill_viridis_c(option = "D", na.value = "grey85") +
    ggplot2::theme_bw() +
    ggplot2::theme(
      legend.position = "right",
      axis.title = ggplot2::element_text(size = 14),
      axis.text = ggplot2::element_text(size = 11)
    ) +
    ggplot2::labs(
      x = "Time",
      y = "Period (hours)",
      fill = fill_lab,
      title = if (is.null(main)) {
        switch(type,
               coherence = "Wavelet coherence",
               cross_power = "Cross-wavelet power",
               phase = "Wavelet phase difference")
      } else main
    )

  # COI overlay
  if (isTRUE(show_coi)) {
    coi_list <- lapply(pair, function(pp) {
      res <- x$results[[pp]]
      wc <- res$wavelet
      if (is.null(wc$coi.1) || is.null(wc$coi.2) || is.null(wc$axis.1) || is.null(wc$axis.2)) return(NULL)

      tz_use <- attr(res$time, "tzone")
      if (is.null(tz_use) || length(tz_use) == 0) tz_use <- "UTC"

      x_num <- stats::approx(
        x = as.numeric(wc$axis.1),
        y = as.numeric(res$time),
        xout = as.numeric(wc$coi.1),
        rule = 2
      )$y

      period_native <- 2^as.numeric(wc$coi.2)
      period_hours <- dmwc_period_to_hours(period_native, x$time_unit)
      period_hours <- pmin(pmax(period_hours, y_rng[1]), y_rng[2])

      tibble::tibble(
        TIME = as.POSIXct(x_num, origin = "1970-01-01", tz = tz_use),
        period_hours = period_hours,
        pair = pp
      )
    })

    coi_dat <- dplyr::bind_rows(coi_list)
    if (nrow(coi_dat) > 0) {
      p <- p + ggplot2::geom_polygon(
        data = coi_dat,
        mapping = ggplot2::aes(x = TIME, y = period_hours, group = pair),
        inherit.aes = FALSE,
        fill = coi_fill,
        alpha = coi_alpha,
        colour = NA
      )
    }
  }

  # significance contour
  if (isTRUE(show_sig) && any(is.finite(dat$pval))) {
    sig_dat <- dat[is.finite(dat$pval), , drop = FALSE]
    p <- p + ggplot2::geom_contour(
      data = sig_dat,
      mapping = ggplot2::aes(x = TIME, y = period_hours, z = pval),
      breaks = siglvl,
      colour = sig_color,
      linewidth = sig_size,
      inherit.aes = FALSE
    )
  }

  # phase arrows
  if (type == "phase") {
    phase_list <- lapply(pair, function(pp) {
      res <- x$results[[pp]]
      wc <- res$wavelet
      ang <- as.matrix(wc$Angle)
      coh <- if (!is.null(wc$Coherence)) as.matrix(wc$Coherence) else NULL

      n_p <- min(nrow(ang), length(res$period_hours))
      n_t <- min(ncol(ang), length(res$time))

      # thin grid for plotting arrows
      p_idx <- unique(round(seq(1, n_p, length.out = min(20, n_p))))
      t_idx <- unique(round(seq(1, n_t, length.out = min(30, n_t))))

      grd <- expand.grid(pi = p_idx, ti = t_idx, KEEP.OUT.ATTRS = FALSE)

      ph <- ang[cbind(grd$pi, grd$ti)]
      if (!is.null(coh)) {
        cv <- coh[cbind(grd$pi, grd$ti)]
      } else {
        cv <- rep(1, length(ph))
      }

      tibble::tibble(
        TIME = res$time[grd$ti],
        period_hours = res$period_hours[grd$pi],
        angle = ph,
        coherence = cv,
        pair = pp
      )
    })

    phase_dat <- dplyr::bind_rows(phase_list)
    phase_dat <- phase_dat[is.finite(phase_dat$angle) &
                             is.finite(phase_dat$period_hours) &
                             phase_dat$period_hours > 0 &
                             is.finite(phase_dat$coherence) &
                             phase_dat$coherence >= 0.5, , drop = FALSE]

    if (nrow(phase_dat) > 0) {
      # arrow size in data units
      xspan <- as.numeric(diff(range(dat$TIME, na.rm = TRUE)), units = "secs")
      yspan <- diff(y_rng)
      dx <- (xspan / 40) * cos(phase_dat$angle)
      dy <- (yspan / 40) * sin(phase_dat$angle)

      phase_dat$TIME2 <- phase_dat$TIME + dx
      phase_dat$period2 <- phase_dat$period_hours + dy

      p <- p + ggplot2::geom_segment(
        data = phase_dat,
        mapping = ggplot2::aes(x = TIME, y = period_hours, xend = TIME2, yend = period2),
        inherit.aes = FALSE,
        colour = "black",
        linewidth = 0.25,
        arrow = grid::arrow(length = grid::unit(0.08, "cm"))
      )
    }
  }

  if (isTRUE(log_period) && all(y_rng > 0)) {
    p <- p + ggplot2::scale_y_log10(limits = y_rng)
  } else {
    p <- p + ggplot2::coord_cartesian(ylim = y_rng, expand = FALSE)
  }

  if (length(pair) > 1 && isTRUE(facet)) {
    p <- p + ggplot2::facet_wrap(stats::as.formula("~ pair"), scales = "free_x", ncol = 1)
  }

  p
}
