import random
from itertools import count
from retro_gamer import GameEnvironment, GameMetadata

class QLearning:
    """
    Trains a policy to play a game using Q-learning. 
    The Q-table maintains an estimate of the quality of every action at every state, 
    and the estimates are improved through training. 

    Arguments:
        episodes (int): How many episodes to run.
        alpha (float): Learning rate.
        gamma (float): Discount factor.
        epsilon (float): Starting exploration rate.
        epsilon_decay (float): Multiply epsilon by this each episode.
        epsilon_min (float): Epsilon never falls below this.
        max_steps_per_episode (int): Safety cutoff. Without this, a lucky
            random walk that keeps finding food (each pickup restores more
            energy than a turn costs) can make an episode run far longer
            than intended, or even effectively forever.
    """

    def __init__(
        self, episodes=1000, alpha=0.1, gamma=0.95, epsilon=1.0, 
        epsilon_decay=0.995, epsilon_min=0.05, max_steps_per_episode=500
    ):
        self.episodes = episodes
        self.alpha = alpha
        self.gamma = gamma
        self.epsilon = epsilon
        self.epsilon_decay = epsilon_decay
        self.epsilon_min = epsilon_min
        self.max_steps_per_episode = max_steps_per_episode

    def train(self, env, actions):
        """Trains the policy by updating estimates of the quality of state/actions
        in the Q-table.

        Arguments:
            env (retro_gamer.GameEnvironment): a game environment, which allows
                stepping through a game one action at a time.
        """
        self.Q = {}
        for episode in range(self.episodes):
            state = env.reset()
            turn = 0
            total_reward = 0
            for turn in range(self.max_steps_per_episode):
                action = self.choose_action(state, actions)
                next_state, reward, done = env.step(action)
                self.update_q(state, action, reward, next_state, actions)
                state = next_state
                total_reward += reward
            self.epsilon = max(self.epsilon_min, self.epsilon * self.epsilon_decay)
            self.report_progress(episode, total_reward)
        return self.Q

    def choose_action(self, state, actions):
        """Choose an action using an epsilon-greedy policy.
        With probability `self.epsilon`, return a random action from `actions`.
        Otherwise, return the action with the highest Q-value in `self.Q`
        for the given `state`. If a (state, action) pair has not been seen
        before, treat its Q-value as 0.
        """
        if random.random() < self.epsilon:
            return random.choice(actions)
        else:
            action_qualities = [(self.Q.get((state, a), 0), a) for a in actions]
            best_q, best_action = sorted(action_qualities, reverse=True)[0]
            return best_action

    def update_q(self, state, action, reward, next_state, next_actions):
        """Update an entry in self.Q. The update rule is: 

            Q(s, a) <- Q(s, a) + alpha * (r + gamma * max_a' Q(s', a') - Q(s, a))

        where:
            s, a                the state we were in and the action we took
            r                   the reward we received
            s'                  the next state 
            max_a' Q(s', a')    the best possible Q-value of actions from the next state
            alpha               the learning rate (use self.alpha)
            gamma               the discount factor (use self.gamma)
        """
        future_qs = [(self.Q.get((next_state, a), 0), a) for a in next_actions]
        best_future_q, best_future_action = sorted(future_qs, reverse=True)[0]
        q = self.Q.get((state, action), 0)
        self.Q[(state, action)] = q + self.alpha * (reward + self.gamma * best_future_q - q)

    def report_progress(self, episode, total_reward):
        if (episode + 1) % 100 == 0:
            print(
                f"Episode {episode + 1:5d}  "
                f"reward={total_reward:6.1f}  "
                f"epsilon={self.epsilon:.3f}  "
                f"q_entries={len(self.Q)}"
            )