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

It contains:
  * the gridworld layout and a few geometry helpers (which cell is where,
    where an intended move lands, which directions are "perpendicular"),
  * the 3-state machine-maintenance MDP from the chapter, already built,
  * small policy helpers (uniform, deterministic, the gridworld presets),
  * a Monte Carlo episode sampler, used by the tests to check that your
    linear-algebra values really are expected returns,
  * a pretty-printer for gridworld values.

Conventions (the same as in the chapter and the interactive demo):
  * States are numbered row-major: s = row * N_COLS + col, row 0 at the top.
  * Actions are 0 = up, 1 = right, 2 = down, 3 = left.
  * An MDP is two numpy arrays:
        P[s, a, s2]  probability of landing in s2 after action a in state s
        R[s, a]      expected reward for taking action a in state s
  * A (stochastic) policy is an array pi[s, a] whose rows sum to 1.
"""
from __future__ import annotations

import numpy as np

# --------------------------------------------------------------------------
# The gridworld
# --------------------------------------------------------------------------
# . = ordinary cell, # = wall, G = goal (+1 on entering), P = pit (-1 on
# entering), S = start (an ordinary cell where episodes begin).
LAYOUT = [
    "...G",
    ".#.P",
    "S...",
]
N_ROWS = len(LAYOUT)
N_COLS = len(LAYOUT[0])
N_STATES = N_ROWS * N_COLS
N_ACTIONS = 4
ACTION_NAMES = ["up", "right", "down", "left"]
ARROWS = ["↑", "→", "↓", "←"]
DELTAS = [(-1, 0), (0, 1), (1, 0), (0, -1)]  # (d_row, d_col) for each action

# Reward collected on the step that ENTERS a cell of this kind.
EXIT_REWARD = {"G": 1.0, "P": -1.0}


def state_index(row: int, col: int) -> int:
    """(row, col) -> state number."""
    return row * N_COLS + col


def cell_of(s: int) -> tuple[int, int]:
    """State number -> (row, col)."""
    return divmod(s, N_COLS)


def kind(s: int) -> str:
    """The layout character of state s: '.', '#', 'G', 'P' or 'S'."""
    r, c = cell_of(s)
    return LAYOUT[r][c]


def is_terminal(s: int) -> bool:
    """Goal and pit end the episode. The wall cell is not a real state: nothing
    ever moves into it. We keep it (as an absorbing dummy) only so that the
    state number is simply row * N_COLS + col."""
    return kind(s) in ("G", "P", "#")


def move(s: int, a: int) -> int:
    """Where you end up if action a from state s goes exactly as intended.
    Moving off the grid or into the wall leaves you where you were ("bump")."""
    r, c = cell_of(s)
    dr, dc = DELTAS[a]
    r2, c2 = r + dr, c + dc
    if not (0 <= r2 < N_ROWS and 0 <= c2 < N_COLS) or LAYOUT[r2][c2] == "#":
        return s
    return state_index(r2, c2)


def perpendicular(a: int) -> tuple[int, int]:
    """The two actions at right angles to a (where you go if you slip)."""
    return ((a + 1) % 4, (a + 3) % 4)


START = next(state_index(r, c) for r in range(N_ROWS) for c in range(N_COLS) if LAYOUT[r][c] == "S")


# --------------------------------------------------------------------------
# The machine-maintenance MDP from the chapter (already built for you)
# --------------------------------------------------------------------------
MACHINE_STATES = ["new", "worn", "broken"]
MACHINE_ACTIONS = ["run", "repair"]


def machine_mdp() -> tuple[np.ndarray, np.ndarray]:
    """The 3-state machine from the chapter. Returns P with shape (3, 2, 3)
    and R with shape (3, 2). States: 0 new, 1 worn, 2 broken.
    Actions: 0 run, 1 repair."""
    P = np.zeros((3, 2, 3))
    R = np.zeros((3, 2))
    # run
    P[0, 0] = [0.6, 0.4, 0.0]; R[0, 0] = 2.0     # new: earns 2, wears out w.p. 0.4
    P[1, 0] = [0.0, 0.5, 0.5]; R[1, 0] = 0.0     # worn: earns 0, breaks w.p. 0.5
    P[2, 0] = [0.0, 0.0, 1.0]; R[2, 0] = 0.0     # broken: nothing happens
    # repair (always gives you a new machine)
    P[0, 1] = [1.0, 0.0, 0.0]; R[0, 1] = -1.0    # pointless tune-up
    P[1, 1] = [1.0, 0.0, 0.0]; R[1, 1] = -1.0    # tune-up of a worn machine
    P[2, 1] = [1.0, 0.0, 0.0]; R[2, 1] = -10.0   # major repair
    return P, R


# --------------------------------------------------------------------------
# Policies
# --------------------------------------------------------------------------
def uniform_policy(n_states: int, n_actions: int) -> np.ndarray:
    """pi[s, a] = 1 / n_actions everywhere."""
    return np.full((n_states, n_actions), 1.0 / n_actions)


def deterministic_policy(actions, n_actions: int) -> np.ndarray:
    """One-hot policy: actions[s] is the action taken in state s."""
    actions = np.asarray(actions)
    pi = np.zeros((len(actions), n_actions))
    pi[np.arange(len(actions)), actions] = 1.0
    return pi


def _grid_policy(rows: list[str]) -> np.ndarray:
    """Build a deterministic gridworld policy from arrow strings like 'RRR.'.
    U/R/D/L are actions; anything else (terminals, wall) gets 'up'."""
    code = {"U": 0, "R": 1, "D": 2, "L": 3}
    acts = [code.get(ch, 0) for row in rows for ch in row]
    return deterministic_policy(acts, N_ACTIONS)


# The same presets as the chapter's demo.
SHORTEST_PATH = _grid_policy(["RRR.", "U.U.", "RRUL"])
SAFE_DETOUR = _grid_policy(["RRR.", "U.U.", "ULLL"])
ALWAYS_RIGHT = _grid_policy(["RRRR", "RRRR", "RRRR"])
MACHINE_RUN_TO_FAILURE = deterministic_policy([0, 0, 1], 2)  # run, run, repair
MACHINE_MAINTAIN = deterministic_policy([0, 1, 1], 2)        # run, repair, repair


# --------------------------------------------------------------------------
# Monte Carlo: sample episodes and average their returns
# --------------------------------------------------------------------------
def sample_return(P, R, pi, gamma, s0, rng, max_steps=200) -> float:
    """Run one episode from s0 under pi and return the discounted return
    G_0 = r_0 + gamma r_1 + gamma^2 r_2 + ... (truncated after max_steps).
    Each step adds the expected reward R[s, a] rather than a random one; by
    linearity of expectation the AVERAGE return is unchanged, which is all
    the tests check."""
    s, G, disc = s0, 0.0, 1.0
    for _ in range(max_steps):
        a = rng.choice(pi.shape[1], p=pi[s])
        G += disc * R[s, a]
        disc *= gamma
        s = rng.choice(P.shape[2], p=P[s, a])
        if disc < 1e-12:
            break
    return G


def monte_carlo_value(P, R, pi, gamma, s0, n_episodes=4000, seed=0) -> tuple[float, float]:
    """Average return from s0 over n_episodes, and the standard error."""
    rng = np.random.default_rng(seed)
    Gs = np.array([sample_return(P, R, pi, gamma, s0, rng) for _ in range(n_episodes)])
    return float(Gs.mean()), float(Gs.std(ddof=1) / np.sqrt(n_episodes))


# --------------------------------------------------------------------------
# Printing
# --------------------------------------------------------------------------
def format_grid(V, pi=None) -> str:
    """A text picture of the gridworld with V(s) in each cell (and the
    policy's most likely arrow if pi is given)."""
    lines = []
    for r in range(N_ROWS):
        cells = []
        for c in range(N_COLS):
            s = state_index(r, c)
            k = kind(s)
            if k == "#":
                cells.append("   ####   ")
            elif k in ("G", "P"):
                cells.append(f"   {'+1' if k == 'G' else '-1'} {k}   ")
            else:
                arrow = ARROWS[int(np.argmax(pi[s]))] if pi is not None else " "
                if pi is not None and np.allclose(pi[s], pi[s][0]):
                    arrow = "*"
                cells.append(f" {arrow}{V[s]:+7.3f} ")
        lines.append("|".join(cells))
    return "\n".join(lines)