#' Timescale Decomposition
#'
#' @author Javier Corvillo, \email{javier.corvillo@bsc.es}
#'
#' @references Campos, D. A., Olmo, M. E., Cos, P., Muñoz, Á. G., & Doblas-Reyes,
#' F. J. (2025). Regional aspects of observed temperature and precipitation trends
#' in the western Mediterranean: Insights from a timescale decomposition analysis.
#' Journal of Geophysical Research: Atmospheres, 130, e2024JD042637.
#' https://doi.org/10.1029/2024JD042637.
#'
#' @description This function decomposes grid point time series into trend, 
#' decadal and residual components using LOESS smoothing and Butterworth filtering.
#' If the frequency of the given time series object is higher than yearly data 
#' (e.g daily, monthly data), the seasonal component is also calculated. The 
#' function returns either variance percentages for each component or their values
#' across spatial grids.
#'
#' @param data A numeric array containing time-series data. The array must
#' include at least the named dimension specified in 'ftime_dim'.
#'
#' @param ftime_dim A character string indicating the name of forecast time
#' dimension. The default value is 'time'.
#'
#' @param fq An integer indicating the number of observations per year (default
#' 12, for monthly data). The decadal cutoff for the Butterworth filter is 10
#' times the 'fq' value.
#'
#' @param lambda A non-negative integer that controls how close the optimal span
#' selection in the LOESS smoothing will be to the decadal cutoff. The default
#' value is 1. Values of lambda close to zero may result in the trend component
#' capturing more of the decadal component (trend is smooth and less linear),
#' while higher values may result in the trend component capturing less of the
#' decadal component (trend is stringent and more linear).
#'
#' @param method Character string indicating the method to use for selecting the
#' optimal span. Must be either \code{"gcv"} (Generalized Correlated
#' Cross-Validation) (default), \code{"aic"} (Akaike Information Criterion), or
#' \code{"rsq"} (R Squared).
#'
#' @param return_type Character string indicating the output format. Must be
#' either \code{"percentages"} (default) for variance percentages of each
#' component or \code{"values"} for actual component time series.
#'
#' @param ncores An integer indicating the number of cores to use for parallel
#' computation. The default value is NULL.
#'
#' @return A named list with the following elements:
#' \describe{
#'   \item{pct_trend / comp_trend}{Trend component:
#'     \itemize{
#'       \item If \code{return_type = "percentages"}: percentages of each component.
#'       \item If \code{return_type = "values"}: time series of each component.
#'     }}
#'   \item{pct_seasonal / comp_seasonal}{Seasonal component (same dimensional structure as above).}
#'   \item{pct_decadal / comp_decadal}{Decadal component (same dimensional structure as above).}
#'   \item{pct_residual / comp_residual}{Residual component (same dimensional structure as above).}
#' }
#' 
#' @examples
#' # Time-based decomposition
#' data <- array(rnorm(120 * 3 * 2), dim = c(time = 120, lat = 3, lon = 2))
#'
#' # Add some variability to the data
#' data <- data +
#'   array(rep(seq(1, 120) * 0.05, each = 3 * 2),
#'         dim = c(time = 120, lat = 3, lon = 2)) + # Trend
#'   array(rep(sin(seq(1, 120) / 12 * 2 * pi) * 2, each = 3 * 2),
#'         dim = c(time = 120, lat = 3, lon = 2)) + # Seasonal
#'   array(rep(sin(seq(1, 120) / 60 * 2 * pi) * 1.5, each = 3 * 2),
#'         dim = c(time = 120, lat = 3, lon = 2)) # Decadal
#'
#' result <- TimescaleDecomposition(data)
#' str(result)
#'
#' @importFrom multiApply Apply
#' @importFrom zoo na.approx
#' @importFrom signal butter filtfilt
#' @importFrom stats var ts stl loess median
#' @export

TimescaleDecomposition <- function(data, ftime_dim = "time", fq = 12,
                                   lambda = 1, method = "gcv",
                                   return_type = "percentages", ncores = NULL) {
  # ===== Function checks =====

  # Check data is a numeric array
  if (!is.numeric(data) || !is.array(data)) {
    stop("'data' must be a numeric array.")
  }

  # Check if ftime_dim is a string
  if (!is.character(ftime_dim) || length(ftime_dim) != 1) {
    stop("'ftime_dim' must be a single character string.")
  }

  # Check if ftime_dim is in data
  if (!(ftime_dim %in% names(dim(data)))) {
    stop(sprintf("The forecast time dimension '%s' is not found in 'data'.", ftime_dim))
  }

  # Checks for fq
  if (!is.numeric(fq) || fq <= 0 || fq != round(fq)) {
    stop("'fq' must be a positive integer")
  }

  # Checks for lambda
  if (!is.numeric(lambda) || lambda < 0) {
    stop("'lambda' must be a positive value")
  }

  decadal_cutoff <- fq * 10

  if (!is.numeric(decadal_cutoff) ||
    decadal_cutoff <= 0 ||
    decadal_cutoff != round(decadal_cutoff) ||
    decadal_cutoff > dim(data)[ftime_dim]) {
    stop("'data' needs to have at least 10 years of values.")
  }

  if (method != "gcv" && method != "aic" && method != "rsq") {
    stop("'method' must be either 'gcv', 'aic', or 'rsq'")
  }

  if (return_type != "percentages" && return_type != "values") {
    stop("'return_type' must be either 'percentages' or 'values'")
  }

  res <- Apply(
    data = list(data = data),
    target_dims = list(c(ftime_dim)),
    fun = .TimescaleDecomposition,
    fq = fq,
    method = method,
    lambda = lambda,
    decadal_cutoff = decadal_cutoff,
    return_type = return_type,
    ncores = ncores
  )

  return(res)
}

.TimescaleDecomposition <- function(data, fq,
                                    lambda,
                                    method,
                                    decadal_cutoff,
                                    return_type) {
  # Loop through each grid point and decompose the time series into trend, seasonal,
  # and decadal components

  if (all(is.na(data))) {
    pct_trend <- NA
    pct_seasonal <- NA
    pct_decadal <- NA
    pct_residual <- NA
    comp_trend <- NA
    comp_seasonal <- NA
    comp_decadal <- NA
    comp_residual <- NA
  } else {
    data <- ifelse(is.na(data), 0, data)
    optimal_span <- .select_optimal_span(data, method, decadal_cutoff, lambda)

    # Create time series object
    signal_ts <- ts(data, frequency = fq, start = c(1, 1))

    if (fq == 1) {
      # Performing LOESS manually for yearly data:
      df <- data.frame(time = seq_along(signal_ts), data = signal_ts)

      model <- suppressWarnings(
        loess(data ~ time, data = df, span = optimal_span / length(signal_ts))
      )

      trend <- predict(model, df)
      residual_signal <- signal_ts - trend

      bf <- butter(2, (1 / decadal_cutoff) / (1 / 2), type = "low")
      decadal <- filtfilt(bf, residual_signal)
      residual <- residual_signal - decadal
      total_var <- var(residual, na.rm = TRUE) +
        var(decadal, na.rm = TRUE) +
        var(trend, na.rm = TRUE)

      if (return_type == "percentages") {
        # Assign percentages to arrays
        pct_trend <- var(trend,
          na.rm = TRUE
        ) / total_var * 100
        pct_seasonal <- 0
        pct_decadal <- var(decadal,
          na.rm = TRUE
        ) / total_var * 100
        pct_residual <- var(residual,
          na.rm = TRUE
        ) / total_var * 100
      } else {
        # Store component values (time series)
        comp_trend <- trend
        comp_seasonal <- NA
        comp_decadal <- decadal
        comp_residual <- residual
      }
    } else {
      # Apply STL decomposition to extract trend and seasonal components
      components <- stl(signal_ts,
        s.window = "per",
        t.window = optimal_span, na.action = na.exclude
      )

      residual_signal <- components$time.series[, 3]

      if (any(is.na(residual_signal))) {
        # Interpolate NAs for Butterworth filter
        residual_signal <- na.approx(residual_signal, na.rm = FALSE)
      }
      # Apply Butterworth filter to extract decadal component
      bf <- butter(2, (1 / decadal_cutoff) / (1 / 2), type = "low")
      decadal <- filtfilt(bf, residual_signal)
      residual <- components$time.series[, 3] - decadal
      total_var <- var(components$time.series[, 1], na.rm = TRUE) +
        var(components$time.series[, 2], na.rm = TRUE) +
        var(residual, na.rm = TRUE) +
        var(decadal, na.rm = TRUE)
      if (return_type == "percentages") {
        # Assign percentages to arrays
        pct_trend <- var(components$time.series[, 2],
          na.rm = TRUE
        ) / total_var * 100
        pct_seasonal <- var(components$time.series[, 1],
          na.rm = TRUE
        ) / total_var * 100
        pct_decadal <- var(decadal,
          na.rm = TRUE
        ) / total_var * 100
        pct_residual <- var(residual,
          na.rm = TRUE
        ) / total_var * 100
      } else {
        # Store component values (time series)
        comp_trend <- components$time.series[, 2]
        comp_seasonal <- components$time.series[, 1]
        comp_decadal <- decadal
        comp_residual <- residual
      }
    }
  }

  # Return results as a named list based on return type
  if (return_type == "percentages") {
    list(
      pct_trend = pct_trend,
      pct_seasonal = pct_seasonal,
      pct_decadal = pct_decadal,
      pct_residual = pct_residual
    )
  } else {
    list(
      comp_trend = comp_trend,
      comp_seasonal = comp_seasonal,
      comp_decadal = comp_decadal,
      comp_residual = comp_residual
    )
  }
}


.select_optimal_span <- function(data, method, decadal_cutoff, lambda) {
  span_choices <- (decadal_cutoff + 1):length(data)
  scores <- numeric(length(span_choices))
  n <- length(data)

  # Standardize input so metric scale is comparable across variables.
  data_iqr <- stats::IQR(data, na.rm = TRUE)
  if (is.na(data_iqr) || data_iqr == 0) {
    data_scaled <- data - median(data, na.rm = TRUE)
  } else {
    data_scaled <- (data - median(data, na.rm = TRUE)) / data_iqr
  }

  df <- data.frame(time = seq_along(data_scaled), data = data_scaled)

  for (i in seq_along(span_choices)) {
    span_val <- span_choices[i]

    # 1. Fit the model
    model <- suppressWarnings(
      loess(data ~ time, data = df, span = span_val / n)
    )

    # 2. Calculate the base metric
    if (method == "gcv") {
      base_metric <- mean(model$residuals^2) / (1 - model$enp / n)^2
    } else if (method == "aic") {
      sigma2 <- sum(model$residuals^2) / n
      base_metric <- n * log(sigma2) + 2 * model$enp
    } else if (method == "rsq") {
      ss_total <- sum((df$data - mean(df$data))^2)
      base_metric <- 1 - (sum(model$residuals^2) / ss_total)
    }

    # Apply a frequency penalty based on lambda
    dist_from_cutoff <- span_val - decadal_cutoff
    penalty <- lambda / dist_from_cutoff
    if (method == "rsq") {
      scores[i] <- base_metric - (penalty / 100) # Rsq is 0-1, so scale penalty
    } else {
      scores[i] <- base_metric + penalty
    }
  }

  # 4. Find optimal
  if (method == "rsq") {
    return(span_choices[which.max(scores)])
  } else {
    return(span_choices[which.min(scores)])
  }
}
