def build_grid_mdp(size=4, gamma=0.9):
    n_states = size * size
    n_actions = 4  # up, down, left, right
    goal = n_states - 1
    P, R = {}, {}

    def to_rc(s):
        return divmod(s, size)

    def to_s(r, c):
        return r * size + c

    deltas = {0: (-1, 0), 1: (1, 0), 2: (0, -1), 3: (0, 1)}

    for s in range(n_states):
        for a in range(n_actions):
            if s == goal:
                P[s, a] = {s: 1.0}
                R[s, a] = {s: 0.0}
                continue
            r, c = to_rc(s)
            dr, dc = deltas[a]
            nr = max(0, min(size - 1, r + dr))
            nc = max(0, min(size - 1, c + dc))
            s_next = to_s(nr, nc)
            reward = 10.0 if s_next == goal else -1.0
            P[s, a] = {s_next: 1.0}
            R[s, a] = {s_next: reward}

    return MDP(n_states, n_actions, P, R, gamma)


if __name__ == "__main__":
    mdp = build_grid_mdp(size=4)

    V_vi, policy_vi = value_iteration(mdp)
    V_pi, policy_pi = policy_iteration(mdp)

    print("Value iteration and policy iteration agree on V:",
          np.allclose(V_vi, V_pi, atol=1e-4))
    print("Value iteration and policy iteration agree on policy:",
          np.array_equal(policy_vi, policy_pi))

    arrows = {0: "^", 1: "v", 2: "<", 3: ">"}
    size = 4
    for r in range(size):
        row = []
        for c in range(size):
            s = r * size + c
            row.append("G" if s == mdp.n_states - 1 else arrows[policy_vi[s]])
        print(" ".join(row))
