# ============================================================================
# A hidden Markov model for daily rainfall
# ============================================================================
#
# Daily precipitation is modeled directly, in mm, with the amount on each day
# drawn from a distribution that depends on an unobserved weather regime.
#
#   1. A look at the data
#   2. The emission model, and the likelihood by the forward recursion
#   3. Estimation, by numerical optimization
#   4. Filtering
#   5. Smoothing
#   6. Viterbi decoding
#   7. Interpreting the states
#   8. The same machinery, from a package
#
# Runs top to bottom in under a minute.
#
# Station: San Luis Obispo Poly (USC00047851), 2.5 km from campus, daily record
# back to 1893 with essentially no gaps.

## ---------------------------------------------------------------------------
## 0. DATA
## ---------------------------------------------------------------------------

url <- paste0("https://www.ncei.noaa.gov/data/global-historical-climatology-",
              "network-daily/access/USC00047851.csv")
raw <- read.csv(url)

date <- as.Date(raw$DATE)
prcp <- raw$PRCP / 10                          # tenths of a mm -> mm

keep <- date >= as.Date("2000-01-01") & date <= as.Date("2025-12-31")
date <- date[keep]
prcp <- prcp[keep]

# contiguous runs of observed days: no NA, no calendar gaps
starts_block <- c(TRUE, is.na(prcp[-1]) | is.na(prcp[-length(prcp)]) |
                        diff(date) != 1)
blocks <- split(prcp, cumsum(starts_block))
blocks <- lapply(blocks, function(b) b[!is.na(b)])
blocks <- Filter(function(b) length(b) >= 2, blocks)

dates  <- split(date, cumsum(starts_block))    # same split, for later plots
dates  <- dates[names(blocks)]

length(blocks)                                 # how many runs
sum(lengths(blocks))                           # days used
mean(unlist(blocks) > 0)                       # share of wet days

## ---------------------------------------------------------------------------
## 1. A LOOK AT THE DATA
## ---------------------------------------------------------------------------

# (a) the raw series: long flat stretches broken by a handful of bursts, so
#     whatever we fit has to carry structure across days
window <- date >= as.Date("2016-07-01") & date < as.Date("2019-07-01")
plot(date[window], prcp[window], type = "h",
     xlab = "", ylab = "precipitation (mm)", main = "daily rainfall, 2016-2019")

# (b) frequency and intensity rise and fall together: the case for one latent
#     regime driving both, rather than two unrelated seasonal effects
month    <- as.integer(format(date, "%m"))
wet_rate <- tapply(prcp > 0, month, mean, na.rm = TRUE)
wet_mean <- tapply(ifelse(prcp > 0, prcp, NA), month, mean, na.rm = TRUE)
op <- par(mar = c(4, 4.2, 3, 4))               # room for the second axis
plot(1:12, wet_rate, type = "b", pch = 19, ylim = c(0, 0.5),
     xlab = "month", ylab = "share of days with rain",
     main = "frequency and intensity")

# the amounts are in mm, so put them on their own axis on the right
to_left_axis <- max(wet_rate) / max(wet_mean)   # mm -> the left axis's scale
lines(1:12, wet_mean * to_left_axis, type = "b", pch = 1, col = "red")
mm_ticks <- pretty(c(0, wet_mean))
axis(4, at = mm_ticks * to_left_axis, labels = mm_ticks, col = "red",
     col.axis = "red")
mtext("mean mm on wet days", side = 4, line = 2.4, col = "red")
par(op)

# (c) dry spells against independent days: if every day were wet with the same
#     probability, spells would be geometric. The gap is what a persistent
#     hidden state is for.
dry_spells <- unlist(lapply(blocks, function(b) {
  r <- rle(b == 0)
  r$lengths[r$values]
}))
p_wet    <- mean(unlist(blocks) > 0)
spell_k  <- 1:60
observed <- sapply(spell_k, function(k) mean(dry_spells >= k))
if_indep <- (1 - p_wet)^(spell_k - 1)          # geometric survival

plot(spell_k, observed, type = "l", log = "y", ylim = c(1e-3, 1),
     xlab = "spell length k (days)", ylab = "P(dry spell >= k)",
     main = "dry spells")
lines(spell_k, if_indep, col = "red")
text(42, 0.15, "observed", adj = 0)
text(6, 0.004, "if days were independent", col = "red", adj = 0)

# (d) wet-day amounts with one gamma fitted by moments: skewed, so not a normal,
#     and one gamma misses the tail, so not one regime either
wet_amounts <- unlist(blocks)[unlist(blocks) > 0]
mom_shape   <- mean(wet_amounts)^2 / var(wet_amounts)
mom_rate    <- mean(wet_amounts) / var(wet_amounts)

hist(wet_amounts, breaks = 60, freq = FALSE, col = "grey85", border = "grey60",
     xlab = "mm on wet days", main = "wet-day amounts")
curve(dgamma(x, shape = mom_shape, rate = mom_rate), add = TRUE, col = "red")
text(60, 0.03, "marginal gamma fit", col = "red", adj = 0)

# (e) compare distribution of rainfall in a streak vs isolated days

all_days <- unlist(blocks)
is_wet   <- all_days > 0

wet_neighbours <- rep(0, length(all_days))     # wet days among the +/- 3 around
for (lag in c(-3, -2, -1, 1, 2, 3)) {
  shifted <- c(rep(NA, max(0, -lag)),
               is_wet[seq_len(length(is_wet) - abs(lag)) + max(0, lag)],
               rep(NA, max(0, lag)))
  wet_neighbours <- wet_neighbours + ifelse(is.na(shifted), 0, shifted)
}

isolated <- all_days[is_wet & wet_neighbours <= 1]
in_spell <- all_days[is_wet & wet_neighbours >= 3]
c(isolated = length(isolated), in_spell = length(in_spell))
c(mean_isolated = mean(isolated), mean_in_spell = mean(in_spell))

# P(amount > x)
grid      <- 10^seq(log10(0.2), log10(120), length.out = 120)
above_iso <- sapply(grid, function(z) mean(isolated > z))
above_spl <- sapply(grid, function(z) mean(in_spell > z))

plot(grid, above_iso, type = "l", log = "x", ylim = c(0, 1),
     xlab = "x (mm, log scale)", ylab = "P(amount > x)",
     main = "wet-day amounts, by what surrounds them")
lines(grid, above_spl, col = "red")
text(5, 0.72, "spell", col = "red", adj = 0)
text(1.1, 0.28, "isolated", adj = 0)

## ---------------------------------------------------------------------------
## 2. THE EMISSION MODEL, AND THE LIKELIHOOD
## ---------------------------------------------------------------------------
##
## Each regime j gets its own rainfall distribution: dry with probability
## 1 - w_j, and otherwise a gamma amount with its own mean and shape.
##
##     f_j(0) = 1 - w_j
##     f_j(x) = w_j * P(a gamma amount rounds to x)                 for x > 0

emission_matrix <- function(x, w, shape, mean_mm, half_width = 0.05) {
  m <- length(w)
  out <- matrix(0, m, length(x))
  for (j in 1:m) {
    rate  <- shape[j] / mean_mm[j]
    upper <- pgamma(x + half_width, shape = shape[j], rate = rate)
    lower <- pgamma(pmax(x - half_width, 0), shape = shape[j], rate = rate)
    out[j, ] <- ifelse(x == 0, 1 - w[j], w[j] * (upper - lower))
  }
  out                                          # m x length(x)
}

# parameters to try out: a storm regime and a dry regime
P_try     <- matrix(c(0.80, 0.20,
                      0.05, 0.95), nrow = 2, byrow = TRUE)
w_try     <- c(0.70, 0.03)                     # P(rain) in each regime
shape_try <- c(0.8, 0.8)
mean_try  <- c(12, 2)                          # mm on a wet day
delta_try <- c(0.25, 0.75)

x <- blocks[[which.max(lengths(blocks))]]      # the longest run of days
F_mat <- emission_matrix(x, w_try, shape_try, mean_try)

# the forward recursion, one step at a time, rescaled so nothing underflows
alpha   <- delta_try * F_mat[, 1]
scale_1 <- sum(alpha)
alpha   <- alpha / scale_1
log_lik <- log(scale_1)

for (t in 2:length(x)) {
  alpha   <- (alpha %*% P_try) * F_mat[, t]    # advance a step, then reweight
  scale_t <- sum(alpha)
  alpha   <- alpha / scale_t
  log_lik <- log_lik + log(scale_t)
}
log_lik

## --- Wrap it up ----------

forward_loglik <- function(x, P, delta, w, shape, mean_mm) {
  F_mat   <- emission_matrix(x, w, shape, mean_mm)
  alpha   <- delta * F_mat[, 1]
  log_lik <- log(sum(alpha))
  alpha   <- alpha / sum(alpha)
  for (t in 2:length(x)) {
    alpha   <- (alpha %*% P) * F_mat[, t]
    log_lik <- log_lik + log(sum(alpha))
    alpha   <- alpha / sum(alpha)
  }
  log_lik
}

# independent runs, so log-likelihoods add
sum(sapply(blocks, forward_loglik, P = P_try, delta = delta_try,
           w = w_try, shape = shape_try, mean_mm = mean_try))

## ---------------------------------------------------------------------------
## 3. ESTIMATION, BY NUMERICAL OPTIMIZATION
## ---------------------------------------------------------------------------
##
## The likelihood is computable, so hand it to an optimizer. Everything is
## constrained -- rows of P on the simplex, w in (0, 1), shape and mean
## positive -- so optimize unconstrained working parameters and transform back.

softmax <- function(z) exp(z) / sum(exp(z))

unpack <- function(par, m) {
  n_trans <- m * (m - 1)

  logit_P <- matrix(0, m, m)                   # diagonal logit fixed at 0
  if (n_trans > 0) logit_P[row(logit_P) != col(logit_P)] <- par[1:n_trans]
  P <- t(apply(logit_P, 1, softmax))

  rest    <- par[n_trans + seq_len(3 * m)]
  w       <- plogis(rest[1:m])                 # rain probability
  shape   <- exp(rest[m + (1:m)])
  mean_mm <- exp(rest[2 * m + (1:m)])

  # stationary start; NA if the optimizer wanders into a reducible P
  delta <- tryCatch(solve(t(diag(m) - P + 1), rep(1, m)),
                    error = function(e) rep(NA_real_, m))
  list(P = P, delta = delta, w = w, shape = shape, mean_mm = mean_mm)
}

negloglik <- function(par, m) {
  th <- unpack(par, m)
  if (anyNA(th$delta) || any(th$delta < 0)) return(1e10)   # reducible chain
  v  <- -sum(sapply(blocks, forward_loglik, P = th$P, delta = th$delta,
                    w = th$w, shape = th$shape, mean_mm = th$mean_mm))
  if (!is.finite(v)) 1e10 else v
}

fit_hmm <- function(m, n_starts = 2, seed = 1) {
  set.seed(seed)
  n_par <- m * (m - 1) + 3 * m
  start_1 <- c(rep(-2, m * (m - 1)),                     # persistent regimes
               qlogis(seq(0.5, 0.1, length.out = m)),    # wetter to drier
               log(rep(0.8, m)),
               log(seq(10, 2, length.out = m)))
  runs <- lapply(1:n_starts, function(s) {
    start <- if (s == 1) start_1 else start_1 + rnorm(n_par, sd = 0.5)
    optim(start, negloglik, m = m, method = "BFGS", control = list(maxit = 300))
  })
  best <- runs[[which.min(sapply(runs, `[[`, "value"))]]
  c(unpack(best$par, m), list(loglik = -best$value, k = n_par,
                              convergence = best$convergence))
}

fit2 <- fit_hmm(m = 2)

round(fit2$P, 3)                               # regime persistence
round(data.frame(rain_prob = fit2$w, shape = fit2$shape,
                 mean_mm = fit2$mean_mm, stationary = fit2$delta), 3)
fit2$loglik

# The fitted rainfall distribution in each state.
storm_state <- which.max(fit2$w)               # whichever state is the wet one
dry_state   <- which.min(fit2$w)

grid_mm    <- 10^seq(log10(0.05), log10(150), length.out = 200)
tail_state <- function(j) {
  fit2$w[j] * (1 - pgamma(grid_mm, shape = fit2$shape[j],
                          rate  = fit2$shape[j] / fit2$mean_mm[j]))
}
tail_storm <- tail_state(storm_state)
tail_dry   <- tail_state(dry_state)
tail_mix   <- fit2$delta[storm_state] * tail_storm +
              fit2$delta[dry_state]   * tail_dry
tail_obs   <- sapply(grid_mm, function(z) mean(all_days > z))
tail_obs[tail_obs == 0] <- NA                  # log scale drops the empty tail

plot(grid_mm, tail_storm, type = "l", col = "red", log = "xy", ylim = c(1e-4, 1),
     xlab = "x (mm, log scale)", ylab = "P(amount > x | state)",
     main = "fitted rainfall, by state")
lines(grid_mm, tail_dry)
text(0.06, 0.30, "storm", col = "red", adj = 0)
text(0.06, 5e-3, "dry", adj = 0)

plot(grid_mm, tail_obs, type = "l", log = "xy", ylim = c(1e-4, 1),
     xlab = "x (mm, log scale)", ylab = "P(amount > x)",
     main = "observed vs fitted")
lines(grid_mm, tail_mix, col = "red")

## ---------------------------------------------------------------------------
## 4. FILTERING
## ---------------------------------------------------------------------------
##
## P(Y_t = j | x_1 ... x_t): what you would believe at the time. Predict with P,
## update with the emissions, normalize.

filter_states <- function(x, P, delta, w, shape, mean_mm) {
  F_mat <- emission_matrix(x, w, shape, mean_mm)
  n <- length(x); m <- nrow(P)
  pi_filt <- matrix(0, n, m)
  c_t     <- numeric(n)

  predicted <- delta
  for (t in 1:n) {
    updated      <- predicted * F_mat[, t]
    c_t[t]       <- sum(updated)               # p(x_t | x_1 ... x_{t-1})
    pi_filt[t, ] <- updated / c_t[t]
    predicted    <- as.vector(pi_filt[t, ] %*% P)
  }
  list(pi = pi_filt, c = c_t, loglik = sum(log(c_t)))
}

filt <- filter_states(x, fit2$P, fit2$delta, fit2$w, fit2$shape, fit2$mean_mm)

# the filter and the likelihood are one computation, read two ways
c(filter  = filt$loglik,
  forward = forward_loglik(x, fit2$P, fit2$delta, fit2$w, fit2$shape, fit2$mean_mm))

## ---------------------------------------------------------------------------
## 5. SMOOTHING
## ---------------------------------------------------------------------------
##
## P(Y_t = j | x_1 ... x_T): the whole record in hand. The backward pass gives
## beta_t(j) = P(x_{t+1} ... x_T | Y_t = j), and gamma is alpha times beta.

smooth_states <- function(x, P, delta, w, shape, mean_mm) {
  F_mat <- emission_matrix(x, w, shape, mean_mm)
  n <- length(x); m <- nrow(P)

  alpha <- matrix(0, n, m); scale_t <- numeric(n)
  alpha[1, ] <- delta * F_mat[, 1]
  scale_t[1] <- sum(alpha[1, ]); alpha[1, ] <- alpha[1, ] / scale_t[1]
  for (t in 2:n) {
    alpha[t, ] <- (alpha[t - 1, ] %*% P) * F_mat[, t]
    scale_t[t] <- sum(alpha[t, ]); alpha[t, ] <- alpha[t, ] / scale_t[t]
  }

  beta <- matrix(0, n, m)
  beta[n, ] <- 1
  for (t in (n - 1):1) {
    beta[t, ] <- (P %*% (F_mat[, t + 1] * beta[t + 1, ])) / scale_t[t + 1]
  }

  gamma <- alpha * beta
  gamma / rowSums(gamma)
}

gamma_x <- smooth_states(x, fit2$P, fit2$delta, fit2$w, fit2$shape, fit2$mean_mm)

## ---------------------------------------------------------------------------
## 6. VITERBI DECODING
## ---------------------------------------------------------------------------
##
## The most likely path, rather than the most likely state day by day

viterbi <- function(x, P, delta, w, shape, mean_mm) {
  F_mat <- emission_matrix(x, w, shape, mean_mm)
  n <- length(x); m <- nrow(P)
  log_P <- log(P); log_F <- log(F_mat)

  best <- matrix(-Inf, n, m)                   # best[t, j] = log d_t(j)
  from <- matrix(0L, n, m)                     # from[t, j] = psi_t(j)
  best[1, ] <- log(delta) + log_F[, 1]

  for (t in 2:n) {
    for (j in 1:m) {
      candidates <- best[t - 1, ] + log_P[, j]
      from[t, j] <- which.max(candidates)
      best[t, j] <- max(candidates) + log_F[j, t]
    }
  }

  path <- integer(n)
  path[n] <- which.max(best[n, ])
  for (t in (n - 1):1) path[t] <- from[t + 1, path[t + 1]]
  path
}

path_x <- viterbi(x, fit2$P, fit2$delta, fit2$w, fit2$shape, fit2$mean_mm)

# pointwise decoding is just the argmax of the smoothed probabilities
pointwise_x <- max.col(gamma_x)
table(viterbi = path_x, pointwise = pointwise_x)

## ---------------------------------------------------------------------------
## 7. INTERPRETING THE STATES
## ---------------------------------------------------------------------------
##
## States are labeled arbitrarily, so order them by something meaningful, then
## describe each one by what it emits, how long it lasts, and when it occurs.

# order the states by how wet they are, and name them
wet_first <- order(fit2$w, decreasing = TRUE)
regime    <- c("storm", "dry")

data.frame(
  regime        = regime,
  rain_prob     = round(fit2$w[wet_first], 3),
  mean_mm_wet   = round(fit2$mean_mm[wet_first], 2),
  dwell_days    = round(1 / (1 - diag(fit2$P)[wet_first]), 1),
  share_of_days = round(fit2$delta[wet_first], 3)
)

# smooth and decode the whole record, not just the longest run
gammas <- lapply(blocks, smooth_states, P = fit2$P, delta = fit2$delta,
                 w = fit2$w, shape = fit2$shape, mean_mm = fit2$mean_mm)
paths  <- lapply(blocks, viterbi, P = fit2$P, delta = fit2$delta,
                 w = fit2$w, shape = fit2$shape, mean_mm = fit2$mean_mm)

storm_prob   <- unlist(lapply(gammas, function(g) g[, wet_first[1]]))
decoded      <- unlist(paths) == storm_state
decoded_date <- as.Date(unlist(dates, use.names = FALSE), origin = "1970-01-01")
decoded_month <- as.integer(format(decoded_date, "%m"))

# how long do the decoded episodes last, in days?
spell_lengths <- rle(decoded)
round(rbind(
  storm = summary(spell_lengths$lengths[spell_lengths$values]),
  dry   = summary(spell_lengths$lengths[!spell_lengths$values])
), 1)

# Does the model reproduce the seasonality? Nothing seasonal was fitted
storm_by_month <- tapply(storm_prob, decoded_month, mean)
implied_wet    <- storm_by_month * fit2$w[wet_first[1]] +
                  (1 - storm_by_month) * fit2$w[wet_first[2]]

plot(1:12, wet_rate, type = "b", pch = 19, ylim = c(0, 0.5), xlab = "month",
     ylab = "share of days with rain", main = "observed vs model-implied")
lines(1:12, implied_wet, type = "b", pch = 1, col = "red")
text(6.5, 0.46, "observed")
text(6.5, 0.38, "implied by the fitted model", col = "red")