# -- Code Cell --
import pandas as pd
import numpy as np
train = pd.read_csv('./p1_train.csv')

# -- Code Cell --
train.head()

# -- Code Cell --
train[train['state'] == train['next_state']]['reward'].unique()

# -- Code Cell --
train['reward'].value_counts()

# -- Code Cell --
GRID= 25
N_STATES= GRID * GRID   
N_ACTIONS  = 4
GAMMA      = 0.90
GOAL_STATE = 500
EPSILON = 1e-4
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 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 --
seen

# -- 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 --
R = np.where(seen,reward_sa / np.where(seen, count_sa, 1.0),-1)

# -- 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 = pd.read_csv("p1_test.csv")

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

# -- Code Cell --
subs = pd.DataFrame({
    "state":test['state'],
    'action':y_pred
}).to_csv('submission_p1.csv',index=False)

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

# -- Code Cell --
