"""A small, dependency-free Q-learning lab. Run with Python 3.10+.

Original teaching example for the reinforcement-learning public guide.
The 100-step cap truncates collection; only the endpoints are terminal.
"""

import random
from statistics import mean

ALPHA = 0.2
GAMMA = 0.95
EPISODES = 600
MAX_STEPS = 100
EVAL_EPISODES = 1000


def step(state, action):
    """Actions: 0 = left, 1 = right. Terminal states: 0 and 6."""
    next_state = state + (-1 if action == 0 else 1)
    terminated = next_state in (0, 6)
    reward = 1.0 if next_state == 6 else -1.0 if next_state == 0 else -0.1
    return next_state, reward, terminated


def greedy(q, state, rng):
    """Break ties randomly instead of always favoring the first action."""
    best = max(q[state])
    return rng.choice([a for a in (0, 1) if q[state][a] == best])


def train(seed):
    rng = random.Random(seed)
    q = [[0.0, 0.0] for _ in range(7)]
    for episode in range(EPISODES):
        state = 3
        epsilon = max(0.05, 1.0 - episode / 400)
        for _ in range(MAX_STEPS):
            action = (rng.randrange(2) if rng.random() < epsilon
                      else greedy(q, state, rng))
            next_state, reward, terminated = step(state, action)
            future = 0.0 if terminated else max(q[next_state])
            target = reward + GAMMA * future
            q[state][action] += ALPHA * (target - q[state][action])
            state = next_state
            if terminated:
                break
        # Hitting MAX_STEPS resets collection, but does not zero bootstrap.
    return q


def evaluate(q, seed):
    """Freeze learning. q=None evaluates uniformly random actions."""
    rng = random.Random(seed)
    wins, truncations, returns, lengths = 0, 0, [], []
    for _ in range(EVAL_EPISODES):
        state, total, terminated = 3, 0.0, False
        for length in range(1, MAX_STEPS + 1):
            action = rng.randrange(2) if q is None else greedy(q, state, rng)
            state, reward, terminated = step(state, action)
            total += reward  # Report undiscounted return separately.
            if terminated:
                break
        wins += state == 6
        truncations += not terminated
        returns.append(total)
        lengths.append(length)
    return wins / EVAL_EPISODES, mean(returns), mean(lengths), truncations


def report(label, result):
    success, avg_return, avg_length, truncated = result
    print(f"{label:12} success={success:6.1%}  return={avg_return:6.3f}  "
          f"steps={avg_length:5.2f}  truncated={truncated}")


if __name__ == "__main__":
    report("random", evaluate(None, 1000))
    for seed in range(5):
        table = train(seed)
        report(f"Q seed {seed}", evaluate(table, 2000 + seed))
    print("Q at start, last seed:", [round(v, 4) for v in table[3]])