# -- Code Cell --
import pandas as pd
import numpy as np

# -- Code Cell --

df_train = pd.read_csv("train.csv")
df_test  = pd.read_csv("test.csv")
df_train.head(5)

# -- Code Cell --
df_train[df_train["state"] == 6]

# -- Code Cell --
GRID       = 32
N_STATES   = GRID * GRID   
N_ACTIONS  = 4
GOAL_STATE = 777 
GAMMA      = 0.95
EPSILON    = 1e-4
STEP_REWARD = -1.0
 
ACTION_NAMES = ["up", "down", "left", "right"]
 
OUTCOME_DIRS = {0:[0,2,3,1], 1:[1,3,2,0], 2:[2,0,1,3], 3:[3,1,0,2]}
DELTAS       = {0:(-1,0), 1:(1,0), 2:(0,-1), 3:(0,1)}
def clip(r, c): return max(0, min(GRID-1, r)) * GRID + max(0, min(GRID-1, c))
 

# -- Code Cell --
count_sa  = np.zeros((N_STATES, N_ACTIONS),           dtype=np.float64)
count_sas = np.zeros((N_STATES, N_ACTIONS, N_STATES), dtype=np.float64)
reward_sa = np.zeros((N_STATES, N_ACTIONS),           dtype=np.float64)
 
for row in df_train.itertuples(index=False):
    s  = int(row.state)
    a  = int(row.action)
    ns = int(row.next_state)
    r  = float(row.reward)
    count_sa [s, a]      += 1
    count_sas[s, a, ns]  += 1
    reward_sa[s, a]      += r
 
seen = count_sa > 0   
 

# -- Code Cell --
P = np.zeros((N_STATES, N_ACTIONS, N_STATES), dtype=np.float64)
 
for s in range(N_STATES):
    row_s, col_s = divmod(s, GRID)
    for a in range(N_ACTIONS):
        if seen[s, a]:
            P[s, a] = count_sas[s, a] / count_sa[s, a]
        else:
            dr, dc = DELTAS[a]
            ns_    = clip(row_s + dr, col_s + dc)
            P[s, a, ns_] = 1.0
 

# -- Code Cell --
P[56,0][P[56, 0] != 0]

# -- Code Cell --
np.where(P[55, 0] != 0)

# -- Code Cell --
R = np.where(seen,reward_sa / np.where(seen, count_sa, 1.0),STEP_REWARD)

# -- Code Cell --
V = np.zeros(N_STATES, dtype=np.float64)
Q = np.zeros((N_STATES, N_ACTIONS), dtype=np.float64)
 
for iteration in range(1, 10_001):
    Q = R + GAMMA * (P @ V)
    Q[GOAL_STATE, :] = 0.0
    V_new = Q.max(axis=1)
    delta = np.abs(V_new - V).max()
    V = V_new
    if iteration <= 5 or iteration % 50 == 0 or delta < EPSILON:
        print(f"  iter {iteration:5d}  |ΔV|_max = {delta:.8f}")
    if delta < EPSILON:
        print(f"\n  ✓ Converged at iteration {iteration}  (ε = {EPSILON})")
        break
policy = Q.argmax(axis=1)  

# -- Code Cell --
test_states = df_test["state"].values
y_pred      = Q[test_states].argmax(axis=1)   # argmax Q for each test state

# -- Code Cell --
pd.DataFrame({
    "state" : test_states,
    "action" : y_pred
}).to_csv("submission.csv")

# -- Code Cell --
!python eval.py