# Workbook 14: fitted VARMA and error-correction comparisons.
# Requires astsa, vars, urca, and MTS. Sourcing this file performs no fits.

.wb14_companion_require <- function(packages) {
  missing <- packages[!vapply(packages, requireNamespace, logical(1), quietly = TRUE)]
  if (length(missing)) stop("Required R packages: ", paste(missing, collapse = ", "), call. = FALSE)
}

wb14_macro_data <- function() {
  .wb14_companion_require("astsa")
  env <- new.env()
  utils::data("econ5", package = "astsa", envir = env)
  levels <- stats::window(env$econ5, end = c(1988, 2))
  x <- cbind(growth = 100 * diff(log(levels[, "gnp"])),
             du = diff(levels[, "unemp"]))
  tt <- as.numeric(stats::time(x))
  labels <- paste0(floor(tt + 1e-8), " Q", stats::cycle(x))
  list(y = as.matrix(x), time = tt, labels = labels,
       units = c(growth = "Quarterly growth (approx. percent)",
                 du = "Unemployment change (percentage points)"),
       source = "astsa::econ5", frequency = 4L)
}

wb14_energy_data <- function() {
  .wb14_companion_require("astsa")
  env <- new.env()
  utils::data("oil", package = "astsa", envir = env)
  utils::data("gas", package = "astsa", envir = env)
  stopifnot(identical(stats::time(env$oil), stats::time(env$gas)),
            all(env$oil > 0), all(env$gas > 0))
  tt <- as.numeric(stats::time(env$oil))
  list(y = cbind(log_oil = log(as.numeric(env$oil)),
                 log_gas = log(as.numeric(env$gas))),
       time = tt,
       labels = sprintf("%d W%02d", floor(tt + 1e-8), stats::cycle(env$oil)),
       units = c(log_oil = "Log oil price", log_gas = "Log gasoline price"),
       source = "astsa::oil and astsa::gas", frequency = 52L)
}

.wb14_companion_validate <- function(data, train) {
  y <- as.matrix(data$y)
  stopifnot(is.numeric(y), ncol(y) == 2L, all(is.finite(y)),
            !is.null(colnames(y)), length(data$labels) == nrow(y),
            length(train) == 1L, train == as.integer(train),
            train >= 15L, train <= nrow(y))
  y
}

# Full profiled Gaussian likelihood, with the ML covariance divisor n.
.wb14_companion_score <- function(e, k) {
  n <- nrow(e); d <- ncol(e)
  covariance <- crossprod(e) / n
  determinant <- determinant(covariance, logarithm = TRUE)
  if (determinant$sign <= 0 || !is.finite(as.numeric(determinant$modulus)))
    stop("Singular fitted innovation covariance.")
  ll <- -n / 2 * (d * (log(2 * pi) + 1) + as.numeric(determinant$modulus))
  data.frame(n = n, k = k, logLik = ll, AIC = -2 * ll + 2 * k,
             BIC = -2 * ll + log(n) * k)
}

.wb14_companion_radius <- function(matrices) {
  d <- nrow(matrices[[1L]]); p <- length(matrices)
  top <- do.call(cbind, matrices)
  companion <- if (p == 1L) top else
    rbind(top, cbind(diag(d * (p - 1L)), matrix(0, d * (p - 1L), d)))
  max(Mod(eigen(companion, only.values = TRUE)$values))
}

.wb14_companion_portmanteau <- function(e, lag, dynamic_lags, center = TRUE) {
  n <- nrow(e); d <- ncol(e)
  stopifnot(lag > dynamic_lags, lag < n)
  if (center) e <- scale(e, center = TRUE, scale = FALSE)
  inverse <- solve(crossprod(e) / n)
  terms <- vapply(seq_len(lag), function(h) {
    ch <- crossprod(e[(h + 1L):n, , drop = FALSE], e[seq_len(n - h), , drop = FALSE]) / n
    sum(diag(t(ch) %*% inverse %*% ch %*% inverse)) / (n - h)
  }, numeric(1))
  statistic <- n^2 * sum(terms)
  df <- d^2 * (lag - dynamic_lags)
  data.frame(PT_lag = lag, PT_statistic = statistic, PT_df = df,
             PT_pvalue = stats::pchisq(statistic, df, lower.tail = FALSE))
}

.wb14_companion_fit_macro <- function(y, model) {
  n <- nrow(y)
  if (model %in% c("VAR1", "VAR2")) {
    p <- as.integer(sub("VAR", "", model))
    object <- vars::VAR(y[seq.int(3L - p, n), , drop = FALSE], p = p, type = "const")
    residual <- as.matrix(stats::residuals(object))
    fit <- list(object = object, residual = residual, k = 4L * p + 5L,
                family = "VAR", order = p, convergence = 0L, message = "",
                ar_radius = max(vars::roots(object)), ma_radius = NA_real_,
                AC_min_sv = NA_real_, AC_condition = NA_real_)
  } else if (model == "VARMA11") {
    # Expose native optimizer diagnostics without changing its fitting algorithm.
    checked <- MTS::VARMA
    body(checked) <- as.call(c(as.list(body(checked)),
      quote(VARMA$optimization <- fit[c("convergence", "message", "objective")]),
      quote(VARMA)))
    invisible(utils::capture.output(object <- checked(y[2:n, , drop = FALSE], p = 1, q = 1)))
    residual <- as.matrix(object$residuals)
    singular_values <- svd(object$Phi - object$Theta)$d
    fit <- list(object = object, residual = residual, k = 13L,
                family = "VARMA", order = 1L,
                convergence = object$optimization$convergence,
                message = object$optimization$message,
                ar_radius = max(Mod(eigen(object$Phi, only.values = TRUE)$values)),
                ma_radius = max(Mod(eigen(object$Theta, only.values = TRUE)$values)),
                AC_min_sv = min(singular_values),
                AC_condition = max(singular_values) / min(singular_values))
  } else stop("Unknown macro model: ", model)
  stopifnot(nrow(residual) == n - 2L, all(is.finite(residual)))
  colnames(fit$residual) <- colnames(y)
  fit$last <- y[n, ]; fit$model <- model
  fit$score <- .wb14_companion_score(residual, fit$k)
  fit
}

.wb14_companion_fit_energy <- function(y, model) {
  n <- nrow(y); d <- ncol(y)
  extra <- c(beta2 = NA_real_, constant = NA_real_, alpha1 = NA_real_, alpha2 = NA_real_)
  if (grepl("^VECM[234]$", model)) {
    p <- as.integer(sub("VECM", "", model))
    jo <- urca::ca.jo(y[seq.int(5L - p, n), , drop = FALSE],
                      type = "trace", ecdet = "const", K = p, spec = "transitory")
    ec <- urca::cajorls(jo, r = 1L)
    object <- vars::vec2var(jo, r = 1L)
    residual <- as.matrix(stats::residuals(ec$rlm))
    stopifnot(max(abs(jo@Z0 - diff(y)[4:(n - 1L), , drop = FALSE])) < 1e-10,
              max(abs(residual - stats::residuals(object))) < 1e-8)
    # alpha: d; normalized beta: d-1; restricted constant: 1.
    k <- d + (d - 1L) + 1L + d^2 * (p - 1L) + d * (d + 1L) / 2L
    extra <- c(beta2 = ec$beta[2L, 1L], constant = ec$beta[3L, 1L],
               alpha1 = unname(coef(ec$rlm)["ect1", 1L]),
               alpha2 = unname(coef(ec$rlm)["ect1", 2L]))
    family <- "VECM"
    ar_radius <- .wb14_companion_radius(object$A)
  } else if (grepl("^DVAR[123]$", model)) {
    p <- as.integer(sub("DVAR", "", model))
    changes <- diff(y[seq.int(4L - p, n), , drop = FALSE])
    object <- vars::VAR(changes, p = p, type = "const")
    residual <- as.matrix(stats::residuals(object))
    stopifnot(max(abs(object$datamat[, seq_len(d)] - diff(y)[4:(n - 1L), , drop = FALSE])) < 1e-10)
    k <- d^2 * p + d + d * (d + 1L) / 2L
    family <- "DVAR"; ar_radius <- max(vars::roots(object))
  } else if (model == "RW") {
    object <- NULL; residual <- diff(y)[4:(n - 1L), , drop = FALSE]
    k <- d * (d + 1L) / 2L; family <- "RW"; p <- 0L; ar_radius <- 1
  } else stop("Unknown energy model: ", model)
  stopifnot(nrow(residual) == n - 4L, all(is.finite(residual)))
  colnames(residual) <- colnames(y)
  list(object = object, residual = residual, k = k, family = family, order = p,
       model = model, last = y[n, ], extra = extra,
       convergence = 0L, message = "", ar_radius = ar_radius,
       ma_radius = NA_real_, AC_min_sv = NA_real_, AC_condition = NA_real_,
       score = .wb14_companion_score(residual, k))
}

.wb14_companion_attempt <- function(y, model, family) {
  # A missing installation is an execution error, not a failed candidate fit.
  .wb14_companion_require(c("vars", if (family == "macro") "MTS" else "urca"))
  warning_text <- character()
  fit <- tryCatch(withCallingHandlers(
    if (family == "macro") .wb14_companion_fit_macro(y, model) else .wb14_companion_fit_energy(y, model),
    warning = function(w) {
      warning_text <<- c(warning_text, conditionMessage(w))
      invokeRestart("muffleWarning")
    }), error = function(e) e)
  error <- if (inherits(fit, "error")) conditionMessage(fit) else ""
  if (!nzchar(error) && fit$convergence != 0L)
    error <- paste("Optimizer did not converge:", fit$message)
  if (!nzchar(error) && family == "macro" &&
      (fit$ar_radius >= 1 || (fit$family == "VARMA" &&
       (fit$ma_radius >= 1 || fit$AC_min_sv <= 1e-8))))
    error <- "Fit is outside the stable, invertible, identified candidate family."
  list(fit = if (inherits(fit, "error")) NULL else fit,
       ok = !nzchar(error), error = error, warnings = paste(unique(warning_text), collapse = " | "))
}

.wb14_companion_status <- function(attempt, model, origin, label) {
  f <- attempt$fit
  data.frame(model = model, origin = origin, origin_label = label,
             ok = attempt$ok, convergence = if (is.null(f)) NA_integer_ else f$convergence,
             ar_radius = if (is.null(f)) NA_real_ else f$ar_radius,
             ma_radius = if (is.null(f)) NA_real_ else f$ma_radius,
             AC_min_sv = if (is.null(f)) NA_real_ else f$AC_min_sv,
             AC_condition = if (is.null(f)) NA_real_ else f$AC_condition,
             warnings = attempt$warnings, error = attempt$error,
             stringsAsFactors = FALSE)
}

.wb14_companion_diagnostic <- function(fit, lag) {
  if (fit$family %in% c("VAR", "VECM", "DVAR")) {
    test <- vars::serial.test(fit$object, lags.pt = lag, type = "PT.adjusted")$serial
    data.frame(PT_lag = lag, PT_statistic = unname(test$statistic),
               PT_df = unname(test$parameter), PT_pvalue = unname(test$p.value))
  } else .wb14_companion_portmanteau(fit$residual, lag,
                dynamic_lags = if (fit$family == "VARMA") 2L else 0L,
                center = fit$family != "RW")
}

.wb14_companion_assessment <- function(data, train, family, lag) {
  y <- .wb14_companion_validate(data, train)
  models <- if (family == "macro") c("VAR1", "VAR2", "VARMA11") else
    c(paste0("VECM", 2:4), paste0("DVAR", 1:3), "RW")
  fits <- setNames(vector("list", length(models)), models)
  comparison <- status <- vector("list", length(models))
  for (i in seq_along(models)) {
    attempt <- .wb14_companion_attempt(y[seq_len(train), , drop = FALSE], models[i], family)
    status[[i]] <- .wb14_companion_status(attempt, models[i], train, data$labels[train])
    f <- attempt$fit; fits[i] <- list(f)
    if (attempt$ok) {
      comparison[[i]] <- cbind(data.frame(model = models[i]), f$score,
                                .wb14_companion_diagnostic(f, lag))
    } else comparison[[i]] <- data.frame(model = models[i], n = NA_integer_, k = NA_integer_,
       logLik = NA_real_, AIC = NA_real_, BIC = NA_real_, PT_lag = lag,
       PT_statistic = NA_real_, PT_df = NA_real_, PT_pvalue = NA_real_)
  }
  comparison <- do.call(rbind, comparison); status <- do.call(rbind, status)
  rownames(comparison) <- rownames(status) <- NULL
  finite_aic <- is.finite(comparison$AIC); finite_bic <- is.finite(comparison$BIC)
  comparison$deltaAIC <- if (any(finite_aic)) comparison$AIC - min(comparison$AIC[finite_aic]) else NA_real_
  comparison$deltaBIC <- if (any(finite_bic)) comparison$BIC - min(comparison$BIC[finite_bic]) else NA_real_
  list(comparison = comparison, fits = fits, status = status,
       failures = status[!status$ok, , drop = FALSE], data = data, train = train,
       response_rows = seq.int(if (family == "macro") 3L else 5L, train))
}

# Common initial response rows are 3:train for macro and 5:train for energy.
wb14_macro_assessment <- function(data = wb14_macro_data(), train = 105L, diagnostic_lag = 12L) {
  .wb14_companion_assessment(data, train, "macro", diagnostic_lag)
}

wb14_energy_assessment <- function(data = wb14_energy_data(), train = 350L, diagnostic_lag = 12L) {
  .wb14_companion_assessment(data, train, "energy", diagnostic_lag)
}

.wb14_companion_forecast <- function(fit, h) {
  d <- length(fit$last)
  if (fit$family == "RW") return(matrix(rep(fit$last, each = h), nrow = h))
  if (fit$family == "VARMA") {
    f <- fit$object; y <- fit$last
    e <- tail(f$residuals, 1L)[1L, ]; prediction <- matrix(0, h, d)
    # MTS uses a minus MA sign; future innovations have conditional mean zero.
    for (j in seq_len(h)) {
      y <- f$Ph0 + drop(f$Phi %*% y) - drop(f$Theta %*% e)
      prediction[j, ] <- y; e <- rep(0, d)
    }
  } else {
    native <- stats::predict(fit$object, n.ahead = h)$fcst
    prediction <- matrix(vapply(native, function(z) z[, "fcst"], numeric(h)), nrow = h)
    if (fit$family == "DVAR") {
      prediction <- matrix(vapply(seq_len(d), function(j) cumsum(prediction[, j]), numeric(h)), nrow = h)
      prediction <- sweep(prediction, 2L, fit$last, "+")
    }
  }
  if (!all(is.finite(prediction))) stop("Nonfinite forecast.")
  prediction
}

.wb14_companion_forecasts <- function(data, train, horizons, origins, family, common_origins) {
  y <- .wb14_companion_validate(data, train); n <- nrow(y)
  stopifnot(length(horizons) > 0L, all(horizons >= 1L), all(horizons == as.integer(horizons)),
            !anyDuplicated(horizons), train + max(horizons) <= n)
  if (is.null(origins)) origins <- seq.int(train, n - if (common_origins) max(horizons) else min(horizons))
  stopifnot(length(origins) > 0L, all(origins == as.integer(origins)), !anyDuplicated(origins),
            all(origins >= train), all(origins + min(horizons) <= n))
  models <- if (family == "macro") c("VAR1", "VAR2", "VARMA11", "Mean") else
    c(paste0("VECM", 2:4), paste0("DVAR", 1:3), "RW")
  predictions <- status <- list(); count <- 0L; fit_count <- 0L
  for (origin in origins) {
    observed <- y[seq_len(origin), , drop = FALSE]
    available <- horizons[origin + horizons <= n]
    for (model in models) {
      if (model == "Mean") {
        attempt <- list(ok = TRUE, error = "", warnings = "", fit = NULL)
        prediction <- matrix(rep(colMeans(observed), each = max(available)), nrow = max(available))
      } else {
        attempt <- .wb14_companion_attempt(observed, model, family)
        prediction <- if (attempt$ok) tryCatch(.wb14_companion_forecast(attempt$fit, max(available)),
                                               error = function(e) e) else NULL
        if (inherits(prediction, "error")) {
          attempt$ok <- FALSE; attempt$error <- conditionMessage(prediction)
        }
      }
      fit_count <- fit_count + 1L
      status[[fit_count]] <- .wb14_companion_status(attempt, model, origin, data$labels[origin])
      for (h in available) for (j in seq_len(ncol(y))) {
        value <- if (attempt$ok) prediction[h, j] else NA_real_
        count <- count + 1L
        predictions[[count]] <- data.frame(model = model, origin = origin,
          origin_label = data$labels[origin], h = h, target = origin + h,
          target_label = data$labels[origin + h], coordinate = colnames(y)[j],
          observed = y[origin + h, j], forecast = value,
          error = y[origin + h, j] - value, fit_ok = attempt$ok)
      }
    }
  }
  predictions <- do.call(rbind, predictions); status <- do.call(rbind, status)
  groups <- split(predictions, interaction(predictions$model, predictions$h, predictions$coordinate, drop = TRUE))
  scores <- do.call(rbind, lapply(groups, function(z) {
    valid <- is.finite(z$error)
    data.frame(model = z$model[1L], h = z$h[1L], coordinate = z$coordinate[1L],
      n_expected = nrow(z), n_forecasts = sum(valid), n_failed = sum(!valid),
      RMSE = if (all(valid)) sqrt(mean(z$error^2)) else NA_real_,
      MAE = if (all(valid)) mean(abs(z$error)) else NA_real_,
      bias = if (all(valid)) mean(z$error) else NA_real_)
  }))
  rownames(scores) <- NULL
  benchmark <- if (family == "macro") "Mean" else "RW"
  scores$RMSE_over_benchmark <- mapply(function(h, coordinate, rmse) {
    rmse / scores$RMSE[scores$model == benchmark & scores$h == h & scores$coordinate == coordinate]
  }, scores$h, scores$coordinate, scores$RMSE)
  list(predictions = predictions, scores = scores, status = status,
       failures = status[!status$ok, , drop = FALSE], train = train,
       origins = origins, horizons = horizons, benchmark = benchmark)
}

# Fixed candidates are refitted at expanding origins; failures retain NA rows.
# A group with failed forecasts has NA scores rather than a changed score window.
wb14_macro_forecasts <- function(data = wb14_macro_data(), train = 105L,
                                 horizons = c(1L, 4L), origins = NULL) {
  .wb14_companion_forecasts(data, train, horizons, origins, "macro", common_origins = TRUE)
}

wb14_energy_forecasts <- function(data = wb14_energy_data(), train = 350L,
                                  horizons = c(1L, 4L), origins = NULL) {
  .wb14_companion_forecasts(data, train, horizons, origins, "energy", common_origins = FALSE)
}