# ============================================================================
# Hidden Markov models
# ============================================================================
#
# Three examples:
#
#   1. A toy example       We generate the states as well as the observations,
#                          so everything the model estimates can be checked
#                          against the truth. Use it to see what a fitted model
#                          contains and how to get things out of it.
#   2. S&P 500 returns     Choosing an emission distribution by looking at the
#                          data before fitting anything.
#   3. Australian rainfall Your turn.

library(depmixS4)
library(tidyverse)

cols <- c("#D55E00", "grey50", "#0072B2")      # low, middle, high throughout

## ---------------------------------------------------------------------------
## 1. A TOY EXAMPLE
## ---------------------------------------------------------------------------

## --- Simulate ----------

set.seed(1)

# transition matrix
P_true  <- matrix(c(0.33, 0.55, 0.12,          # from the low state
                    0.44, 0.52, 0.04,          # from the middle state
                    0.45, 0.00, 0.55),         # from the high state
                  nrow = 3, byrow = TRUE)

# parameters for emission distribution
mu_true <- c(3.9, 4.3, 4.8)
sd_true <- c(0.19, 0.17, 0.16)

# path length
n       <- 126

# simulate states
state <- integer(n)
state[1] <- 2
for (t in 2:n) {
  state[t] <- sample(1:3, size = 1, prob = P_true[state[t - 1], ])
}
state

# simulate emissions
y <- rnorm(n, mean = mu_true[state], sd = sd_true[state])

# store as dataframe
toy <- data.frame(t = 1:n, yt = y, st = state)

# plot
ggplot(toy, aes(x = t, y = yt)) +
  geom_path(color = 'grey') +
  geom_point(aes(color = factor(st))) +
  scale_color_discrete(palette = cols) +
  theme_minimal()


## --- Fit ----------

# documentation
?depmix
?GLMresponse

# specify model
mod_toy <- depmix(yt ~ 1, data = toy, nstates = 3, family = gaussian())

# fit
fit_toy <- fit(mod_toy, verbose = FALSE)

# fit summary
summary(fit_toy)

# vector of parameter estimates
getpars(fit_toy)

# asymptotic confidence intervals (same order as)
confint(fit_toy)

## --- compare estimates with truth ----------

# check estimated mean for each state
state_means <- function(fit) {
  pars <- getpars(fit)
  unname(pars[names(pars) == "(Intercept)"])
}
state_means(fit_toy)

# order by mean for comparison
rbind(truth = sort(mu_true), 
      fitted = sort(state_means(fit_toy)))


## --- Decoding ----------

# viterbi decoding
vd_toy <- posterior(fit_toy, type = "viterbi")

# relabel states ordered by mean
order_toy <- order(state_means(fit_toy))
path_toy <- match(vd_toy$state, order_toy)

# confusion matrix
table(truth = state, decoded = path_toy)

# overall accuracy
mean(path_toy == state)

# visual: colour is the DECODED state, open circles are the ones it got wrong
toy$dt <- path_toy

ggplot(toy, aes(x = t, y = yt)) +
  geom_path(color = 'grey') +
  geom_point(aes(color = factor(dt), shape = dt == st)) +
  scale_color_discrete(palette = cols) +
  scale_shape_manual(values = c(1, 19)) +
  theme_minimal()


## --- Filtering and smoothing ----------

filt   <- posterior(fit_toy, type = "filtering")[, order_toy]
smooth <- posterior(fit_toy, type = "smoothing")[, order_toy]

plot(1:n, smooth[, 3], type = "l", lwd = 2, ylim = c(0, 1),
     xlab = "t", ylab = "P(high state)",
     main = "filtering vs smoothing")
lines(1:n, filt[, 3], col = "grey55")
points((1:n)[state == 3], rep(1.02, sum(state == 3)), pch = "|", xpd = NA)
text(38, 0.92, "smoothing", adj = 0)
text(38, 0.82, "filtering", col = "grey55", adj = 0)

# underlying recursions are available directly
fb <- forwardbackward(fit_toy)
str(fb)

## ---------------------------------------------------------------------------
## 2. S&P 500 RETURNS
## ---------------------------------------------------------------------------

## --- Data ----------

# monthly log returns on the S&P 500, 1950-2012, shipped with depmixS4
data(sp500)

# store as dataframe
sp <- data.frame(date = as.Date(rownames(sp500)), rt = sp500$logret)

# log returns over time
ggplot(sp, aes(x = date, y = rt)) +
  geom_segment(aes(xend = date, yend = 0), color = 'grey40') +
  theme_minimal()


## --- Choosing an emission distribution ----------

# marginally, against a single normal
ggplot(sp, aes(x = rt)) +
  geom_histogram(aes(y = after_stat(density)), bins = 40,
                 fill = 'grey90', color = 'grey60') +
  stat_function(fun = dnorm, args = list(mean = mean(sp$rt), sd = sd(sp$rt))) +
  theme_minimal()

# kurtosis, against 3 for a normal -- heavy tailed
mean((sp$rt - mean(sp$rt))^4) / sd(sp$rt)^4


## --- Fit ----------

# specify model
mod_sp <- depmix(rt ~ 1, data = sp, nstates = 2, family = gaussian())

# to model volatility regimes, constrain emissions to equal means
# the means are positions 7 and 9 of getpars()
eq <- c(1, 1,  1, 1, 1, 1,  2, 1,  2, 1)

# add equality constraints and use a general optimizer (ignore error)
fit_sp <- fit(mod_sp, equal = eq, method = 'rsolnp',
              verbose = FALSE, solnpcntrl = list(trace = 0))

# fit summary
summary(fit_sp)


## --- Inspect ----------

# one mean, two standard deviations
state_sds <- function(fit) {
  pars <- getpars(fit)
  unname(pars[names(pars) == "sd"])
}
state_sds(fit_sp)

# inferred annual volatility
data.frame(mean = state_means(fit_sp),
           sd   = state_sds(fit_sp)) |>
  mutate(sd_annual_pct = 100 * sd * sqrt(12))

# identify turbulent state
turbulent <- which.max(state_sds(fit_sp))

# persistence: the chance a state stays put. the run of consecutive months in a
# dwell: expected run length
# stationary: persistence if the states were iid
P_sp <- matrix(getpars(fit_sp)[3:6], 2, 2, byrow = TRUE)

data.frame(sd          = state_sds(fit_sp),
           persistence = diag(P_sp),
           dwell       = 1 / (1 - diag(P_sp)),
           stationary  = c(P_sp[2, 1], P_sp[1, 2]) / (P_sp[1, 2] + P_sp[2, 1]))


## --- Decoding ----------

# viterbi decoding
vd_sp <- posterior(fit_sp, type = "viterbi")

# add to dataframe
sp$st <- vd_sp$state

# turbulent stretches, as date ranges, padded so one-month spells stay visible
spells <- sp |>
  mutate(run = cumsum(st != lag(st, default = first(st)))) |>
  filter(st == turbulent) |>
  group_by(run) |>
  summarize(from = min(date) - 20, to = max(date) + 20)

# returns with the turbulent stretches shaded
ggplot(sp, aes(x = date, y = rt)) +
  geom_rect(data = spells, inherit.aes = FALSE,
            aes(xmin = from, xmax = to, ymin = -Inf, ymax = Inf),
            fill = 'mistyrose') +
  geom_segment(aes(xend = date, yend = 0), color = 'grey40') +
  theme_minimal()


## --- Diagnostic: does the emission distribution fit within each state? ----------

# standardize each observation by the mean and sd of its decoded state
sp$z <- (sp$rt - state_means(fit_sp)[sp$st]) / state_sds(fit_sp)[sp$st]

# compare with normal quantiles
ggplot(sp, aes(sample = z)) +
  stat_qq(size = 0.8) +
  stat_qq_line() +
  facet_wrap(~ st) +
  theme_minimal()

# state 2 is still a little heavy tailed...
sp |>
  group_by(st) |>
  summarize(n = n(), kurtosis = mean((z - mean(z))^4) / sd(z)^4)

# ...but it's really just one outlier...
slice_min(sp, abs(z), n = nrow(sp) - 1) |>
  group_by(st) |>
  summarize(n = n(), kurtosis = mean((z - mean(z))^4) / sd(z)^4)

# which was the month with the largest one-day decline in history
slice_max(sp, abs(z))

## ---------------------------------------------------------------------------
## 3. AUSTRALIAN RAINFALL -- YOUR TURN
## ---------------------------------------------------------------------------
##
## Total September-November rainfall (mm) averaged over Australia, 1900-2025
## The oni column measures Pacific Ocean temperature in spring and is how
## climatologists define El Nino (+0.5) and La Nina (-0.5)

url <- "https://tdruiz-stat5710.share.connect.posit.cloud/data/australia-rainfall.csv"
aus <- read.csv(url)

# rainfall over time
ggplot(aus, aes(x = year, y = rain)) +
  geom_segment(aes(xend = year, yend = 0), color = 'grey40') +
  theme_minimal()

# log rainfall
aus$lr <- log(aus$rain)

# 1. Fit a three-state model to lr, with an appropriate emission distribution

# 2. Look at the fitted model.

# 3. Put intervals on the three means. Are the states clearly separated?

# 4. Order the states, decode the series, and visualize.

# 5. Compare states by ONI. Do they seem to track ENSO?

# 6. Based on the model find expected return times for El Niño and La Niña.