# Simulated gauge readings of water height in centimetres relative to a reference.
# One time step is one day; the state contains height and its daily change.
# State at day 1: mean (0, 0), variance diag(4, 0.25).
# Daily observation-error variance is 1; q is the slope-increment variance.

trend_sample <- function(n = 88L, q = 0.04, seed = 5421503L) {
  stopifnot(n >= 2L, q >= 0)
  set.seed(seed)
  level <- slope <- numeric(n)
  level[1] <- rnorm(1, sd = 2)
  slope[1] <- rnorm(1, sd = 0.5)
  for (t in 2:n) {
    level[t] <- level[t - 1L] + slope[t - 1L]
    slope[t] <- slope[t - 1L] + rnorm(1, sd = sqrt(q))
  }
  data.frame(day = seq_len(n), x = level + rnorm(n), level = level, slope = slope)
}

# F = (1, 0), so the innovation variance uses only P[1,1] and the
# covariance of each state coordinate with the reading is P[,1].
trend_filter <- function(x, q) {
  stopifnot(length(x) > 0L, all(is.finite(x)), length(q) == 1L, q >= 0)
  G <- rbind(c(1, 1), c(0, 1))
  Q <- diag(c(0, q))
  m <- c(level = 0, slope = 0)
  P <- diag(c(4, 0.25))
  loglik <- 0
  for (t in seq_along(x)) {
    if (t > 1L) {
      m[] <- G %*% m
      P <- G %*% P %*% t(G) + Q
    }
    error <- x[t] - m[1]
    variance <- P[1, 1] + 1
    loglik <- loglik + dnorm(error, sd = sqrt(variance), log = TRUE)
    K <- P[, 1] / variance
    m <- m + K * error
    P <- P - tcrossprod(K) * variance
  }
  list(m = m, P = P, q = q, loglik = unname(loglik))
}

# Forecast the actual observation, including its independent error.
trend_forecast <- function(fit, horizons = 1:8) {
  stopifnot(length(horizons) > 0L, all(horizons >= 1),
            all(horizons == floor(horizons)))
  h <- horizons
  center <- unname(fit$m[1] + h * fit$m[2])
  state_var <- fit$P[1, 1] + 2 * h * fit$P[1, 2] + h^2 * fit$P[2, 2] +
    fit$q * (h - 1) * h * (2 * h - 1) / 6
  sd <- sqrt(state_var + 1)
  data.frame(h = h, mean = center, sd = sd,
             lower = center + qnorm(0.025) * sd,
             upper = center + qnorm(0.975) * sd)
}

trend_plot <- function(record, origin, forecasts) {
  days <- origin + forecasts$h
  past <- seq.int(max(1L, origin - 24L), origin)
  old <- par(mar = c(3.6, 4.0, 0.9, 0.7), mgp = c(2.3, 0.65, 0), las = 1)
  on.exit(par(old))
  plot(range(past, days), range(record$x[c(past, days)], forecasts$lower,
                                 forecasts$upper), type = "n", bty = "l",
       xlab = "Day", ylab = "Gauge reading (cm)")
  polygon(c(days, rev(days)), c(forecasts$lower, rev(forecasts$upper)),
          col = "#D8EAF3", border = NA)
  abline(v = origin, lty = 3, col = "#667780")
  points(past, record$x[past], pch = 16, cex = 0.68, col = "#3F4B53")
  lines(days, forecasts$mean, lwd = 2.1, col = "#075882")
  points(days, record$x[days], pch = 1, cex = 0.9, lwd = 1.3, col = "#803A22")
  legend("topleft", c("Forecast mean", "Later reading"), lty = c(1, NA),
         pch = c(NA, 1), col = c("#075882", "#803A22"),
         lwd = c(2.1, 1.3), bty = "n", cex = 0.86)
}

# Compare the last two one-step predictions with the subsequent measurements.
nile_flow_plot <- function(flow, forecast) {
  years <- as.integer(time(flow))
  observed <- as.numeric(flow)
  target <- tail(years, 2)
  last_two <- tail(observed, 2)
  ink <- "#173753"
  forecast_color <- "#A75226"
  old <- par(mar = c(3.8, 4.3, 0.9, 0.8), mgp = c(2.3, 0.65, 0),
             las = 1, bty = "l")
  on.exit(par(old))
  plot(years, observed, type = "n", xaxt = "n", yaxt = "n", xlab = "Year",
       ylab = expression("Annual flow (" * 10^8 ~ m^3 * ")"),
       ylim = extendrange(c(observed, forecast), f = .12))
  ticks <- pretty(range(observed, forecast), n = 4)
  abline(h = ticks, col = "#E4E8EB")
  axis(2, at = ticks)
  axis(1, at = c(years[seq(1L, length(years) - 2L, by = 2L)], target))
  lines(years, observed, col = "#A3B3BB")
  segments(target, last_two, target, forecast, col = forecast_color,
           lty = 3, lwd = 1.7)
  points(years, observed, pch = 16, cex = .8, col = ink)
  points(target, forecast, pch = 2, cex = 1.15, lwd = 1.8, col = forecast_color)
  legend("topleft", legend = c("Recorded flow", "Forecast from preceding year"),
         col = c(ink, forecast_color), pch = c(16, 2), horiz = TRUE,
         bty = "n", cex = .83)
}