"""Lesson 1 solution — Asteroid Hop: value iteration.

See lesson_1_asteroid_hop.py for the full assignment. This is the
solution Owen wrote during Session 1, lightly cleaned up.

Converged result (59 sweeps to delta < 1e-9; the greedy policy stopped
changing at sweep 6 -- see the lesson's "policy freezes before values
converge" point):

    V* = [0.522, 0.548, 0.571, 0.619, 0.667,
          0.709, 0.780, 0.855, 0.900, 1.000]   (V[10] = 0, terminal)

    policy: JUMP for states 0-7, THRUST for 8-9
"""

GAMMA = 0.9
N_STATES = 11  # positions 0..10
TERMINAL = 10
THRUST = "THRUST"
JUMP = "JUMP"
ACTIONS = [THRUST, JUMP]


def outcomes(s, a):
    """All (prob, s_next, reward) outcomes of taking action `a` in state `s`.

    The +1 is collected on the transition INTO the terminal state;
    the terminal state itself is never updated and keeps V = 0.
    """
    if a == THRUST:
        s_next = min(s + 1, TERMINAL)
        return [(1.0, s_next, 1.0 if s_next == TERMINAL else 0.0)]
    # JUMP: +3 with p = 0.6, misfire -2 (floored at 0) with p = 0.4
    hit = min(s + 3, TERMINAL)
    miss = max(s - 2, 0)
    return [
        (0.6, hit, 1.0 if hit == TERMINAL else 0.0),
        (0.4, miss, 0.0),
    ]


def q_value(s, a, V):
    """One-step lookahead: E[r + GAMMA * V(s_next)] for action `a` in state `s`.

    GAMMA multiplies V[s_next] -- the future -- not the immediate reward.
    """
    return sum(p * (r + GAMMA * V[s_next]) for p, s_next, r in outcomes(s, a))


def value_iteration():
    """Sweep until the largest single-state change drops below 1e-9."""
    V = [0.0] * N_STATES
    sweep = 0
    while True:
        sweep += 1
        # Synchronous update: compute the whole sweep from the OLD values,
        # then swap. Updating V in place also converges (Gauss-Seidel
        # value iteration, often faster) -- but then later states in the
        # sweep would read this sweep's values, and the one-ring-per-sweep
        # picture printed below would lie to you.
        new_V = list(V)
        for s in range(TERMINAL):  # state 10 is terminal: V[10] stays 0
            new_V[s] = max(q_value(s, a, V) for a in ACTIONS)
        delta = max(abs(new_V[s] - V[s]) for s in range(N_STATES))
        V = new_V
        print(f"sweep {sweep:2d}: " + " ".join(f"{v:.3f}" for v in V))
        if delta < 1e-9:
            return V


def greedy_policy(V):
    """For each non-terminal state, the action with the highest Q."""
    return {s: max(ACTIONS, key=lambda a: q_value(s, a, V)) for s in range(TERMINAL)}


if __name__ == "__main__":
    V = value_iteration()
    print("Optimal policy:")
    for s, a in sorted(greedy_policy(V).items()):
        print(f"  state {s}: {a}")