# =============================================================================
# Estimation for discrete-time Markov chains: weather / precipitation
# =============================================================================
#
# A toy example, then a real data example, in the order the theory comes up
# in lecture. See also nyctaxi-models.R and geolife-models.R for the other
# two real data examples (sibling scripts, split out of what was originally
# one combined dtmc-estimation.R).
#
#   Toy example.  Simulated weather   3 states, no missing data. Table of
#                                      counts -> transition probabilities,
#                                      checked against the known DGP.
#   Real data.    SLO precipitation   3 states. Conditional MLE by hand and by
#                                      package, diagnostics, time homogeneity,
#                                      fit by season, interval estimates,
#                                      predictions.
#
# FIRST DRAFT: originally an archived script, repaired enough to run top to
# bottom. Comments flag [CUT?] expensive or peripheral parts, [FIX] repairs
# to the originals, and [NEW] pieces that weren't in the originals.

library(tidyverse)
library(markovchain)

## =============================================================================
## TOY EXAMPLE: A SIMULATED THREE-STATE WEATHER CHAIN
## =============================================================================

set.seed(2087)
states_toy <- c("sun", "cloud", "rain")

# the data-generating process 
# transitions between sun and rain mostly pass through cloud
P_true <- matrix(
  c(0.70, 0.25, 0.05,   # from sun
    0.30, 0.40, 0.30,   # from cloud
    0.05, 0.35, 0.60),  # from rain
  nrow = 3, byrow = TRUE,
  dimnames = list(states_toy, states_toy)
)
P_true

# simulate one long realization
n_toy <- 500
sim_toy <- character(n_toy)
sim_toy[1] <- "sun"
for (t in 2:n_toy) {
  sim_toy[t] <- sample(states_toy, size = 1, prob = P_true[sim_toy[t - 1], ])
}

## CONDITIONAL LIKELIHOOD ESTIMATION -------------------------------------------

# table of transition counts
origin_toy <- sim_toy[-n_toy]
destination_toy <- sim_toy[-1]
counts_toy <- table(factor(origin_toy, levels = states_toy),
                     factor(destination_toy, levels = states_toy))
counts_toy

# row normalization to get estimate
P_hat_toy <- prop.table(counts_toy, margin = 1)
round(P_hat_toy, 2)

# compare the estimate to the data-generating process
round(unclass(P_hat_toy) - P_true, 2)

# bonus: same comparison for the stationary distribution
pi_true_toy <- Re(eigen(t(P_true))$vectors[, 1])
pi_true_toy <- setNames(pi_true_toy / sum(pi_true_toy), states_toy)
pi_hat_toy <- Re(eigen(t(P_hat_toy))$vectors[, 1])
pi_hat_toy <- setNames(pi_hat_toy / sum(pi_hat_toy), states_toy)
rbind(true = pi_true_toy, estimated = pi_hat_toy)

## DIAGNOSTIC CHECKS -----------------------------------------------------------

# data adequacy (denominator check)
# need enough transitions observed out of each state to estimate probabilities
rowSums(counts_toy)

# fitted long-run (stationary) shares vs. the empirical time shares
# mismatch would suggest the chain hasn't mixed / 500 steps isn't enough
rbind(
  empirical = table(sim_toy) |> prop.table(),
  fitted    = pi_hat_toy
)

# does simulating from the fitted chain reproduce persistence?
get_streak_lengths <- function(seq, target) {
  r <- rle(seq)                  # run-length encoding
  r$lengths[r$values == target]  # lengths of runs where the value == target
}

set.seed(5309)
sim_toy_check <- character(n_toy)
sim_toy_check[1] <- sim_toy[1]
for (t in 2:n_toy) {
  sim_toy_check[t] <- sample(states_toy, size = 1, prob = P_hat_toy[sim_toy_check[t - 1], ])
}
rbind(
  observed  = c(mean = mean(get_streak_lengths(sim_toy, "rain")),
                median = median(get_streak_lengths(sim_toy, "rain"))),
  simulated = c(mean = mean(get_streak_lengths(sim_toy_check, "rain")),
                median = median(get_streak_lengths(sim_toy_check, "rain")))
)

## A FEW NEAT THINGS YOU CAN DO WITH A FITTED MODEL ----------------------------

# share of time in each state, long run (the stationary distribution)
pi_hat_toy
table(sim_toy) |> prop.table() # note this differs

# starting from a sunny day, how many days until it rains?
hitting_time <- function(P, target) {
  states <- rownames(P)
  other <- setdiff(states, target)
  A <- diag(length(other)) - P[other, other, drop = FALSE]
  h <- solve(A, rep(1, length(other)))
  setNames(h, other)
}
hitting_time(unclass(P_hat_toy), target = "rain")

# once it starts raining, how many consecutive rainy days?
1 / (1 - P_hat_toy["rain", "rain"])

## SAME THINGS, WITH markovchain PACKAGE FUNCTIONS -----------------------------

# data adequacy: createSequenceMatrix() == table of counts, one call
createSequenceMatrix(sim_toy)
rowSums(createSequenceMatrix(sim_toy))

# estimate transition probabilities
fit_toy <- markovchainFit(data = sim_toy, method = "mle")
fit_toy$estimate       # same as P_hat_toy
steadyStates(fit_toy$estimate)  # same as pi_hat_toy

# simulate from the fitted model
set.seed(5309)
sim_toy_check_pkg <- markovchainSequence(
  n = n_toy,
  markovchain = fit_toy$estimate,
  t0 = sim_toy[1],
  include = TRUE
)
rbind(
  observed  = c(mean = mean(get_streak_lengths(sim_toy, "rain")),
                median = median(get_streak_lengths(sim_toy, "rain"))),
  simulated = c(mean = mean(get_streak_lengths(sim_toy_check_pkg, "rain")),
                median = median(get_streak_lengths(sim_toy_check_pkg, "rain")))
)

# hitting time: meanFirstPassageTime() gives every pair of states at once
meanFirstPassageTime(fit_toy$estimate)["sun", "rain"]  # same as hitting_time() above

# streak length has no dedicated package function -- for a discrete-time
# chain it's still just the geometric mean, 1 / (1 - p_ii)
1 / (1 - fit_toy$estimate["rain", "rain"])

## =============================================================================
## DAILY PRECIPITATION AT SAN LUIS OBISPO
## =============================================================================

## DATA ------------------------------------------------------------------------

# noaa daily summaries for a slo station (downloads ~1 MB)
raw <- read_csv('https://www.ncei.noaa.gov/data/global-historical-climatology-network-daily/access/US1CASL0018.csv')

# 12ish years of data
raw$DATE |> range()

# reindex onto the full daily calendar so unreported days show up as NA
full_dates <- tibble(DATE = seq(min(raw$DATE), max(raw$DATE), by = "day"))
n_days_possible <- nrow(full_dates)
n_rows_present <- nrow(raw)
c(days_possible = n_days_possible, rows_present = n_rows_present,
  pct_days_missing = 1 - n_rows_present / n_days_possible)

raw <- full_dates |>
  left_join(raw, by = "DATE") |>
  arrange(DATE)

# visualize -- tons of missing data
set.seed(4821)
raw |>
  drop_na(PRCP) |>
  arrange(DATE) |>
  mutate(
    gap = as.numeric(DATE - lag(DATE)),
    run = cumsum(is.na(gap) | gap > 5),
    PRCP_jit = ifelse(PRCP == 0, runif(n(), 0, 1), PRCP)
  ) |>
  ggplot(aes(x = yday(DATE), y = PRCP)) +
  geom_path(aes(group = interaction(year(DATE), run)), alpha = 0.5) +
  geom_point(aes(y = PRCP_jit), alpha = 0.3, size = 0.7) +
  facet_wrap(~year(DATE)) +
  scale_y_sqrt() +
  theme_minimal()

# how is missingness spread across the calendar year?
raw |>
  mutate(month = month(DATE, label = TRUE), missing = is.na(PRCP)) |>
  group_by(month) |>
  summarise(n = n(), n_missing = sum(missing), pct_missing = mean(missing), .groups = "drop") |>
  ggplot(aes(x = month, y = pct_missing)) +
  geom_col() +
  scale_y_continuous(labels = scales::percent) +
  labs(x = NULL, y = "% days missing") +
  theme_minimal()

# convert to precipitation states by thresholding
prcp <- raw$PRCP

cutoff <- prcp %>%
  .[!is.na(.) & . > 0] %>%
  quantile(probs = 0.75, na.rm = TRUE) %>%
  as.numeric()

prcp_state <- tibble(prcp = prcp) %>%
  mutate(
    state = case_when(
      is.na(prcp)        ~ NA_character_,
      prcp <= 0          ~ "none",
      prcp <= cutoff     ~ "light",
      TRUE               ~ "heavy"
    ),
    state = fct_relevel(state, "none", "light", "heavy")
  ) %>%
  pull(state) |>
  as.character()

# quick check
head(prcp_state, 20)

# share of days in each state
table(prcp_state, useNA = 'always') |> prop.table()

## CONDITIONAL LIKELIHOOD ESTIMATION (PACKAGE) ---------------------------------

# quick and dirty -- ignore NAs and stitch across gaps
# (NOT CORRECT, but works okay if NA rate is low)
fit_approx <- prcp_state[!is.na(prcp_state)] |>
  markovchainFit(method = 'mle')
fit_approx$estimate

# split into contiguous blocks and remove NAs, then fit
# (CORRECT)
x <- prcp_state
is_miss <- is.na(x)
block_id <- cumsum(is_miss & !dplyr::lag(is_miss, default = FALSE))
blocks <- split(x, block_id)
blocks <- lapply(blocks, \(b) b[!is.na(b)])
blocks <- Filter(\(b) length(b) >= 2, blocks)
fit <- markovchainFit(data = blocks, method = 'mle')

# transition probabilities
fit$estimate

# stationary distribution
steadyStates(fit$estimate)
steadyStates(fit_approx$estimate) # for comparison

## DIAGNOSTICS -----------------------------------------------------------------

# data adequacy check
# (looking for sufficient effective sample size)
transition_counts <- fct_na_value_to_level(prcp_state, level = 'unobserved') |>
  createSequenceMatrix()
rowSums(transition_counts) # this is the main check
transition_counts # can also look for row sparsity

# compare stationary distribution to empirical marginals
# (mismatch indicates inhomogeneity or low mixing)
table(prcp_state) |> prop.table()
steadyStates(fit$estimate)

# simulation to check run lengths for dry spells
# (looking to see if the model reproduces time behavior of data)
set.seed(12526)

get_spells <- function(seq, target = "none") {
  r <- rle(seq)                  # run-length encoding
  r$lengths[r$values == target]  # lengths of runs where the value == target
}

# helper function
spell_summary <- function(spells) {
  c(n = length(spells),
    mean = mean(spells),
    median = median(spells),
    max = max(spells))
}

obs_spells <- unlist(lapply(blocks, get_spells))
sim_spells <- unlist(lapply(blocks, function(b) {
  sim_seq <- markovchainSequence(
    n = length(b),
    markovchain = fit$estimate,
    t0 = b[1],
    include = TRUE
  )
  get_spells(sim_seq, target = "none")
}))
rbind(
  observed  = spell_summary(obs_spells),
  simulated = spell_summary(sim_spells)
)

# the model clearly underpredicts persistence. one
# obvious candidate explanation is that we're pooling over seasons

## REFIT BY SEASON -------------------------------------------------------------

# check monthly averages
monthly_precip <- raw |>
  mutate(month = month(DATE, label = TRUE)) |>
  group_by(month) |>
  summarise(mean_prcp = mean(PRCP, na.rm = TRUE), .groups = "drop")

ggplot(monthly_precip, aes(x = month, y = mean_prcp)) +
  geom_col() +
  labs(x = NULL, y = "mean daily PRCP", title = "Average daily precipitation by month") +
  theme_minimal()

# two seasons with cutoff at shoulder (Mar/Apr and Oct/Nov)
d <- as.Date(raw$DATE)
m <- as.integer(format(d, "%m"))
season <- ifelse(m %in% c(11, 12, 1, 2, 3), "wet", "dry")
table(season)

# create blocks within a season by breaking on missing data and nonconsecutive dates
make_blocks <- function(x, d) {
  new_block <- is.na(x) | is.na(dplyr::lag(x)) | (as.numeric(d - dplyr::lag(d)) != 1)
  blks <- split(x, cumsum(new_block)) |> lapply(\(b) b[!is.na(b)])
  Filter(\(b) length(b) >= 2, blks)
}

blocks_by_season <- split(seq_along(prcp_state), season) |>
  lapply(function(idx) {
  make_blocks(x[idx], d[idx])
})

str(blocks_by_season, max.level = 1)

# refit within each season
fit_by_season <- lapply(blocks_by_season, function(blks) {
  markovchainFit(data = blks, method = 'mle')
})

# inspect differences
lapply(fit_by_season, \(f) f$estimate)
lapply(fit_by_season, \(f) steadyStates(f$estimate))

# formal test of whether transition probabilities differ by season
verifyHomogeneity(
  inputList = list(wet = createSequenceMatrix(blocks_by_season$wet),
                    dry = createSequenceMatrix(blocks_by_season$dry)),
  method = "G"
)

## RECHECK DIAGNOSTICS ---------------------------------------------------------

# data adequacy -- note the dry season has almost no "heavy" transitions to
# learn from, so that row of the dry-season fit is essentially unidentified
rowSums(createSequenceMatrix(blocks_by_season$dry))
rowSums(createSequenceMatrix(blocks_by_season$wet))

# stationary distribution vs. empirical marginals
table(unlist(blocks_by_season$dry)) |> prop.table()
steadyStates(fit_by_season$dry$estimate)

table(unlist(blocks_by_season$wet)) |> prop.table()
steadyStates(fit_by_season$wet$estimate)

# dry-spell simulation check, redone within each season: does simulating
# from that season's fit reproduce that season's persistence?
simulate_spells <- function(blks, mc, target = "none") {
  unlist(lapply(blks, function(b) {
    sim_seq <- markovchainSequence(n = length(b), markovchain = mc, t0 = b[1], include = TRUE)
    get_spells(sim_seq, target = target)
  }))
}

set.seed(90142)
obs_spells_dry <- unlist(lapply(blocks_by_season$dry, get_spells))
sim_spells_dry <- simulate_spells(blocks_by_season$dry, fit_by_season$dry$estimate)
rbind(observed = spell_summary(obs_spells_dry), simulated = spell_summary(sim_spells_dry))

obs_spells_wet <- unlist(lapply(blocks_by_season$wet, get_spells))
sim_spells_wet <- simulate_spells(blocks_by_season$wet, fit_by_season$wet$estimate)
rbind(observed = spell_summary(obs_spells_wet), simulated = spell_summary(sim_spells_wet))

# better than pooled, but observed spells still run a bit longer so there's 
# likely some additional persistence (e.g. multi-day storm systems) that the
# model can't capture

## INTERVAL ESTIMATES FOR ONE ROW (WET SEASON) ---------------------------------

transition_counts_wet <- createSequenceMatrix(blocks_by_season$wet)
transition_counts_wet
p_hat_wet <- prop.table(transition_counts_wet, margin = 1)
p_hat_wet

## Row "none" of P, four ways.
i     <- "none"
N_i   <- transition_counts_wet[i, ]            # counts out of state i
n_i   <- sum(N_i)                              # effective sample size for the row
p_i   <- p_hat_wet[i, ]
k_st  <- ncol(p_hat_wet)                       # number of states

# (a) Wald interval from the asymptotic normal approximation
se_i <- sqrt(p_i * (1 - p_i) / n_i)
wald <- cbind(lower = p_i - 1.96 * se_i, upper = p_i + 1.96 * se_i)
wald

# (b) exact binomial interval, treating N_ij ~ binomial(N_i., p_ij)
exact <- sapply(N_i, \(x) binom.test(x, n_i)$conf.int) |> t()
colnames(exact) <- c("lower", "upper")
exact

# (c) posterior credible interval; the marginals of a Dirichlet are beta
alpha_j <- 1                                   # 1 = Laplace
a_post  <- N_i + alpha_j
b_post  <- n_i + k_st * alpha_j - a_post
credible <- cbind(lower = qbeta(0.025, a_post, b_post),
                  upper = qbeta(0.975, a_post, b_post))
credible

# (d) same thing by sampling from the posterior, which is what you need for a
#     joint region
post_draws <- replicate(2000, {
  g <- rgamma(k_st, shape = a_post, rate = 1)  # dirichlet via gamma draws
  g / sum(g)
}) |> t()
colnames(post_draws) <- names(N_i)
apply(post_draws, 2, quantile, probs = c(0.025, 0.975)) |> t()

# (e) package version: multinomialConfidenceIntervals() gives *joint* intervals
mci_wet <- multinomialConfidenceIntervals(
  transitionMatrix = fit_by_season$wet$estimate@transitionMatrix,
  countsTransitionMatrix = transition_counts_wet,
  confidencelevel = 0.95
)
cbind(lower = mci_wet$lowerEndpointMatrix[i, ], 
      upper = mci_wet$upperEndpointMatrix[i, ])

## PREDICTIONS (WET SEASON) ----------------------------------------------------

# starting from "none", how many days until it starts raining?
hitting_time <- function(P, target) {
  states <- rownames(P)
  other <- setdiff(states, target)
  A <- diag(length(other)) - P[other, other, drop = FALSE]
  h <- solve(A, rep(1, length(other)))
  setNames(h, other)
}

P_wet <- unclass(p_hat_wet)
hitting_time(P_wet, target = c("light", "heavy"))

# meanFirstPassageTime() only works for a single-state target (not a set)
meanFirstPassageTime(fit_by_season$wet$estimate)["none", c("light", "heavy")]

# expected number of heavy-precipitation days in the next n days?
# expectedRewards() sums a per-state reward (here, 1 for "heavy", 0 otherwise) 
# over n steps to give sum_{k=1}^n (P^k)[start, target], with less code
mc_wet <- fit_by_season$wet$estimate
reward_heavy <- setNames(as.numeric(states(mc_wet) == "heavy"), states(mc_wet))
setNames(expectedRewards(mc_wet, n = 30, rewards = reward_heavy), states(mc_wet))

# sanity check against the long run
n <- 30
pi_wet <- steadyStates(fit_by_season$wet$estimate)[1, ]
per_start <- setNames(expectedRewards(mc_wet, n = n, rewards = reward_heavy), states(mc_wet))
sum(pi_wet[names(per_start)] * per_start)  # pi-weighted average of per_start
n * pi_wet["heavy"]                        # should match almost exactly