# -- Code Cell --
import pandas as pd

# -- Code Cell --
df=pd.read_csv('escape_logs.csv')

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

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

# -- Code Cell --
df[df['reward'] == 300]['next_state'].unique()

# -- Code Cell --
answers = []
for i in range(len(df)):
    if (abs(df['next_state'][i]) > abs(df['current_state'][i])) and abs((df['next_state'][i]) - abs(df['current_state'][i]))>2:
        # print(abs(df['next_state'][i]) - abs(df['current_state'][i]))
        answers.append(abs(df['next_state'][i]) - abs(df['current_state'][i]))
    if (abs(df['next_state'][i]) < abs(df['current_state'][i])) and (abs(df['current_state'][i]) -abs(df['next_state'][i]))>2:
        # print(abs(df['current_state'][i]) -abs(df['next_state'][i]))
        answers.append(abs(df['current_state'][i]) -abs(df['next_state'][i]))


# -- Code Cell --
answers

# -- Code Cell --
import numpy as np
np.unique(answers)

# -- Code Cell --
GRID = 12
N_STATE = GRID**2
N_ACTIONS = 4
GAMMA = 0.95
TARGETS =[132,11]
OUTCOME_DIRS = {0:[0,3,1,2],1:[1,2,0,3],2:[2,3,1,0], 3:[3,2,0,1]}
DELTA = {0:(1,0),1:(0,1),2:(-1,0),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_STATE,N_ACTIONS))
count_sas = np.zeros((N_STATE,N_ACTIONS,N_STATE))
reward_sas = np.zeros((N_STATE,N_ACTIONS,N_STATE))
for row in df.itertuples(index=False):
    s = row.current_state
    r = row.reward
    a = row.action
    ns = row.next_state
    count_sa[s,a] +=1
    count_sas[s,a,ns] +=1
    reward_sas[s,a,ns] +=r
    
    seen = count_sas>0

# -- Code Cell --
P = np.zeros((N_STATE,N_ACTIONS, N_STATE))
for s in range(N_STATE):
    s_row, s_col = divmod(s,GRID)
    for a in range(N_ACTIONS):
        if count_sa[s,a]>0:
            P[s,a] = count_sas[s,a] / count_sa[s,a]
        else:
            dr,dc= DELTA[a]
            ns_ = clip(s_row + dr, s_col + dc)
            P[s,a,ns_] = 1.0

# -- Code Cell --
R = np.where(seen,reward_sas / np.where(seen, count_sas, 1.0),-1)

# -- Code Cell --
Q = np.zeros((N_STATE,N_ACTIONS))
V = np.zeros((N_STATE))
for i in range(10000):
    Q = np.sum(P * (R + GAMMA * V[None, None, :]), axis=2)
    Q[TARGETS, :] = 0 
    V_new = Q.max(1)
    delta = np.abs(V_new - V).max()
    V = V_new
    if delta<1e-3:
        print(f"converged at {i}")
        break

# -- Code Cell --
V

# -- Code Cell --
query = pd.read_csv('query_states.csv')

# -- Code Cell --
state_test = query['state_id']
preds = []
for state in state_test:
    pred = V[state]
    preds.append(pred)
preds

# -- Code Cell --
state_test

# -- Code Cell --
subs = pd.DataFrame({
    "id": range(len(preds)),
    "v_star":preds  
}).to_csv("predictions.csv",index=False)

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

# -- Code Cell --
