## FUNCTION: PLOT N-STEP TRANSITION PROBABILITIES
plot_pn <- function(P, n_max, include_n0 = TRUE, lwd = 2) {
  # Plot n -> (P^n)_{ij} in an m x m panel layout (like P itself):
  # row = start i, col = destination j. One global x/y label.
  
  m <- nrow(P)
  n0 <- if (include_n0) 0L else 1L
  n  <- n0:n_max
  
  # Compute powers P^n (store as powers[k, i, j] where k indexes n)
  powers <- array(NA_real_, dim = c(length(n), m, m))
  Pn <- diag(m)  # P^0
  
  k <- 1L
  for (t in 0:n_max) {
    if (t >= n0) {
      powers[k, , ] <- Pn
      k <- k + 1L
    }
    Pn <- Pn %*% P
  }
  
  # Plot: m-by-m panels, color constant within each row i
  op <- par(no.readonly = TRUE)
  on.exit(par(op), add = TRUE)
  
  par(mfrow = c(m, m),
      mar = c(1.0, 1.0, 0.6, 0.4),
      oma = c(3, 3, 0.8, 0.5))
  
  row_col <- seq_len(m)
  
  for (i in seq_len(m)) {
    for (j in seq_len(m)) {
      xaxt <- if (i == m) "s" else "n"
      yaxt <- if (j == 1) "s" else "n"
      
      plot(n, powers[, i, j],
           type = "l", lwd = lwd, col = row_col[i],
           xlab = "", ylab = "",
           xlim = range(n), ylim = c(0, 1),
           xaxt = xaxt, yaxt = yaxt)
    }
  }
  
  mtext("n steps", side = 1, outer = TRUE, line = 1.2)
  mtext(expression((P^n)[ij]), side = 2, outer = TRUE, line = 1.2)
  
  invisible(powers)
}

## EXAMPLE: TWO STATE CHAIN

# function to generate transition matrix
make_Pts <- function(p, q) {
  matrix(c(p, 1 - p,
           1 - q, q),
         nrow = 2, byrow = TRUE)
}

# case 1: stay put
make_Pts(1, 1) |>
  plot_pn(n_max = 10)

# case 2: absorbing state
make_Pts(0.5, 1) |>
  plot_pn(n_max = 10)

# case 3: absorbing state + "sticky" state
make_Pts(0.8, 1) |>
  plot_pn(n_max = 20)

# case 4: "bouncy" state
make_Pts(0.5, 0.01) |>
  plot_pn(n_max = 15)

# case 5: oscillator
make_Pts(0, 0) |>
  plot_pn(n_max = 10)

# EXAMPLE: RANDOM WALKS ON 3D CUBE
states <- c("000","001","010","011","100","101","110","111")
m <- length(states)
P <- matrix(c(
  # to:   000   001   010   011   100   101   110   111
          0,    1/3,  1/3,    0,  1/3,    0,    0,    0,  # from 000
          1/3,    0,    0,  1/3,    0,  1/3,    0,    0,  # from 001
          1/3,    0,    0,  1/3,    0,    0,  1/3,    0,  # from 010
          0,    1/3,  1/3,    0,    0,    0,    0,  1/3,  # from 011
          1/3,    0,    0,    0,    0,  1/3,  1/3,    0,  # from 100
          0,    1/3,    0,    0,  1/3,    0,    0,  1/3,  # from 101
          0,      0,  1/3,    0,  1/3,    0,    0,  1/3,  # from 110
          0,      0,    0,  1/3,    0,  1/3,  1/3,    0   # from 111
), nrow = 8, byrow = TRUE, dimnames = list(states, states))

plot_pn(P, n_max = 20)

# "lazy" random walk: stays put with probability 0.5
P_lazy <- 0.5 * diag(m) + 0.5 * P
plot_pn(P_lazy, n_max = 20)

# even lazier
P_verylazy <- 0.9 * diag(m) + 0.1 * P
plot_pn(P_verylazy, n_max = 20)