#' Compute state transition probability matrix for a given lag
#'
#' @param x Numeric vector, the time series.
#' @param n_s Integer, number of states (quantile bins). Default 15.
#' @param lag Integer, lag order for transition. Default 1.
#'
#' @return A matrix of dimension n_s x n_s with row-wise transition probabilities.
#' @export
#'
#' @examples
#' x <- arima.sim(model = list(ar = 0.5), n = 1000)
#' P <- transition_matrix(x, n_s = 15, lag = 1)
transition_matrix <- function(x, n_s = 15, lag = 1) {
  states <- discretize_quantile(x, n_s)
  n <- length(states)
  if (n <= lag) stop("Time series too short for the specified lag")

  from <- states[1:(n - lag)]
  to   <- states[(1 + lag):n]

  counts <- table(from, to)
  P <- matrix(0, n_s, n_s)
  row_idx <- as.integer(rownames(counts))
  col_idx <- as.integer(colnames(counts))
  for (i in seq_along(row_idx)) {
    P[row_idx[i], col_idx] <- counts[i, ]
  }

  row_sums <- rowSums(P)
  P[row_sums > 0, ] <- P[row_sums > 0, ] / row_sums[row_sums > 0]
  return(P)
}
