"""Helper code for the Chapter 4 lab. You do NOT need to edit this file.

Every MDP here is returned as plain numpy arrays, the same format as Chapter 3:

    P[s, a, s2]  probability of landing in s2 after taking action a in state s
    R[s, a]      expected reward for taking action a in state s
                 (= sum over s2 of P[s, a, s2] * reward for entering s2)

Terminal states (goal, pit / hole) are *absorbing*: every action keeps you
there with reward 0, so their value is always 0. The reward for reaching a
goal or a pit is paid on the step that enters it.

Gridworld actions are numbered 0 = up, 1 = right, 2 = down, 3 = left.
(Gymnasium's FrozenLake numbers them differently; we use this order everywhere
in the course demos.)
"""
from dataclasses import dataclass, field

import numpy as np

ACTIONS = [(-1, 0), (0, 1), (1, 0), (0, -1)]  # up, right, down, left as (d_row, d_col)
ARROWS = ["↑", "→", "↓", "←"]

# The gridworld from the chapter's demo. '.' free, '#' wall, 'G' goal (+1), 'P' pit (-1).
DEMO_GRID = [
    "......",
    ".###..",
    ".....G",
    ".PPPP.",
    "......",
]

# Gymnasium's FrozenLake maps. 'S' start (an ordinary free cell for DP), 'F' frozen
# (free), 'H' hole (terminal, reward 0), 'G' goal (terminal, reward +1).
FROZEN_LAKE_4x4 = ["SFFF", "FHFH", "FFFH", "HFFG"]
FROZEN_LAKE_8x8 = [
    "SFFFFFFF",
    "FFFFFFFF",
    "FFFHFFFF",
    "FFFFFHFF",
    "FFFHFFFF",
    "FHHFFFHF",
    "FHFFHFHF",
    "FFFHFFFG",
]


@dataclass
class MDP:
    """A finite MDP as arrays, plus enough bookkeeping to draw gridworlds."""
    P: np.ndarray                     # (S, A, S)
    R: np.ndarray                     # (S, A)
    terminal: np.ndarray              # (S,) bool
    names: list = field(default_factory=list)   # a label per state
    layout: list | None = None        # gridworld rows (strings), if a grid
    cells: list | None = None         # (row, col) of each state, if a grid

    @property
    def n_states(self):
        return self.P.shape[0]

    @property
    def n_actions(self):
        return self.P.shape[1]


def make_gridworld(layout, slip=0.0, step_reward=0.0, goal_reward=1.0, pit_reward=-1.0):
    """Build a gridworld MDP from a list of strings.

    Characters: '.', 'F', 'S' free; '#' wall (not a state); 'G' goal; 'P' or 'H' pit/hole.
    `slip` is the total probability of moving perpendicular to the intended
    direction (split equally between the two sides). Bumping into a wall or the
    edge leaves you where you are. Entering G pays `goal_reward`, entering P/H
    pays `pit_reward`, any other move pays `step_reward`.
    """
    n_rows, n_cols = len(layout), len(layout[0])
    cells = [(r, c) for r in range(n_rows) for c in range(n_cols) if layout[r][c] != "#"]
    index = {rc: i for i, rc in enumerate(cells)}
    S, A = len(cells), 4
    P = np.zeros((S, A, S))
    R = np.zeros((S, A))
    terminal = np.zeros(S, dtype=bool)

    def entry_reward(ch):
        if ch == "G":
            return goal_reward
        if ch in "PH":
            return pit_reward
        return step_reward

    for (r, c), s in index.items():
        if layout[r][c] in "GPH":
            terminal[s] = True
            P[s, :, s] = 1.0          # absorbing, reward 0
            continue
        for a in range(A):
            outcomes = [(a, 1.0 - slip), ((a + 1) % 4, slip / 2), ((a + 3) % 4, slip / 2)]
            for d, p in outcomes:
                if p == 0:
                    continue
                nr, nc = r + ACTIONS[d][0], c + ACTIONS[d][1]
                if not (0 <= nr < n_rows and 0 <= nc < n_cols) or layout[nr][nc] == "#":
                    nr, nc = r, c
                s2 = index[(nr, nc)]
                P[s, a, s2] += p
                R[s, a] += p * entry_reward(layout[nr][nc])
    names = [f"({r},{c})" for r, c in cells]
    return MDP(P, R, terminal, names, list(layout), cells)


def make_frozen_lake(layout=FROZEN_LAKE_4x4, slippery=True):
    """FrozenLake: on ice you move as intended with prob 1/3 and to each side with prob 1/3."""
    return make_gridworld(layout, slip=2 / 3 if slippery else 0.0, step_reward=0.0,
                          goal_reward=1.0, pit_reward=0.0)


def make_corridor():
    """The chapter's corridor:   +0.5 <- [A] <-> [B] <-> [C] -> +1

    States 0, 1, 2 are A, B, C; state 3 is the absorbing 'exit'.
    Actions: 0 = left, 1 = right. Moves are deterministic.
    Leaving A to the left pays +0.5, leaving C to the right pays +1, other moves pay 0.
    """
    S, A = 4, 2
    P = np.zeros((S, A, S))
    R = np.zeros((S, A))
    EXIT = 3
    P[0, 0, EXIT], R[0, 0] = 1.0, 0.5   # A, left: small exit
    P[0, 1, 1] = 1.0                    # A, right -> B
    P[1, 0, 0] = 1.0                    # B, left  -> A
    P[1, 1, 2] = 1.0                    # B, right -> C
    P[2, 0, 1] = 1.0                    # C, left  -> B
    P[2, 1, EXIT], R[2, 1] = 1.0, 1.0   # C, right: big exit
    P[EXIT, :, EXIT] = 1.0              # exit is absorbing
    terminal = np.array([False, False, False, True])
    return MDP(P, R, terminal, ["A", "B", "C", "exit"])


def random_policy(n_states, n_actions):
    """The uniform random policy as an (S, A) array of probabilities."""
    return np.full((n_states, n_actions), 1.0 / n_actions)


def deterministic_policy(actions, n_actions):
    """Turn a list/array of action indices (one per state) into a one-hot (S, A) policy."""
    actions = np.asarray(actions)
    pi = np.zeros((len(actions), n_actions))
    pi[np.arange(len(actions)), actions] = 1.0
    return pi


def render(mdp, V=None, policy=None):
    """Return a text picture of a gridworld: arrows for the policy, or values."""
    lines = []
    for r, row in enumerate(mdp.layout):
        out = []
        for c, ch in enumerate(row):
            if ch == "#":
                out.append("  ###  ")
                continue
            s = mdp.cells.index((r, c))
            if mdp.terminal[s]:
                out.append(f"  {ch*3}  ")
            elif V is not None and policy is None:
                out.append(f"{V[s]:7.3f}")
            elif policy is not None:
                a = int(np.argmax(policy[s]))
                val = f"{V[s]:.2f}" if V is not None else ""
                out.append(f"{ARROWS[a]}{val:>6}")
            else:
                out.append("   .   ")
        lines.append(" ".join(out))
    return "\n".join(lines)


def plot_values(ax, mdp, V, policy=None, title=None):
    """Draw a gridworld's values as a heatmap (matplotlib axis) with policy arrows."""
    from matplotlib.patches import Rectangle

    n_rows, n_cols = len(mdp.layout), len(mdp.layout[0])
    grid = np.full((n_rows, n_cols), np.nan)
    for s, (r, c) in enumerate(mdp.cells):
        grid[r, c] = V[s]
    ax.imshow(grid, cmap="viridis", vmin=min(0.0, np.nanmin(grid)), vmax=max(1e-9, np.nanmax(grid)))
    for s, (r, c) in enumerate(mdp.cells):
        ch = mdp.layout[r][c]
        if mdp.terminal[s]:
            ax.text(c, r, ch, ha="center", va="center", color="white", fontsize=10, weight="bold")
            continue
        ax.text(c, r + 0.3, f"{V[s]:.2f}", ha="center", va="center", color="white", fontsize=6)
        if policy is not None:
            a = int(np.argmax(policy[s]))
            dr, dc = ACTIONS[a]
            ax.arrow(c - 0.18 * dc, r - 0.18 * dr - 0.08, 0.3 * dc, 0.3 * dr, head_width=0.15,
                     head_length=0.1, color="white", length_includes_head=True)
    for r in range(n_rows):
        for c in range(n_cols):
            if mdp.layout[r][c] == "#":
                ax.add_patch(Rectangle((c - 0.5, r - 0.5), 1, 1, color="0.3"))
    ax.set_xticks([])
    ax.set_yticks([])
    if title:
        ax.set_title(title, fontsize=9)