import numpy as np


class GridWorld:
    """A simple deterministic gridworld with a cliff."""

    def __init__(self, width=6, height=4):
        self.width, self.height = width, height
        self.start = (0, 0)
        self.goal = (width - 1, 0)
        self.cliff = {(x, 0) for x in range(1, width - 1)}
        self.actions = [(0, 1), (0, -1), (-1, 0), (1, 0)]  # U D L R

    def reset(self):
        self.pos = self.start
        return self.pos

    def step(self, action_idx):
        dx, dy = self.actions[action_idx]
        x, y = self.pos
        nx = min(max(x + dx, 0), self.width - 1)
        ny = min(max(y + dy, 0), self.height - 1)
        next_pos = (nx, ny)

        if next_pos in self.cliff:
            self.pos = self.start
            return self.pos, -100.0, False
        if next_pos == self.goal:
            self.pos = next_pos
            return self.pos, -1.0, True
        self.pos = next_pos
        return self.pos, -1.0, False


def epsilon_greedy(action_values, n_actions, epsilon, rng):
    """action_values: the Q-table row for the current state alone
    (already indexed by the caller), i.e. Q[s] with shape
    (n_actions,) -- not the full Q-table."""
    if rng.random() < epsilon:
        return rng.integers(n_actions)
    return int(np.argmax(action_values))


def q_learning(
    env, n_episodes=500, alpha=0.5, gamma=1.0, epsilon=0.1, seed=0
):
    rng = np.random.default_rng(seed)
    Q = np.zeros((env.width, env.height, len(env.actions)))
    returns = []
    for _ in range(n_episodes):
        state = env.reset()
        total_reward = 0.0
        done = False
        while not done:
            a = epsilon_greedy(
                Q[state[0], state[1]], len(env.actions), epsilon, rng
            )
            next_state, reward, done = env.step(a)
            best_next = np.max(Q[next_state[0], next_state[1]])
            td_target = reward + gamma * best_next
            td_error = td_target - Q[state[0], state[1], a]
            Q[state[0], state[1], a] += alpha * td_error
            state = next_state
            total_reward += reward
        returns.append(total_reward)
    return Q, returns


def sarsa(
    env, n_episodes=500, alpha=0.5, gamma=1.0, epsilon=0.1, seed=0
):
    rng = np.random.default_rng(seed)
    Q = np.zeros((env.width, env.height, len(env.actions)))
    returns = []
    for _ in range(n_episodes):
        state = env.reset()
        a = epsilon_greedy(
            Q[state[0], state[1]], len(env.actions), epsilon, rng
        )
        total_reward = 0.0
        done = False
        while not done:
            next_state, reward, done = env.step(a)
            a_next = epsilon_greedy(
                Q[next_state[0], next_state[1]],
                len(env.actions), epsilon, rng,
            )
            next_q = Q[next_state[0], next_state[1], a_next]
            td_target = reward + gamma * next_q
            td_error = td_target - Q[state[0], state[1], a]
            Q[state[0], state[1], a] += alpha * td_error
            state, a = next_state, a_next
            total_reward += reward
        returns.append(total_reward)
    return Q, returns
