# ============================================================================
# Long-run behavior of Markov chains: building the machinery
# ============================================================================

## ---------------------------------------------------------------------------
## 1. GAMBLER'S RUIN
## ---------------------------------------------------------------------------
##
## This is gamblersruin.R from HW1, with k and n renamed. It plays one gambler
## to the end and reports only whether they were ruined.

gamble <- function(i0, k, p) {
  stake <- i0
  
  # play turns until win or bust
  while (stake > 0 && stake < k) {
    stake <- stake + sample(c(-1, 1), size = 1, prob = c(1 - p, p))
  }
  
  # outcome (win or go bust)
  if (stake == 0) 1 else 0
}

i0 <- 2 # start with $2
k  <- 5 # play until $5 or bust
p  <- 0.5 # chance of winning $1 each turn

gamble(i0, k, p) # a single game

set.seed(52710)
mean(replicate(2000, gamble(i0, k, p)))   # estimated ruin probability
(k - i0) / k                              # exact, for a fair game

## --- Drawing one step ----------

## Used by both simulators below: given u ~ Uniform(0,1), find where it lands in
## the cumulative sum of a probability vector.
draw_index <- function(probs, u) sum(u >= cumsum(probs)) + 1L

## --- Same rules, but record the sample path for n turns ----------

gamble_path <- function(i0, k, p, nturn) {
  x <- integer(nturn + 1)         # storage
  x[1] <- i0                      # initial state
  for (t in seq_len(nturn)) {     # for each turn...
    stake <- x[t]                 # origin state
    x[t + 1] <- if (stake == 0 || stake == k) { 
      stake
    } else {
      stake + c(-1L, 1L)[draw_index(c(1 - p, p), runif(1))]
    }                             # destination state
  }
  x                               # return path
}

gamble_path(i0, k, p, nturn = 20) # example usage

## --- Same simulation, but driven by P ----------

gr_matrix <- function(k, p) {
  m <- k + 1
  P <- matrix(0, m, m, dimnames = list(0:k, 0:k))
  P[1, 1] <- 1                     # state 0 is absorbing
  P[m, m] <- 1                     # state k is absorbing
  for (s in 1:(k - 1)) {           # interior states: down 1-p, up p
    P[s + 1, s]     <- 1 - p
    P[s + 1, s + 2] <- p
  }
  P
}

P <- gr_matrix(k, p)                # generate transition matrix from k, p
P                                   # inspect
rowSums(P)                          # every row sums to 1 (P is stochastic)

## --- The general simulator ----------

sim_path <- function(P, path_length, i0) {
  m   <- nrow(P)                                # size of state space
  idx <- integer(path_length + 1)
  idx[1] <- i0 + 1                              # state -> row index
  for (t in seq_len(path_length)) {
    idx[t + 1] <- draw_index(P[idx[t], ], runif(1))
  }
  idx - 1                                       # row index -> state
}

sim_path(P, path_length = 20, i0 = i0)     # example usage

## Check: same rules, same draws, different implementation -- so from the same
## seed the two simulators return the identical path, not just similar summaries.
set.seed(42); a <- gamble_path(i0, k, p, nturn = 30)
set.seed(42); b <- sim_path(P, path_length = 30, i0 = i0)
identical(a, b)
rbind(gamble_path = a, sim_path = b)

## And the ruin probability still lands where it should.
mean(replicate(2000, tail(sim_path(P, path_length = 200, i0 = i0), 1)) == 0)
(k - i0) / k

## Q: What do you expect the long-run behavior to be?

## --- Numerically approximate the limit of P^n ----------

mat_pow <- function(P, n) {
  R <- diag(nrow(P))
  dimnames(R) <- dimnames(P)          # diag() drops the state names
  for (i in seq_len(n)) R <- R %*% P
  R
}

round(mat_pow(P, 5),   3)
round(mat_pow(P, 50),  3)
round(mat_pow(P, 500), 3)

## Q1. The rows settle down, so a limit exists. Does it depend on starting state?
## Q2. How do you interpret this limit? Probability of... what?

## ---------------------------------------------------------------------------
## 2. SIMULATING LONG-RUN BEHAVIOR
## ---------------------------------------------------------------------------
##
## One path tells us very little. We want thousands. Rather than calling
## sim_path in a loop, vectorize.

n     <- 40      # steps per path
n_sim <- 2000    # paths per starting state

sim_many <- function(P, path_length, i0, n_sim) {
  X <- matrix(0L, nrow = n_sim, ncol = path_length + 1)
  X[, 1] <- i0 + 1                              # row indices again
  for (t in seq_len(path_length)) {
    cur <- X[, t]
    u   <- runif(n_sim)
    for (s in unique(cur)) {                    # one draw per occupied state
      rows <- which(cur == s)
      X[rows, t + 1] <- findInterval(u[rows], cumsum(P[s, ])) + 1L
    }
  }
  X - 1L                                        # back to states 0..k
}

X <- sim_many(P, path_length = n, i0 = i0, n_sim = n_sim)
dim(X)                                          # simulations x (steps + 1)
X[1:5, 1:12]                                    # first 11 transitions of first 5 paths


## --- Plot a few hundred of them; not useful as a summary ----------

n_show <- min(300, nrow(X))
steps  <- 0:(ncol(X) - 1)

matplot(steps, t(X[seq_len(n_show), ]), type = "l", lty = 1,
        col = adjustcolor("black", 0.08),
        xlab = "n (steps)", ylab = "state", yaxt = "n",
        main = paste0(n_show, " paths starting from ", i0))
axis(2, at = 0:k, las = 1)

## --- Building better summaries ----------

# occupancy indicator
# I0[r, n, j] = 1 if the r-th path sits in state j after n steps
I0 <- array(0L, dim = c(nrow(X), ncol(X), k + 1),
            dimnames = list(sim  = NULL,
                            step = 0:(ncol(X) - 1),
                            state = 0:k))
for (j in 0:k) I0[, , j + 1] <- (X == j)

dim(I0)            # simulations x turns x states
I0[1:10, 15, ]     # states of first ten paths after 15 steps

# state counts by step across paths
N0 <- colSums(I0, dims = 1)         # step x state
N0[1:6, ]

# estimate the 10-step transition probabilities empirically
round(N0["10", ] / nrow(X), 3)

# compare with analytical solution
round(mat_pow(P, 10)[as.character(i0), ], 3)

## Q. Do you think this depends on starting state?

## --- Simulate from EVERY starting state ----------

# storage
A <- array(0L, dim = c(n_sim, k + 1, n + 1),
           dimnames = list(sim   = NULL,
                           start = 0:k,
                           step  = 0:n))

# simulate n_sim paths of length n from each starting state
set.seed(83126)
for (s in 0:k) {
  A[, s + 1, ] <- sim_many(P, path_length = n, i0 = s, n_sim = n_sim)
}


# A[r, i, t] is the state at step t-1 of the r-th chain started in state i-1.
dim(A)                              # path x starting state x step
A[1:3, "2", 1:10]                   # three paths that began with a stake of 2
A[1:3, , 10]                        # state of three paths from each start at step 10

# Sanity check: share of paths starting from 2 that go bust, should match before
mean(A[, "2", dim(A)[3]] == 0)
(k - i0) / k

# occupancy indicator
# I[r, i, n, j] = 1 if the r-th path, started in i, sits in j after n steps
I <- array(0L, dim = c(n_sim, k + 1, n + 1, k + 1),
           dimnames = list(sim = NULL, start = 0:k, step = 0:n, state = 0:k))
for (j in 0:k) I[, , , j + 1] <- (A == j)

# some example inspections
I[1:3, 2, 5, ]    # state of first three paths starting in state 2 after 5 steps
I[1, , 5, ]       # transitions of one path started in each state after 5 steps
sum(I[, 2, 2, 3]) # how many paths went from 2 to 3 in 2 steps

# transition counts
# N[i, n, j] = how many of the simulated paths went from i to j in n steps
N <- colSums(I, dims = 1)
dim(N)                                   # start x step x state

# Fixing a step gives an estimate of the n-step transition matrix
round(N[, "10", ] / n_sim, 3)
round(mat_pow(P, 10), 3) # theory

round(N[, "40", ] / n_sim, 3)
round(mat_pow(P, 40), 3) # theory

## --- Watching p^(n)_ij settle, or not ----------
##
## One line per starting state: the exact value from mat_pow as a curve, the
## simulated estimate from N as points on top. Fixing a destination j lets all
## k + 1 starting states share one pair of axes.

cols <- hcl.colors(k + 1, "Dark 3")

exact_curves <- function(Pmat, j, n_max) {
  out <- matrix(NA_real_, n_max + 1, nrow(Pmat))
  R <- diag(nrow(Pmat))
  dimnames(R) <- dimnames(Pmat)
  for (t in 0:n_max) {
    out[t + 1, ] <- R[, as.character(j)]
    R <- R %*% Pmat
  }
  out
}

panel_pn <- function(Nc, Pmat, j = 0, main = "", legend = TRUE) {
  n_max <- dim(Nc)[2] - 1
  ex  <- exact_curves(Pmat, j, n_max)
  sim <- sapply(0:k, function(i) Nc[as.character(i), , as.character(j)] / n_sim)

  matplot(0:n_max, ex, type = "l", lty = 1, lwd = 1.4, col = cols,
          ylim = c(0, 1.3), yaxt = "n", xlab = "n (steps)",
          ylab = bquote(p[list(i, .(j))]^(n)), main = main)
  matpoints(0:n_max, sim, pch = 16, cex = 0.55, col = cols)
  axis(2, at = seq(0, 1, 0.2), las = 1)

  if (legend)
    legend("top", paste("start", 0:k), col = cols, lwd = 1.4,
           bty = "n", cex = 0.75, ncol = 3)
}

panel_pn(N, P, j = 0, main = "absorbing: settles, but not together")

## Points a little off the curves at large n are ordinary simulation error, not
## a mismatch: once a path absorbs it stops moving, so N/n_sim stops changing
## and a single unlucky draw shows up as a fixed offset all the way across.

## ---------------------------------------------------------------------------
## 3. REFLECTING BARRIERS
## ---------------------------------------------------------------------------
##
## Suppose you're addicted: if you go bust you borrow $1 to keep playing;
## if you win you buy a $1 snack and then keep playing.

reflect_matrix <- function(k, p) {
  m <- k + 1
  P <- matrix(0, m, m, dimnames = list(0:k, 0:k))
  P[1, 2]     <- 1                 # at 0, bounce up
  P[m, m - 1] <- 1                 # at k, bounce down
  for (s in 1:(k - 1)) {           # interior states unchanged
    P[s + 1, s]     <- 1 - p
    P[s + 1, s + 2] <- p
  }
  P
}

Pr <- reflect_matrix(k, p)
Pr
rowSums(Pr)

# limit via matrix powers
round(mat_pow(Pr, 5),   3) # what explains the zeroes?
round(mat_pow(Pr, 50),  3) # does P^n converge?
round(mat_pow(Pr, 51),  3)

# simulate from every starting state
Ar <- array(0L, dim = c(n_sim, k + 1, n + 1),
            dimnames = list(sim = NULL, start = 0:k, step = 0:n))

set.seed(83126)
for (s in 0:k) {
  Ar[, s + 1, ] <- sim_many(Pr, path_length = n, i0 = s, n_sim = n_sim)
}

# occupancy indicator
Ir <- array(0L, dim = c(n_sim, k + 1, n + 1, k + 1),
            dimnames = list(sim = NULL, start = 0:k, step = 0:n, state = 0:k))
for (j in 0:k) Ir[, , , j + 1] <- (Ar == j)

# transition counts
Nr <- colSums(Ir, dims = 1)

round(Nr[, "39", ] / n_sim, 3)
round(mat_pow(Pr, 39), 3) # theory

round(Nr[, "40", ] / n_sim, 3)
round(mat_pow(Pr, 40), 3) # theory

panel_pn(Nr, Pr, j = 0, main = "reflecting: never settles")

## ---------------------------------------------------------------------------
## 4. HOLDING AT THE BARRIERS
## ---------------------------------------------------------------------------
##
## One more swap: if you win or go bust, you take a break or keep playing
## with equal probability.

hold_matrix <- function(k, p) {
  m <- k + 1
  P <- matrix(0, m, m, dimnames = list(0:k, 0:k))
  P[1, 1] <- 1 - p                 # at 0: lose -> stay
  P[1, 2] <- p                     #       win  -> up
  P[m, m - 1] <- 1 - p             # at k: lose -> down
  P[m, m]     <- p                 #       win  -> stay
  for (s in 1:(k - 1)) {           # interior states unchanged
    P[s + 1, s]     <- 1 - p
    P[s + 1, s + 2] <- p
  }
  P
}

Ph <- hold_matrix(k, p)
Ph
rowSums(Ph)

# limit via matrix powers
round(mat_pow(Ph, 5),   3)
round(mat_pow(Ph, 50),  3)
round(mat_pow(Ph, 500), 3)

# simulate from every starting state
Ah <- array(0L, dim = c(n_sim, k + 1, n + 1),
            dimnames = list(sim = NULL, start = 0:k, step = 0:n))

set.seed(83126)
for (s in 0:k) {
  Ah[, s + 1, ] <- sim_many(Ph, path_length = n, i0 = s, n_sim = n_sim)
}

# occupancy indicator
Ih <- array(0L, dim = c(n_sim, k + 1, n + 1, k + 1),
            dimnames = list(sim = NULL, start = 0:k, step = 0:n, state = 0:k))
for (j in 0:k) Ih[, , , j + 1] <- (Ah == j)

# transition counts
Nh <- colSums(Ih, dims = 1)

round(Nh[, "10", ] / n_sim, 3)
round(mat_pow(Ph, 10), 3) # theory

round(Nh[, "40", ] / n_sim, 3)
round(mat_pow(Ph, 40), 3) # theory

panel_pn(Nh, Ph, j = 0, main = "holding: settles together")

## Q. Now the powers do converge. Does the limit depend on the starting state?


## --- All three together ----------

op <- par(mfrow = c(1, 3), mar = c(4, 4.2, 3, 1))
panel_pn(N,  P,  j = 0, main = "absorbing",  legend = FALSE)
panel_pn(Nr, Pr, j = 0, main = "reflecting", legend = FALSE)
panel_pn(Nh, Ph, j = 0, main = "holding")
par(op)