# 值迭代：4×3 網格世界（Russell & Norvig 的經典例子）
# 每走一步報酬 r，(3,2) 終點 +1、(3,1) 終點 −1，(1,1) 是牆；往意圖方向走的機率 0.8，往左右兩側各 0.1
import numpy as np

ACTS = {"↑": (0, 1), "→": (1, 0), "↓": (0, -1), "←": (-1, 0)}
WALL, TERM = (1, 1), {(3, 2): 1.0, (3, 1): -1.0}
STATES = [(x, y) for y in range(3) for x in range(4) if (x, y) != WALL]

def move(s, a):
    n = (s[0] + ACTS[a][0], s[1] + ACTS[a][1])
    return n if 0 <= n[0] < 4 and 0 <= n[1] < 3 and n != WALL else s      # 撞牆就留在原地

def q(V, s, a, r, gamma=1.0):
    names = list(ACTS); i = names.index(a)
    outcomes = [(a, 0.8), (names[(i - 1) % 4], 0.1), (names[(i + 1) % 4], 0.1)]
    return sum(p * (r + gamma * V[move(s, b)]) for b, p in outcomes)

def value_iteration(r):
    V = {s: TERM.get(s, 0.0) for s in STATES}
    for it in range(1, 1000):
        newV = {s: V[s] if s in TERM else max(q(V, s, a, r) for a in ACTS) for s in STATES}   # Bellman 最佳方程式
        delta = max(abs(newV[s] - V[s]) for s in STATES)
        V = newV
        if delta < 1e-10:
            break
    pi = {s: max(ACTS, key=lambda a: q(V, s, a, r)) for s in STATES if s not in TERM}
    return V, pi, it

for r in (-0.04, -2.0):
    V, pi, it = value_iteration(r)
    print(f"每步報酬 r = {r}：{it} 次迭代收斂")
    for y in (2, 1, 0):
        vals = "  ".join("  牆  " if (x, y) == WALL else f"{V[(x, y)]:6.3f}" for x in range(4))
        arrows = " ".join("■" if (x, y) == WALL else ("+1" if (x, y) == (3, 2) else "-1" if (x, y) == (3, 1) else pi[(x, y)]) for x in range(4))
        print(f"  {vals}    {arrows}")