import numpy as np


# --- ` ---
N = 3
S = (0, 0)  # X^[g
G = (2, 2)  # S[
H = (2, 0)  # 
W = {(1, 1)}  # 
A = [(-1, 0), (0, 1), (1, 0), (0, -1)]  # E
MAX_STEPS = 50  # 1Gs\[h̍ős
reward_G = 100.0  # ViS[j
reward_H = -100.0  # Vij
reward_N = -1.0  # Viʏj


def sid(s):
    return s[0] * N + s[1]


def sid2pos(i):
    return (i // N, i % N)


def step(pos, a):
    r, c = pos
    dr, dc = A[a]
    nr, nc = r + dr, c + dc
    if not (0 <= nr < N and 0 <= nc < N) or (nr, nc) in W:
        nr, nc = r, c
    new = (nr, nc)
    if new == G:
        return new, reward_G, True
    if new == H:
        return new, reward_H, True
    return new, reward_N, False


# --- QwK ---
nS, nA = N * N, 4
Q = np.zeros((nS, nA))
alpha, gamma = 0.2, 0.99
eps_start, eps_end = 1.0, 0.05
episodes = 500
rng = np.random.default_rng(7)


def eps(ep):
    t = min(ep, episodes)
    return eps_start + (eps_end - eps_start) * (t / episodes)


def print_Q(ep):
    print(f"\n=== Episode {ep} ===")
    print("State:  [, , , ]")
    for state in range(nS):
        r, c = sid2pos(state)
        # ǂ͕\XLbv
        if (r, c) in W:
            print(f"{(r, c)}: WALL")
        else:
            print(f"{(r, c)}: {Q[state]}")


# --- wK[v ---
for ep in range(1, episodes+1):
    pos = S
    for _ in range(MAX_STEPS):
        s = sid(pos)
        if rng.random() < eps(ep):
            a = rng.integers(nA)
        else:
            a = np.argmax(Q[s])
        pos2, r, done = step(pos, a)
        s2 = sid(pos2)
        td = r if done else r + gamma * np.max(Q[s2])
        Q[s, a] += alpha * (td - Q[s, a])
        pos = pos2
        if done:
            break


    # --- Qe[u\ ---
    if ep <= 10 or ep % 50 == 0:
        print_Q(ep)


# --- wK̕\i}bvj ---
ar = {0:"", 1:"", 2:"", 3:""}
policy = np.argmax(Q, axis = 1)
grid = [["  "] * N for _ in range(N)]
for r in range(N):
    for c in range(N):
        if (r, c) in W:
            grid[r][c] = "W"
        elif (r, c) == H:
            grid[r][c] = "H"
        elif (r, c) == G:
            grid[r][c] = "G"
        else:
            grid[r][c] = ar[policy[sid((r,c))]]
print("\nywK z")
print("\n".join(" ".join(row) for row in grid))


# --- ȈՕ]i×~sł̓Bj ---
def greedy_episode():
    pos = S
    for _ in range(MAX_STEPS):
        a = np.argmax(Q[sid(pos)])
        pos, _, done = step(pos, a)
        if done:
            return pos == G
    return False


trials = 100
succ = sum(greedy_episode() for _ in range(trials))
print(f"\ny]zS[B: {succ/trials:.2f} ({succ}/{trials})")