# -- Code Cell --
import pandas as pd
test = pd.read_csv('test_data.csv')

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

# -- Code Cell --
GRID = 6
TARGET_CELL = 35
N_STATES = 36
N_ACTIONS = {"up",'down',"left",'right'} #0,1,2,3
MOVES = {0:(-1,0), 1:(0,1), 2:(1,0), 3:(0,-1)}
GAMMA = 0.9
STEP_REWARDS = -0.01

# -- Code Cell --
def clip(r,c):
    return max(0, min(GRID-1, r))*GRID + max(0, min(GRID-1, c))

# -- Code Cell --
import numpy as np
valori = [0, 1, 2]
prob = [0.8, 0.1, 0.1]
x = np.random.choice(valori, p=prob)
right_turn = {0:1, 1:2, 2:3, 3:0}
left_turn  = {0:3, 1:0, 2:1, 3:2}

# -- Code Cell --
cell_types = test.set_index('cell_index')['type'].to_dict()

# -- Code Cell --
def next_state(s, a):
    r, c = s // GRID, s % GRID
    dr, dc = MOVES[a]
    nr, nc = r+dr, c+dc
    if 0 <= nr < GRID and 0 <= nc < GRID and cell_types[nr*GRID+nc] != 1:
        return nr*GRID+nc
    return s

# -- Code Cell --
def transitions(s, a):
    res = {}
    for ns, p in [(next_state(s,a), 0.8), (next_state(s,right_turn[a]), 0.1), (next_state(s,left_turn[a]), 0.1)]:
        res[ns] = res.get(ns, 0) + p
    return res

# -- Code Cell --
V = np.zeros(N_STATES)
policy = np.zeros(N_STATES, dtype=int)

# -- Code Cell --
for _ in range(10000):
    delta = 0
    for s in range(N_STATES):
        if cell_types[s] in (1, 3):
            continue
        best_v, best_a = -np.inf, 0
        for a in range(4):
            q = 0
            for ns, p in transitions(s, a).items():
                if cell_types[ns] == 3:
                    q += p * 1.0          
                else:
                    q += p * (STEP_REWARDS + GAMMA * V[ns])
            if q > best_v:
                best_v, best_a = q, a
        delta = max(delta, abs(best_v - V[s]))
        V[s], policy[s] = best_v, best_a
    if delta < 1e-9:
        break

V[35] = 1.0

# -- Code Cell --
rows = []
for s in range(N_STATES):
    rows.append({'subtaskID': 1, 'datapointID': s, 'answer': float(V[s])})
for s in range(N_STATES):
    if cell_types[s] in (1, 3): 
        continue
    rows.append({'subtaskID': 2, 'datapointID': s, 'answer': int(policy[s])})
pd.DataFrame(rows).to_csv('subs.csv', index=False)