MDPs & the Bellman Equations

"""MDPs and the Bellman equations, transcribed literally.

No library. A tiny MDP held as nested lists, so every claim in the
theory becomes something that can be computed rather than quoted:

    V*(s)  = max_a [ R(s,a) + gamma * sum_s' P(s'|s,a) V*(s') ]
    V^pi(s) =  sum_a pi(a|s) [ R(s,a) + gamma * sum_s' P(s'|s,a) V^pi(s') ]

Four things to read for.

1. The two equations are structurally different. The second is LINEAR
   in V, and solve_policy_value_exactly() inverts it by Gaussian
   elimination with no iteration at all. The first has a max in it and
   is nonlinear, so value_iteration() can only approach it. That single
   difference is why every method in this category alternates between a
   cheap evaluation step and an expensive improvement step.

2. contraction_check() measures the inequality the whole formalism
   rests on: one backup shrinks the sup-norm gap between any two value
   functions by at least gamma.

3. value_iteration() returns a real error BOUND, not a sense that it
   converged. If two successive iterates differ by eps then the current
   estimate is within gamma*eps/(1-gamma) of the truth.

4. markov_violation_demo() aggregates two states that behave
   differently and solves again. Nothing errors. The iteration
   converges cleanly to the fixed point of a process nobody wanted, and
   the derived policy is wrong -- which is the most dangerous failure
   here precisely because the numbers look healthy.
"""

Kernel = list[list[list[float]]]   # P[state][action][next_state]
Reward = list[list[float]]         # R[state][action]
Policy = list[list[float]]         # pi[state][action]


def sup_norm(left: list[float], right: list[float]) -> float:
    """Largest absolute difference over states.

    The norm the contraction holds in, and not an arbitrary choice: the
    guarantees below are all sup-norm statements, and they say nothing
    about a least-squares projection -- which is exactly the gap where
    deep reinforcement learning loses its stability argument.
    """
    return max(abs(a - b) for a, b in zip(left, right))


def expected_next_value(
    kernel: Kernel, values: list[float], state: int, action: int
) -> float:
    total = 0.0
    for next_state, probability in enumerate(kernel[state][action]):
        total += probability * values[next_state]
    return total


def bellman_optimality_backup(
    kernel: Kernel, reward: Reward, values: list[float], gamma: float
) -> list[float]:
    """One application of T: the max makes this nonlinear.

    Every state updated from the same input values, which is what makes
    this an operator rather than a sweep -- and what makes the
    contraction argument apply cleanly.
    """
    updated = []
    for state, action_rewards in enumerate(reward):
        best = max(
            action_rewards[action] + gamma * expected_next_value(kernel, values, state, action)
            for action in range(len(action_rewards))
        )
        updated.append(best)
    return updated


def bellman_expectation_backup(
    kernel: Kernel, reward: Reward, policy: Policy, values: list[float], gamma: float
) -> list[float]:
    """One application of T^pi: no max, therefore linear in V.

    Read against the function above. The only difference is that the
    max became an average under a fixed policy, and that is the entire
    reason this equation has a closed-form solution and the other does
    not.
    """
    updated = []
    for state, action_probabilities in enumerate(policy):
        total = 0.0
        for action, probability in enumerate(action_probabilities):
            if probability == 0.0:
                continue
            total += probability * (
                reward[state][action]
                + gamma * expected_next_value(kernel, values, state, action)
            )
        updated.append(total)
    return updated


def solve_policy_value_exactly(
    kernel: Kernel, reward: Reward, policy: Policy, gamma: float
) -> list[float]:
    """V^pi = (I - gamma P_pi)^-1 R_pi, by Gaussian elimination.

    Written out because the fact that this is possible at all is the
    structural point. Policy evaluation is a linear system with as many
    equations as states; no iteration, no tolerance, no convergence
    criterion. The optimality equation admits nothing of the kind at
    any size.
    """
    n_states = len(reward)

    # Build (I - gamma P_pi) augmented with R_pi.
    matrix = []
    for state in range(n_states):
        row = [0.0] * (n_states + 1)
        row[state] = 1.0

        for action, probability in enumerate(policy[state]):
            if probability == 0.0:
                continue
            row[n_states] += probability * reward[state][action]
            for next_state, transition in enumerate(kernel[state][action]):
                row[next_state] -= gamma * probability * transition
        matrix.append(row)

    # Forward elimination with partial pivoting.
    for column in range(n_states):
        pivot_row = max(range(column, n_states), key=lambda r: abs(matrix[r][column]))
        if abs(matrix[pivot_row][column]) < 1e-12:
            raise ValueError(
                f"(I - gamma P_pi) is singular at column {column}; this happens at "
                "gamma = 1 with a policy that never terminates, where the equation "
                "has no solution rather than a hard one"
            )
        matrix[column], matrix[pivot_row] = matrix[pivot_row], matrix[column]

        for row in range(column + 1, n_states):
            factor = matrix[row][column] / matrix[column][column]
            if factor == 0.0:
                continue
            for position in range(column, n_states + 1):
                matrix[row][position] -= factor * matrix[column][position]

    # Back substitution.
    values = [0.0] * n_states
    for state in reversed(range(n_states)):
        total = matrix[state][n_states]
        for column in range(state + 1, n_states):
            total -= matrix[state][column] * values[column]
        values[state] = total / matrix[state][state]
    return values


def value_iteration(
    kernel: Kernel, reward: Reward, gamma: float, tolerance: float = 1e-8, max_sweeps: int = 10_000
):
    """Iterate T to its fixed point, with an actual error bound.

    The contraction gives more than "it converged": if successive
    iterates differ by eps in sup-norm, the current estimate is within
    gamma*eps/(1-gamma) of V*. That bound is computable, free, and
    almost never reported.

    Note the sweep count against gamma. At 0.9 each backup removes a
    tenth of the error; at 0.999 it removes a thousandth. The hundredfold
    difference in sweeps is the problem being harder, not the solver
    being slow.
    """
    if not 0.0 <= gamma < 1.0:
        raise ValueError(
            f"gamma must lie in [0, 1); got {gamma}. At exactly 1 the backup is not a "
            "contraction and none of the guarantees below hold"
        )

    values = [0.0] * len(reward)
    for sweep in range(max_sweeps):
        updated = bellman_optimality_backup(kernel, reward, values, gamma)
        gap = sup_norm(updated, values)
        values = updated

        if gap < tolerance:
            return {
                "values": values,
                "sweeps": sweep + 1,
                "final_gap": gap,
                # The bound, from the contraction. Not an estimate.
                "error_bound": gamma * gap / (1.0 - gamma),
            }

    return {
        "values": values,
        "sweeps": max_sweeps,
        "final_gap": gap,
        "error_bound": gamma * gap / (1.0 - gamma),
    }


def greedy_policy(kernel: Kernel, reward: Reward, values: list[float], gamma: float) -> list[int]:
    """Act greedily on a value function.

    Tie-breaking is a policy decision, not an implementation detail:
    exact ties are common in structured problems, and min() here picks
    the lowest index deterministically so a re-solve does not silently
    produce a different policy with the same value.
    """
    actions = []
    for state, action_rewards in enumerate(reward):
        scored = [
            (
                action_rewards[action] + gamma * expected_next_value(kernel, values, state, action),
                action,
            )
            for action in range(len(action_rewards))
        ]
        best_value = max(score for score, _ in scored)
        actions.append(min(action for score, action in scored if score == best_value))
    return actions


def contraction_check(
    kernel: Kernel, reward: Reward, gamma: float, first: list[float], second: list[float]
) -> dict[str, float]:
    """The inequality the entire formalism rests on, measured.

        ||T V - T U||_inf  <=  gamma * ||V - U||_inf

    Everything else -- that a solution exists, that it is unique, that
    iteration reaches it, at what rate -- is Banach's fixed-point
    theorem applied to this one line.
    """
    before = sup_norm(first, second)
    after = sup_norm(
        bellman_optimality_backup(kernel, reward, first, gamma),
        bellman_optimality_backup(kernel, reward, second, gamma),
    )
    return {
        "gap_before": before,
        "gap_after": after,
        "observed_modulus": after / before if before > 0.0 else 0.0,
        "gamma": gamma,
        "holds": after <= gamma * before + 1e-12,
    }


def effective_horizon(gamma: float) -> float:
    """Roughly how many steps the discount actually looks ahead.

    1/(1-gamma). Worth computing rather than intuiting: gamma of 0.99 is
    a hundred-step horizon, which is a different problem from the
    ten-step one at 0.9 rather than a more patient version of it.
    """
    return 1.0 / (1.0 - gamma)


def markov_violation_demo(gamma: float = 0.9) -> dict[str, object]:
    """The failure with no diagnostic.

    Two states that behave differently are merged into one, as any
    coarse state representation does. The aggregated MDP is perfectly
    well-formed: the kernel normalizes, the backup contracts, value
    iteration converges, and the bound is satisfied.

    It converges to the value function of the aggregated process, which
    is not the value function of the original. The derived policy is
    wrong and nothing anywhere says so -- which is why the Markov claim
    has to be argued substantively rather than checked numerically.
    """
    # Three states. From s0, action 0 leads to s1 (good), action 1 to s2
    # (bad). s1 and s2 are absorbing with very different rewards.
    kernel: Kernel = [
        [[0.0, 1.0, 0.0], [0.0, 0.0, 1.0]],
        [[0.0, 1.0, 0.0], [0.0, 1.0, 0.0]],
        [[0.0, 0.0, 1.0], [0.0, 0.0, 1.0]],
    ]
    reward: Reward = [[0.0, 0.0], [1.0, 1.0], [-1.0, -1.0]]

    true_solution = value_iteration(kernel, reward, gamma)
    true_policy = greedy_policy(kernel, reward, true_solution["values"], gamma)

    # Now aggregate s1 and s2 into a single state, averaging what an
    # observer who cannot tell them apart would see.
    aggregated_kernel: Kernel = [
        [[0.0, 1.0], [0.0, 1.0]],
        [[0.0, 1.0], [0.0, 1.0]],
    ]
    aggregated_reward: Reward = [[0.0, 0.0], [0.0, 0.0]]

    aggregated_solution = value_iteration(aggregated_kernel, aggregated_reward, gamma)
    aggregated_policy = greedy_policy(
        aggregated_kernel, aggregated_reward, aggregated_solution["values"], gamma
    )

    return {
        "markov_values": true_solution["values"],
        "markov_policy": true_policy,
        "aggregated_values": aggregated_solution["values"],
        "aggregated_policy": aggregated_policy,
        "note": (
            "both runs converged inside their error bounds; the aggregated one "
            "solved a different problem and reported no difficulty doing so"
        ),
    }

"""The formalism with its preconditions enforced.

The literal version accepted all of these without complaint:

  * a transition row summing to 0.97, which converges cleanly to the
    value of a process that leaks probability;
  * a discount of exactly 1, where the backup is not a contraction and
    every guarantee silently stops applying;
  * a state with no available action, where the max is over an empty
    set;
  * a kernel, reward and policy whose shapes disagree, which indexes
    into whatever happens to be there.

Each is checked once, at construction. That placement is the point: an
MDP is a contract, and a contract validated per call is one that gets
validated in three places and skipped in the fourth.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import NamedTuple

import numpy as np
import numpy.typing as npt

FloatArray = npt.NDArray[np.float64]
IntArray = npt.NDArray[np.int64]

# Rows may drift by this much from summing to one before it is an error
# rather than floating-point noise.
KERNEL_TOLERANCE = 1e-9


class MdpError(Exception):
    """Base for every violated precondition of the formalism."""


class MalformedKernel(MdpError):
    """A transition row does not describe a distribution.

    Worth its own type because the failure is silent: probability that
    leaks out of a row is never redistributed and never reported, and
    the fixed point that results is the value function of a process
    that quietly loses mass.
    """


class InvalidDiscount(MdpError):
    """Gamma outside [0, 1), where the contraction does not hold."""


class NoAvailableAction(MdpError):
    """A state with an empty action set; the max has nothing to range over."""


class SingularEvaluation(MdpError):
    """(I - gamma P_pi) is not invertible.

    This is not a numerical difficulty to work around. At gamma of one
    with a policy that never terminates, the expectation equation has no
    solution, and the right response is to fix the model.
    """


@dataclass(frozen=True)
class Discount:
    """The discount, with the horizon it actually specifies.

    A value object rather than a float because it is routinely treated
    as a tuning knob, and it is not one: changing it changes the problem
    being solved and the cost of solving it at the same time.
    """

    gamma: float

    def __post_init__(self) -> None:
        if not 0.0 <= self.gamma < 1.0:
            raise InvalidDiscount(
                f"gamma must lie in [0, 1); got {self.gamma}. At exactly 1 the Bellman "
                "backup is not a contraction, so existence, uniqueness and the "
                "convergence rate all stop applying"
            )

    @property
    def effective_horizon(self) -> float:
        """1/(1-gamma) — roughly how many steps this looks ahead."""
        return 1.0 / (1.0 - self.gamma)

    def sweeps_for_tolerance(self, tolerance: float, initial_gap: float = 1.0) -> int:
        """How many backups a given accuracy needs, from the contraction.

        Computable in advance, which makes the cost of a long horizon a
        planning input rather than a surprise: at gamma 0.999 this is
        roughly a thousand times the count at 0.9.
        """
        if tolerance <= 0.0:
            raise ValueError(f"tolerance must be positive, got {tolerance}")
        if self.gamma == 0.0:
            return 1
        return int(np.ceil(np.log(tolerance / initial_gap) / np.log(self.gamma)))


@dataclass(frozen=True)
class MarkovDecisionProcess:
    """States, actions, dynamics, reward — validated once.

    Frozen because an MDP is the problem statement. Mutating the kernel
    after a solve leaves a value function that answers a question no
    longer being asked, and nothing downstream would notice.
    """

    kernel: FloatArray      # (states, actions, states)
    reward: FloatArray      # (states, actions)
    available: npt.NDArray[np.bool_]  # (states, actions)

    def __post_init__(self) -> None:
        if self.kernel.ndim != 3:
            raise MalformedKernel(f"kernel must be (S, A, S); got shape {self.kernel.shape}")

        n_states, n_actions, n_next = self.kernel.shape
        if n_states != n_next:
            raise MalformedKernel(
                f"kernel maps {n_states} states to {n_next}; a transition kernel is square"
            )
        if self.reward.shape != (n_states, n_actions):
            raise MdpError(
                f"reward shape {self.reward.shape} does not match kernel ({n_states}, {n_actions})"
            )
        if self.available.shape != (n_states, n_actions):
            raise MdpError(
                f"availability shape {self.available.shape} does not match the reward"
            )

        if not np.all(np.isfinite(self.reward[self.available])):
            raise MdpError("reward contains a non-finite value on an available action")
        if np.any(self.kernel < 0.0):
            raise MalformedKernel("kernel contains a negative probability")

        # The check that matters most, and the one nothing else catches.
        row_sums = self.kernel.sum(axis=2)
        offenders = np.argwhere(np.abs(row_sums[self.available] - 1.0) > KERNEL_TOLERANCE)
        if offenders.size > 0:
            worst = float(np.max(np.abs(row_sums[self.available] - 1.0)))
            raise MalformedKernel(
                f"a transition row departs from 1 by {worst:.3e}; probability that leaks "
                "is never redistributed, and the fixed point becomes the value of a "
                "process that loses mass"
            )

        empty = np.flatnonzero(~self.available.any(axis=1))
        if empty.size > 0:
            raise NoAvailableAction(
                f"states {empty.tolist()} have no available action; the max in the "
                "optimality equation ranges over an empty set"
            )

    @property
    def n_states(self) -> int:
        return int(self.kernel.shape[0])

    @property
    def n_actions(self) -> int:
        return int(self.kernel.shape[1])

    def action_values(self, values: FloatArray, discount: Discount) -> FloatArray:
        """Q(s, a) for every pair, with unavailable actions masked out.

        Masked rather than dropped so the array shape is constant, and
        masked to negative infinity rather than to a large negative
        number, which would be an exchange rate instead of a prohibition.
        """
        expected = self.kernel @ values
        scored = self.reward + discount.gamma * expected
        return np.where(self.available, scored, -np.inf)


class SolveResult(NamedTuple):
    """A value function, with how much to trust it.

    The bound is free — it falls straight out of the contraction — and
    reporting a value function without it is reporting a number with an
    unstated accuracy.
    """

    values: FloatArray
    sweeps: int
    final_gap: float
    error_bound: float
    converged: bool


def value_iteration(
    mdp: MarkovDecisionProcess,
    discount: Discount,
    tolerance: float = 1e-10,
    max_sweeps: int = 100_000,
) -> SolveResult:
    """Iterate the optimality operator, with a real error bound.

    Guard clauses first, then the loop: the interesting path is the
    iteration and everything else is an early exit.
    """
    if tolerance <= 0.0:
        raise ValueError(f"tolerance must be positive, got {tolerance}")

    values = np.zeros(mdp.n_states, dtype=np.float64)
    gap = np.inf

    for sweep in range(1, max_sweeps + 1):
        updated = mdp.action_values(values, discount).max(axis=1)
        gap = float(np.max(np.abs(updated - values)))
        values = updated

        if gap < tolerance:
            return SolveResult(
                values=values,
                sweeps=sweep,
                final_gap=gap,
                error_bound=discount.gamma * gap / (1.0 - discount.gamma),
                converged=True,
            )

    return SolveResult(
        values=values,
        sweeps=max_sweeps,
        final_gap=gap,
        error_bound=discount.gamma * gap / (1.0 - discount.gamma),
        converged=False,
    )


def evaluate_policy_exactly(
    mdp: MarkovDecisionProcess, policy: FloatArray, discount: Discount
) -> FloatArray:
    """V^pi = (I - gamma P_pi)^-1 R_pi, solved rather than iterated.

    The structural point of the whole formalism: with the max removed
    the equation is linear, and a linear system is solved rather than
    approached. solve() rather than inv(), which is both stabler and the
    idiomatic way to say "I want x, not the inverse".
    """
    if policy.shape != (mdp.n_states, mdp.n_actions):
        raise MdpError(f"policy shape {policy.shape} does not match the MDP")
    if not np.allclose(policy.sum(axis=1), 1.0, atol=KERNEL_TOLERANCE):
        raise MdpError("policy rows must be distributions over actions")
    if np.any(policy[~mdp.available] > 0.0):
        raise MdpError("policy places mass on an unavailable action")

    transition = np.einsum("sa,sat->st", policy, mdp.kernel)
    expected_reward = np.einsum("sa,sa->s", policy, mdp.reward)
    system = np.eye(mdp.n_states) - discount.gamma * transition

    try:
        return np.linalg.solve(system, expected_reward)
    except np.linalg.LinAlgError as error:
        raise SingularEvaluation(
            "(I - gamma P_pi) is singular; at gamma near 1 with a policy that never "
            "terminates the expectation equation has no solution, and the model is "
            "what needs fixing"
        ) from error


def greedy_policy(
    mdp: MarkovDecisionProcess, values: FloatArray, discount: Discount
) -> IntArray:
    """Act greedily, breaking ties by lowest index.

    Tie-breaking is a policy decision rather than an implementation
    detail: exact ties are common in structured problems, and an
    unspecified rule lets a re-solve return a different policy with an
    identical value, which reads as instability and is not.
    """
    return np.argmax(mdp.action_values(values, discount), axis=1).astype(np.int64)


class ContractionReport(NamedTuple):
    gap_before: float
    gap_after: float
    observed_modulus: float
    gamma: float
    holds: bool


def contraction_report(
    mdp: MarkovDecisionProcess, discount: Discount, first: FloatArray, second: FloatArray
) -> ContractionReport:
    """The inequality everything rests on, measured rather than assumed.

        ||T V - T U||_inf  <=  gamma * ||V - U||_inf

    A cheap assertion to keep in a test suite: if it fails, the kernel
    or the discount is malformed in a way the constructor missed, and
    every guarantee downstream is void.
    """
    before = float(np.max(np.abs(first - second)))
    after = float(
        np.max(
            np.abs(
                mdp.action_values(first, discount).max(axis=1)
                - mdp.action_values(second, discount).max(axis=1)
            )
        )
    )
    return ContractionReport(
        gap_before=before,
        gap_after=after,
        observed_modulus=after / before if before > 0.0 else 0.0,
        gamma=discount.gamma,
        holds=after <= discount.gamma * before + 1e-9,
    )


def policy_gap(
    mdp: MarkovDecisionProcess, policy: FloatArray, optimal_values: FloatArray, discount: Discount
) -> float:
    """How much a policy actually loses, in sup-norm.

    Compare policies by their exact value, never by the value functions
    that produced them: two value functions can differ substantially and
    induce the same greedy policy, and two that differ slightly can
    induce different ones. The policy is the artefact; the value
    function is scaffolding.
    """
    return float(np.max(np.abs(optimal_values - evaluate_policy_exactly(mdp, policy, discount))))

"""The Bellman backup, shaped for a kernel that is mostly zeros.

A real MDP has a handful of reachable successors per state-action pair.
A dense (S, A, S) kernel is therefore almost entirely zeros, and a dense
backup spends nearly all of its time multiplying by them. Flattening
into one CSR matrix over state-action pairs turns the expectation into a
single sparse matrix-vector product whose cost is the number of
transitions that exist.

What this does not change: the sweep count still comes from gamma, so a
faster sweep does not shorten a long horizon. It makes each of the
thousand sweeps cheap; it does not make there be fewer of them.
"""

from __future__ import annotations

import heapq
from dataclasses import dataclass

import numpy as np
import numpy.typing as npt
from scipy import sparse

FloatArray = npt.NDArray[np.float64]
IndexArray = npt.NDArray[np.int32]


@dataclass
class SparseMdp:
    """Transitions as one CSR matrix over flattened (state, action) rows.

    transitions is (S·A, S): row s*A + a holds the successor
    distribution for that pair. One sparse matvec then computes the
    expectation for every pair simultaneously, which is the whole
    reason for the flattening.

    The layout is immutable by construction — CSR indices encode the
    sparsity pattern — so a kernel re-estimated on a schedule means a
    rebuild. That is the right trade for a solver and the wrong one for
    a model that is being edited.
    """

    transitions: sparse.csr_matrix
    reward: FloatArray            # (S·A,), flattened to match
    available: npt.NDArray[np.bool_]  # (S·A,)
    n_states: int
    n_actions: int

    @classmethod
    def from_dense(
        cls,
        kernel: FloatArray,
        reward: FloatArray,
        available: npt.NDArray[np.bool_],
    ) -> SparseMdp:
        n_states, n_actions, _ = kernel.shape
        flattened = kernel.reshape(n_states * n_actions, n_states)

        return cls(
            transitions=sparse.csr_matrix(flattened, dtype=np.float64),
            reward=np.ascontiguousarray(reward.reshape(-1), dtype=np.float64),
            available=np.ascontiguousarray(available.reshape(-1)),
            n_states=n_states,
            n_actions=n_actions,
        )

    @property
    def density(self) -> float:
        """What fraction of the kernel is actually non-zero.

        Worth reporting rather than assuming: if this is near one the
        sparse layout is pure overhead and the dense backup was
        correct. It is the number that decides which implementation is
        appropriate, and it is cheap.
        """
        return self.transitions.nnz / (self.transitions.shape[0] * self.n_states)


class BackupWorkspace:
    """Every buffer a sweep needs, allocated once.

    A solver running ten thousand sweeps at a high discount allocates
    nothing after construction. The buffers are not reentrant, which is
    the price: two solves cannot share a workspace.
    """

    def __init__(self, mdp: SparseMdp) -> None:
        self.pairs = np.empty(mdp.n_states * mdp.n_actions, dtype=np.float64)
        self.updated = np.empty(mdp.n_states, dtype=np.float64)
        self.residual = np.empty(mdp.n_states, dtype=np.float64)


def fused_optimality_backup(
    mdp: SparseMdp, values: FloatArray, gamma: float, workspace: BackupWorkspace
) -> FloatArray:
    """One backup: sparse matvec, mask and max, in a single pass.

    The (S·A) action-value array is written into preallocated scratch
    rather than returned, so nothing is materialized per sweep. The
    availability mask is applied as negative infinity rather than as a
    large negative reward, which keeps a prohibition a prohibition
    instead of an exchange rate the optimizer can trade against.
    """
    # Expectation for every (state, action) pair at once. Cost is the
    # number of transitions that exist, not the square of the state
    # count -- which is the entire point of the layout.
    np.multiply(mdp.transitions.dot(values), gamma, out=workspace.pairs)
    np.add(workspace.pairs, mdp.reward, out=workspace.pairs)
    np.copyto(workspace.pairs, -np.inf, where=~mdp.available)

    # Row-wise max over actions, into the preallocated value buffer.
    np.max(
        workspace.pairs.reshape(mdp.n_states, mdp.n_actions), axis=1, out=workspace.updated
    )
    return workspace.updated


def value_iteration(
    mdp: SparseMdp, gamma: float, tolerance: float = 1e-10, max_sweeps: int = 100_000
) -> dict[str, object]:
    """Jacobi iteration with the contraction's error bound.

    Sweep count is set by gamma and the tolerance and not by anything
    done here. That is worth restating because it is the thing a faster
    backup cannot fix: at gamma 0.999 the answer needs roughly a
    thousand times the sweeps of gamma 0.9, and each of them is now
    cheap.
    """
    workspace = BackupWorkspace(mdp)
    values = np.zeros(mdp.n_states, dtype=np.float64)
    gap = np.inf

    for sweep in range(1, max_sweeps + 1):
        updated = fused_optimality_backup(mdp, values, gamma, workspace)
        np.subtract(updated, values, out=workspace.residual)
        gap = float(np.max(np.abs(workspace.residual)))
        values = values.copy() if sweep == 1 else values
        np.copyto(values, updated)

        if gap < tolerance:
            break

    return {
        "values": values,
        "sweeps": sweep,
        "final_gap": gap,
        "error_bound": gamma * gap / (1.0 - gamma),
        "converged": gap < tolerance,
    }


def gauss_seidel_sweep(
    mdp: SparseMdp, values: FloatArray, gamma: float, order: IndexArray
) -> float:
    """In-place backup reading its own updates within the sweep.

    Jacobi computes every new value from the old array; Gauss-Seidel
    reads whatever is current, so information propagates across the
    state space within a single pass instead of one state per sweep.
    On most structured problems that is a meaningful reduction in sweep
    count for no extra arithmetic.

    What it costs is the clean operator semantics: this is no longer an
    application of T to a value function, so the contraction bound needs
    restating and the ordering now affects the trajectory. The fixed
    point is the same; the path to it is not.
    """
    largest_change = 0.0
    indptr = mdp.transitions.indptr
    indices = mdp.transitions.indices
    data = mdp.transitions.data

    for state in order:
        best = -np.inf
        base = int(state) * mdp.n_actions

        for action in range(mdp.n_actions):
            row = base + action
            if not mdp.available[row]:
                continue

            expectation = 0.0
            for position in range(indptr[row], indptr[row + 1]):
                expectation += data[position] * values[indices[position]]

            candidate = mdp.reward[row] + gamma * expectation
            if candidate > best:
                best = candidate

        change = abs(best - values[state])
        if change > largest_change:
            largest_change = change
        values[state] = best

    return largest_change


def prioritized_sweeping(
    mdp: SparseMdp,
    gamma: float,
    predecessors: list[list[int]],
    tolerance: float = 1e-10,
    max_updates: int = 1_000_000,
) -> dict[str, object]:
    """Update the states that are still wrong, in residual order.

    A uniform sweep spends the same effort on a state whose value has
    converged as on one that has not. Prioritized sweeping keeps a queue
    ordered by Bellman residual and pushes a state's predecessors
    whenever its value moves, so work follows the information.

    Whether it pays is a property of the problem: when the residual is
    concentrated -- a goal-directed problem where value propagates
    backwards from a terminal state -- it is a large win; on a uniformly
    wrong value function the heap is pure overhead against a vectorized
    sweep. The density and residual distribution decide, not taste.
    """
    values = np.zeros(mdp.n_states, dtype=np.float64)
    queue: list[tuple[float, int]] = []
    queued: set[int] = set()

    for state in range(mdp.n_states):
        heapq.heappush(queue, (-abs(mdp.reward[state * mdp.n_actions]), state))
        queued.add(state)

    updates = 0
    while queue and updates < max_updates:
        priority, state = heapq.heappop(queue)
        queued.discard(state)

        if -priority < tolerance:
            continue

        previous = values[state]
        gauss_seidel_sweep(mdp, values, gamma, np.array([state], dtype=np.int32))
        updates += 1

        change = abs(values[state] - previous)
        if change < tolerance:
            continue

        # The change propagates backwards, so predecessors are the
        # states whose values are now stale.
        for predecessor in predecessors[state]:
            if predecessor in queued:
                continue
            heapq.heappush(queue, (-gamma * change, predecessor))
            queued.add(predecessor)

    return {"values": values, "updates": updates}


def sweeps_needed(gamma: float, tolerance: float, initial_gap: float = 1.0) -> int:
    """How many backups the tolerance needs, before running any.

    Straight from the contraction, and the honest way to size a job: the
    cost of a long horizon is knowable in advance rather than discovered
    after an hour. Nothing in this file changes this number.
    """
    if not 0.0 <= gamma < 1.0:
        raise ValueError(f"gamma must lie in [0, 1); got {gamma}")
    if gamma == 0.0:
        return 1
    return int(np.ceil(np.log(tolerance / initial_gap) / np.log(gamma)))