# Workbook 14: model assessment for the monthly SOI/recruitment example.
# Source this file to define functions. No data are loaded or models fitted on source.
# Matrices A[[k]] follow Workbook 13: row i predicts series i from lagged series j.

wb14_fish_data <- function() {
  if (!requireNamespace("astsa", quietly = TRUE)) stop("The astsa package is required.")
  e <- new.env(parent = baseenv())
  utils::data("soi", "rec", package = "astsa", envir = e)
  stopifnot(identical(tsp(e$soi), tsp(e$rec)), !anyNA(e$soi), !anyNA(e$rec))
  list(y = cbind(soi = as.numeric(e$soi), rec = as.numeric(e$rec)),
       dates = seq(as.Date("1950-01-01"), by = "month", length.out = length(e$soi)),
       month = as.integer(cycle(e$soi)))
}

wb14_fit_var <- function(data, p, origin = nrow(data$y), response_start = 13,
                         seasonal = TRUE) {
  y <- as.matrix(data$y); d <- ncol(y)
  stopifnot(p >= 1, p == as.integer(p), response_start > p,
            origin >= response_start, origin <= nrow(y), !anyNA(y[1:origin, ]))
  rows <- response_start:origin
  deterministic <- if (seasonal) diag(12)[data$month[rows], , drop = FALSE] else
    matrix(1, length(rows), 1)
  colnames(deterministic) <- if (seasonal) month.abb else "constant"
  lagged <- lapply(seq_len(p), function(k) {
    a <- y[rows - k, , drop = FALSE]
    colnames(a) <- paste0(colnames(y), ".L", k); a
  })
  x <- cbind(deterministic, do.call(cbind, lagged))
  ols <- lm.fit(x, y[rows, , drop = FALSE])
  if (ols$rank < ncol(x)) stop("The lag/seasonal design is rank deficient.")
  b <- matrix(ols$coefficients, nrow = ncol(x), dimnames = list(colnames(x), colnames(y)))
  residuals <- matrix(ols$residuals, ncol = d, dimnames = list(NULL, colnames(y)))
  N <- nrow(residuals); Sigma <- crossprod(residuals) / N
  logdet <- determinant(Sigma, logarithm = TRUE)
  if (logdet$sign <= 0) stop("The fitted innovation covariance is not positive definite.")
  ndet <- ncol(deterministic)
  A <- lapply(seq_len(p), function(k) t(b[ndet + (k - 1) * d + seq_len(d), , drop = FALSE]))
  companion <- do.call(cbind, A)
  if (p > 1) companion <- rbind(companion, cbind(diag(d * (p - 1)), matrix(0, d * (p - 1), d)))
  k <- length(b) + d * (d + 1) / 2
  logLik <- -N / 2 * (d * (1 + log(2 * pi)) + as.numeric(logdet$modulus))
  list(p = p, d = d, origin = origin, response_start = response_start, seasonal = seasonal,
       coefficients = b, A = A, residuals = residuals, dates = data$dates[rows],
       Sigma = Sigma, correlation = cov2cor(Sigma), n = N, k = k, logLik = logLik,
       AIC = -2 * logLik + 2 * k, BIC = -2 * logLik + log(N) * k,
       max_root = max(Mod(eigen(companion, only.values = TRUE)$values)))
}

# C(h)[i,j] uses residual i at the later time and residual j h months earlier.
# This is the Gamma(h) convention in Workbook 10, including C(-h)=t(C(h)).
wb14_residual_correlations <- function(fit, H = 24) {
  e <- fit$residuals; N <- nrow(e); d <- ncol(e)
  stopifnot(H >= 0, H == as.integer(H), H < N)
  lag <- -H:H
  C <- array(0, c(d, d, length(lag)), dimnames = list(colnames(e), colnames(e), lag))
  C[, , H + 1] <- crossprod(e) / N
  if (H > 0) for (h in seq_len(H)) {
    ch <- crossprod(e[(h + 1):N, , drop = FALSE], e[1:(N - h), , drop = FALSE]) / N
    C[, , H + 1 + h] <- ch; C[, , H + 1 - h] <- t(ch)
  }
  Sigma <- crossprod(e) / N
  scale <- sqrt(outer(diag(Sigma), diag(Sigma)))
  R <- C
  for (j in seq_along(lag)) R[, , j] <- C[, , j] / scale
  list(lag = lag, covariance = C, correlation = R)
}

wb14_portmanteau <- function(fit, H = 24) {
  stopifnot(H > fit$p, H < fit$n)
  ac <- wb14_residual_correlations(fit, H)
  N <- fit$n; inv <- solve(crossprod(fit$residuals) / N)
  terms <- sapply(seq_len(H), function(h) {
    ch <- matrix(ac$covariance[, , H + 1 + h], fit$d, fit$d)
    sum(diag(t(ch) %*% inv %*% ch %*% inv))
  })
  Q <- N^2 * sum(terms / (N - seq_len(H)))
  df <- fit$d^2 * (H - fit$p)
  data.frame(Q = Q, df = df, p_value = pchisq(Q, df, lower.tail = FALSE))
}

wb14_fit_table <- function(fits, H = 24) {
  do.call(rbind, lapply(fits, function(f) {
    q <- wb14_portmanteau(f, H)
    data.frame(Model = paste0("VAR(", f$p, ")"), p = f$p, N = f$n, k = f$k,
               logLik = f$logLik, AIC = f$AIC, BIC = f$BIC, Max_root = f$max_root,
               Q = q$Q, df = q$df, p_value = q$p_value)
  }))
}

# Marginal standardization preserves contemporaneous cross-series correlation.
wb14_series_labels <- function(x) {
  labels <- c(soi = "SOI", rec = "Recruitment")
  known <- match(x, names(labels)); x[!is.na(known)] <- labels[known[!is.na(known)]]
  x
}

wb14_plot_residuals <- function(fit, data_label = "SOI and recruitment", time_label = "Month") {
  old <- par(mfrow = c(fit$d, 1), mar = c(2.5, 4, 2.3, 1), oma = c(2, 0, 2.2, 0),
             mgp = c(2, .6, 0), cex = .9, cex.axis = .95, cex.lab = 1,
             cex.main = 1.05); on.exit(par(old))
  z <- sweep(fit$residuals, 2, sqrt(colMeans(fit$residuals^2)), "/")
  colors <- c("#00539B", "#C65D21")
  for (j in seq_len(fit$d)) {
    plot(fit$dates, z[, j], type = "l", col = colors[(j - 1) %% 2 + 1],
         xlab = "", ylab = "Standardized residual", main = wb14_series_labels(colnames(z))[j])
    abline(h = 0, col = "gray60")
  }
  mtext(paste0("Real data: ", data_label, " | VAR(", fit$p, ")"), outer = TRUE, font = 2)
  mtext(time_label, side = 1, outer = TRUE, line = .4)
  invisible(z)
}

wb14_plot_correlations <- function(fit, H = 24, lag_unit = "months", data_label = "Real data") {
  ac <- wb14_residual_correlations(fit, H); d <- fit$d
  old <- par(mfrow = c(d, d), mar = c(3.2, 3.6, 2.9, 1), oma = c(2.8, 0, 2.5, 0),
             mgp = c(2, .6, 0), cex = .9, cex.axis = .95, cex.lab = .95,
             cex.main = .95); on.exit(par(old))
  labels <- wb14_series_labels(colnames(fit$residuals))
  shown <- ac$correlation
  for (j in seq_len(d)) shown[j, j, H + 1] <- NA_real_
  lim <- max(.2, abs(shown), na.rm = TRUE) * 1.1
  for (i in seq_len(d)) for (j in seq_len(d)) {
    plot(ac$lag, shown[i, j, ], type = "h", ylim = c(-lim, lim), lwd = 1.4,
         col = "#00539B", xlab = paste0("Signed lag h (", lag_unit, ")"), ylab = "Correlation",
         main = paste0(labels[i], " at t+h\n", labels[j], " at t"))
    abline(h = 0, col = "gray55"); abline(v = 0, col = "gray80", lty = 3)
    abline(h = c(-1, 1) * 1.96 / sqrt(fit$n), lty = 2, col = "#C65D21")
  }
  model_label <- if (fit$p > 0) paste0("VAR(", fit$p, ") residual ") else ""
  mtext(paste0(data_label, ": ", model_label, "auto- and cross-correlations"), outer = TRUE, font = 2)
  mtext("Positive h: the first series is later. Diagonal h=0 omitted.",
        side = 1, outer = TRUE, line = .35, cex = .8)
  mtext("Dashed bands are visual guides.", side = 1, outer = TRUE, line = 1.35, cex = .8)
  invisible(ac)
}

wb14_forecast <- function(fit, data, h = 12) {
  stopifnot(h >= 1, h == as.integer(h))
  y <- as.matrix(data$y)[seq_len(fit$origin), colnames(fit$residuals), drop = FALSE]
  predictions <- matrix(NA_real_, h, fit$d, dimnames = list(NULL, colnames(y)))
  for (step in seq_len(h)) {
    next_month <- ((data$month[fit$origin] + step - 1) %% 12) + 1
    deterministic <- if (fit$seasonal) diag(12)[next_month, ] else 1
    lags <- unlist(lapply(seq_len(fit$p), function(k) y[nrow(y) + 1 - k, ]))
    predictions[step, ] <- as.numeric(c(deterministic, lags) %*% fit$coefficients)
    y <- rbind(y, predictions[step, ])
  }
  if (any(!is.finite(predictions))) stop("The recursive forecast is not finite.")
  predictions
}

wb14_evaluate_fish <- function(data, initial = 302, orders = c(1, 2, 6, 12),
                              horizons = c(1, 6, 12), response_start = 13) {
  stopifnot(initial >= response_start, max(orders) < response_start,
            initial <= nrow(data$y) - max(horizons))
  origins <- initial:(nrow(data$y) - max(horizons))
  records <- selections <- statuses <- list(); nr <- ns <- nf <- 0L
  attempt <- function(expr) tryCatch(expr, error = function(e) e)
  for (origin in origins) {
    joint <- lapply(orders, function(p) attempt(wb14_fit_var(data, p, origin, response_start)))
    names(joint) <- paste0("VAR(", orders, ")")
    joint_paths <- lapply(joint, function(f) if (inherits(f, "error")) f else
      attempt(wb14_forecast(f, data, max(horizons))))
    for (target in colnames(data$y)) {
      single <- data; single$y <- data$y[, target, drop = FALSE]
      ar <- lapply(orders, function(p) attempt(wb14_fit_var(single, p, origin, response_start)))
      names(ar) <- paste0("AR(", orders, ")")
      ar_paths <- lapply(ar, function(f) if (inherits(f, "error")) f else
        attempt(wb14_forecast(f, single, max(horizons))))
      paths <- c(lapply(joint_paths, function(a) if (inherits(a, "error")) a else a[, target]),
                 lapply(ar_paths, function(a) if (inherits(a, "error")) a else as.numeric(a)))
      for (family in c("VAR", "AR")) {
        fits <- if (family == "VAR") joint else ar
        for (label in names(fits)) {
          f <- fits[[label]]; failed <- inherits(f, "error"); nf <- nf + 1L
          statuses[[nf]] <- data.frame(origin = origin, target = target, model = label,
            max_root = if (failed) NA_real_ else f$max_root,
            failure = if (failed) conditionMessage(f) else "")
        }
        for (criterion in c("AIC", "BIC")) {
          value <- sapply(fits, function(f) if (inherits(f, "error")) Inf else f[[criterion]])
          chosen <- if (all(!is.finite(value))) NA_integer_ else which.min(value)
          label <- paste(family, criterion, sep = "-")
          paths[[label]] <- if (is.na(chosen)) simpleError("Every candidate fit failed.") else
            paths[[names(fits)[chosen]]]
          ns <- ns + 1L
          selections[[ns]] <- data.frame(origin = origin, date = data$dates[origin],
            target = target, family = family, criterion = criterion,
            p = if (is.na(chosen)) NA_integer_ else orders[chosen])
        }
      }
      means <- tapply(single$y[seq_len(origin), 1], data$month[seq_len(origin)], mean)
      future_month <- ((data$month[origin] + seq_len(max(horizons)) - 1) %% 12) + 1
      paths[["Seasonal mean"]] <- means[future_month]
      for (model in names(paths)) for (h in horizons) {
        path <- paths[[model]]; failed <- inherits(path, "error"); nr <- nr + 1L
        records[[nr]] <- data.frame(origin = origin, origin_date = data$dates[origin],
          target_date = data$dates[origin + h], target = target, model = model, h = h,
          actual = data$y[origin + h, target], forecast = if (failed) NA_real_ else as.numeric(path[h]),
          failure = if (failed) conditionMessage(path) else "")
      }
    }
  }
  forecasts <- do.call(rbind, records)
  forecasts$error <- forecasts$actual - forecasts$forecast
  groups <- split(forecasts, list(forecasts$target, forecasts$model, forecasts$h), drop = TRUE)
  scores <- do.call(rbind, lapply(groups, function(a) {
    failed <- any(!is.finite(a$error))
    data.frame(target = a$target[1], model = a$model[1], h = a$h[1], n = nrow(a),
      failures = sum(!is.finite(a$error)), RMSE = if (failed) NA_real_ else sqrt(mean(a$error^2)),
      MAE = if (failed) NA_real_ else mean(abs(a$error)))
  }))
  rownames(scores) <- NULL
  list(forecasts = forecasts, scores = scores[order(scores$target, scores$h, scores$RMSE), ],
       selections = do.call(rbind, selections), status = do.call(rbind, statuses),
       origins = origins, horizons = horizons, initial = initial)
}

wb14_plot_rmse <- function(evaluation, target = "rec",
                           models = c("VAR-BIC", "AR-BIC", "VAR-AIC", "Seasonal mean")) {
  scores <- evaluation$scores[evaluation$scores$target == target, ]
  h <- sort(unique(scores$h))
  values <- sapply(models, function(m) scores$RMSE[match(paste(m, h), paste(scores$model, scores$h))])
  palette <- c("VAR-BIC" = "#00539B", "AR-BIC" = "#C65D21", "VAR-AIC" = "#599BAB",
               "Seasonal mean" = "#686F62", "AR-AIC" = "#DAA078")
  colors <- unname(palette[models]); colors[is.na(colors)] <- "gray40"
  symbols <- rep(c(16, 17, 15, 18, 19), length.out = length(models))
  old <- par(mar = c(4.2, 4.3, 3.7, 1), mgp = c(2.4, .7, 0), cex = .9); on.exit(par(old))
  matplot(h, values, type = "b", pch = symbols, lty = 1, lwd = 2, col = colors,
    ylim = c(min(values) * .92, max(values) * 1.23), xaxt = "n",
    xlab = "Forecast horizon (months)", ylab = "RMSE (index units)",
    main = paste("Real data:", wb14_series_labels(target), "forecast errors"))
  axis(1, at = h)
  legend("topleft", legend = models, col = colors, pch = symbols, lty = 1, lwd = 2, ncol = 2, bty = "n", cex = .9)
  mtext(paste(length(evaluation$origins), "shared expanding origins; AIC/BIC reselect at each origin"),
        side = 3, line = .5, cex = .85)
  invisible(values)
}

wb14_plot_forecasts <- function(evaluation, target = "rec", h = 6,
                                models = c("VAR-BIC", "AR-BIC")) {
  f <- evaluation$forecasts
  f <- f[f$target == target & f$h == h & f$model %in% models, ]
  observed <- f[f$model == models[1], ]; colors <- c("#00539B", "#C65D21", "#686F62")
  old <- par(mar = c(4.2, 4.3, 3.2, 1), mgp = c(2.4, .7, 0), cex = .9); on.exit(par(old))
  limits <- range(c(f$actual, f$forecast), finite = TRUE)
  limits[2] <- limits[2] + .2 * diff(limits)
  plot(observed$target_date, observed$actual, type = "l", col = "gray25", lwd = 1.4,
       ylim = limits, xlab = "Target month", ylab = paste(wb14_series_labels(target), "index"),
       main = "")
  title(paste0("Real data: ", h, "-month-ahead ", tolower(wb14_series_labels(target)),
               " forecasts"), line = 1.5)
  for (j in seq_along(models)) {
    a <- f[f$model == models[j], ]
    lines(a$target_date, a$forecast, col = colors[j], lwd = 1.6)
  }
  legend("topleft", c("Observed", models), col = c("gray25", colors[seq_along(models)]),
         lty = 1, lwd = 2, bty = "n", cex = .9, horiz = TRUE)
  invisible(f)
}