# Three settings operate in the cyclic order low, medium, high, low.
wb15_cycle_matrix <- function(stay) {
  matrix(c(stay, 1 - stay, 0,
           0, stay, 1 - stay,
           1 - stay, 0, stay), nrow = 3, byrow = TRUE)
}

wb15_count_record <- function(n = 81L, seed = 1507L) {
  set.seed(seed)
  rates <- c(2, 6, 12)
  transition <- wb15_cycle_matrix(0.8)
  state <- integer(n)
  state[1] <- sample.int(3, 1)
  for (t in seq_len(n - 1L)) {
    state[t + 1L] <- sample.int(3, 1, prob = transition[state[t], ])
  }
  rpois(n, rates[state])
}

# Each row of emission contains p(x_t | S_t = j) over the states j.
wb15_forward <- function(emission, transition, initial) {
  predicted <- initial
  filtered <- matrix(NA_real_, nrow(emission), ncol(emission))
  score <- numeric(nrow(emission))
  for (t in seq_len(nrow(emission))) {
    weights <- predicted * emission[t, ]
    score[t] <- sum(weights)
    if (score[t] <= 0) stop("The observed value has zero predictive probability.")
    filtered[t, ] <- weights / score[t]
    predicted <- drop(filtered[t, ] %*% transition)
  }
  list(filtered = filtered, next_state = predicted, loglik = sum(log(score)))
}

wb15_count_fit <- function(x, stay, rates = c(2, 6, 12)) {
  emission <- outer(x, rates, dpois)
  wb15_forward(emission, wb15_cycle_matrix(stay), rep(1/3, 3))
}

wb15_plot_count_forecast <- function(history, next_weights, rates) {
  old <- par(no.readonly = TRUE)
  on.exit(par(old))
  par(mfrow = c(1, 2), mar = c(3.6, 4.1, 2.1, 0.9),
      mgp = c(2.25, 0.6, 0), las = 1, bty = "l")

  plot(seq_along(history), history, type = "o", pch = 16, cex = 0.43,
       lwd = 1, col = "#31566C", xlab = "Second", ylab = "Number of flashes",
       ylim = c(0, max(history) + 1))
  title("Observed counts", adj = 0, cex.main = 1)

  count <- 0:max(qpois(0.999, rates))
  mass <- drop(outer(count, rates, dpois) %*% next_weights)
  plot(count, mass, type = "h", lwd = 4,
       col = ifelse(count >= 10, "#A34F30", "#31566C"),
       xlab = "Number of flashes", ylab = "Forecast probability",
       ylim = c(0, max(mass) * 1.13))
  title(sprintf("Forecast for second %d", length(history) + 1L),
        adj = 0, cex.main = 1)
}