"""Lesson 1 exercise — Asteroid Hop: value iteration by hand.

THE GAME (single player, so no opponent to worry about yet)

  You pilot a ship across an asteroid field of positions 0..10.
  You start at position 0. Reaching position 10 (or past it -- cap
  at 10) wins the game: reward +1, episode over. Every other reward
  is 0. State 10 is terminal: nothing happens there, ever again.

  Two actions are available in every non-terminal state:

    THRUST -- move +1, guaranteed.
    JUMP   -- move +3 with probability 0.6;
              misfire and slide back -2 with probability 0.4
              (position floors at 0).

  Discount factor GAMMA = 0.9. The discount is what makes the choice
  interesting: without it, every policy wins eventually and all values
  are 1.0. With it, V(s) means roughly "win SOON," so speed vs. risk
  is a real tradeoff.

YOUR JOB

  Implement the three functions below, then run this file. You should
  see value bleed backward from state 10, one ring of states per
  sweep -- the "anchors propagate" claim from the lesson, live.

BEFORE YOU RUN IT, write down two predictions:

  P1: What does V[9] converge to, exactly?
      (You can compute both Q(9, THRUST) and Q(9, JUMP) in your
      head -- which wins?)
  P2: Will the optimal policy be all-THRUST, all-JUMP, or mixed?
      If mixed, which end of the board jumps? Think about what you'd
      do in a real game when you're far behind versus one step from
      winning.

ONE TRAP TO AVOID

  Leave V[10] = 0 forever. The +1 was collected on the transition
  INTO 10 -- terminal states have no future, hence no value. (This
  same off-by-one bites people in real AlphaZero implementations,
  so meet it now on an 11-state toy.)
"""

GAMMA = 0.9
N_STATES = 11  # positions 0..10; state 10 is terminal
THRUST = "THRUST"
JUMP = "JUMP"
ACTIONS = [THRUST, JUMP]


def outcomes(s, a):
    """Return the transition model for taking action `a` in state `s`.

    Returns a list of (prob, s_next, reward) tuples -- every state you
    might land in, the probability of landing there, and the reward
    collected on that transition.

    THRUST has one outcome; JUMP has two. Landing on (or past) 10
    caps at 10 and earns reward +1 -- the +1 belongs in the returned
    tuple's `reward` slot, NOT in V[10] (see the trap above). A
    misfire that would go negative floors at 0, with reward 0.
    The probabilities in the returned list must sum to 1.

    You may assume `s` is non-terminal (s < 10); the caller never
    passes s = 10.
    """
    # your code here
    raise NotImplementedError


def q_value(s, a, V):
    """Expected value of taking action `a` in state `s`, judged by the
    current value estimates `V` (a list of N_STATES floats).

    This is the one-step lookahead from the lesson:

        Q(s, a) = sum over outcomes of  prob * (reward + GAMMA * V[s_next])

    Don't forget the discount.
    """
    # your code here
    raise NotImplementedError


def value_iteration():
    """Run value iteration to convergence and return the final V.

    - Start with V = [0.0] * N_STATES.
    - Sweep: for each NON-terminal state s, replace V[s] with
      max over actions of q_value(s, a, V). (State 10 stays 0 -- see
      the trap in the module docstring.) Compute each sweep from a
      fresh copy of V (a "synchronous" update): updating V in place
      also converges, but then the printout below won't show the
      clean one-ring-per-sweep picture the lesson promised.
    - Track the largest single-state change in each sweep; stop when
      it drops below 1e-9.
    - Print V after every sweep, each value formatted to ~3 decimals.
      Watching the values fill in backward IS the point of the
      exercise, so don't skip the printing.
    """
    # your code here
    raise NotImplementedError


def greedy_policy(V):
    """Return the greedy policy as a dict {state: action} -- for each
    non-terminal state, the action whose q_value is highest.
    """
    # your code here
    raise NotImplementedError


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