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

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

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

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

# -- Code Cell --
df[df['reward'] == 50]['next_state']

# -- Code Cell --
GRID = 16
N_STATES = GRID * GRID
N_ACTIONS = 4
TARGETS = [7,62,136,226,255]
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)}
gamma = 0.95
def clip(r, c): return max(0, min(GRID-1, r)) * GRID + max(0, min(GRID-1, c))

# -- Code Cell --
counts_sa = np.zeros((N_STATES,N_ACTIONS))
counts_sas = np.zeros((N_STATES,N_ACTIONS,N_STATES))
reward_sas = np.zeros((N_STATES,N_ACTIONS, N_STATES))
for row in df.itertuples(index=False):
    s = row.current_state
    a = row.action
    r = row.reward
    ns = row.next_state
    
    counts_sa[s,a] +=1
    counts_sas[s,a,ns] +=1
    reward_sas[s,a, ns] +=r

    seen = counts_sa > 0

# -- Code Cell --
seen

# -- Code Cell --
P = np.zeros((N_STATES, N_ACTIONS, N_STATES))
for s in range(N_STATES):
    s_row, s_col = divmod(s, GRID)
    for a in range(N_ACTIONS):
        if seen > 0:
            P[s,a] = counts_sas[s,a] / counts_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, counts_sas, 1.0),-1)

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

# -- 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 --
subs = pd.DataFrame({
    "state": state_test,
    "answer":preds  
}).to_csv("subs.csv",index=False)

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

# -- Code Cell --
