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

# -- Code Cell --
root_path = "/home/stefan/ioai-prep/kits/algolymp/markov"
seed = 42

# -- Markdown Cell --
# # Data

# -- Code Cell --
GAMMA = 0.9
CONVERGENCE_THRESHOLD = 1e-10
REWARD_GOAL = 1.0
REWARD_STEP = -0.01

# 0: Up, 1: Right, 2: Down, 3: Left
ACTIONS = [(-1, 0), (0, 1), (1, 0), (0, -1)]

# -- Code Cell --
test_df = pd.read_csv(f"{root_path}/test_data.csv")
grid_size = int(np.sqrt(len(test_df))) # 6
grid = test_df["type"].values.reshape((grid_size, grid_size))

# -- Code Cell --
grid

# -- Code Cell --
def get_next_state(r, c, move):
    nr, nc = r + move[0], c + move[1]
    
    # check boundaries and obstacles
    if 0 <= nr < grid_size and 0 <= nc < grid_size and grid[nr, nc] != 1:
        return nr, nc
    return r, c

# -- Markdown Cell --
# # Subtask 1

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

# -- Markdown Cell --
# $$ V_*(s) = \underset{a \in \mathcal{A}(s)}{\mathrm{max}} \sum_{s' \in \mathcal{S}, r \in \mathcal{R}} p(s', r|s,a) [r + \gamma V_*(s')] $$

# -- Code Cell --
max_depth = 1e6

while True:
    max_depth -= 1
    if max_depth == 0:
        break

    delta = 0
    new_V = V.copy()

    for r in range(grid_size):
        for c in range(grid_size):
            # terminal state and obstacles have no value/next moves
            if grid[r, c] == 3 or grid[r, c] == 1:
                continue

            q_values = []
            for a in range(4):
                # stochastic transitions: 0.8 intended, 0.1 CCW, 0.1 CW
                moves = [ACTIONS[a], ACTIONS[(a - 1) % 4], ACTIONS[(a + 1) % 4]]
                probs = [0.8, 0.1, 0.1]

                q_a = 0
                for move, prob in zip(moves, probs):
                    nr, nc = get_next_state(r, c, move) # s'
                    reward = REWARD_GOAL if grid[nr, nc] == 3 else REWARD_STEP # r
                    q_a += prob * (reward + GAMMA * V[nr, nc])

                q_values.append(q_a)

            best_q = max(q_values)
            delta = max(delta, abs(best_q - V[r, c]))
            new_V[r, c] = best_q

    V = new_V

    # check for convergence
    if delta < CONVERGENCE_THRESHOLD:
        break

# -- Code Cell --
subtask1 = []

for r in range(grid_size):
    for c in range(grid_size):
        idx = r * grid_size + c

        subtask1.append(V[r, c])

# -- Markdown Cell --
# # Subtask 2

# -- Markdown Cell --
# $$ q_*(s,a) = \sum_{s' \in \mathcal{S}, r \in \mathcal{R}} p(s', r|s,a) [r + \gamma v_*(s')] $$
# 
# $$ \pi_*(s) = \underset{a}{\mathrm{argmax}} \text{ } q_*(s,a) $$

# -- Code Cell --
# extract the policy
subtask2 = []

for r in range(grid_size):
    for c in range(grid_size):
        idx = r * grid_size + c

        if grid[r, c] in [1, 3]:
            continue 

        q_values = []
        for a in range(4):
            moves = [ACTIONS[a], ACTIONS[(a - 1) % 4], ACTIONS[(a + 1) % 4]]
            probs = [0.8, 0.1, 0.1]

            q_a = 0
            for move, prob in zip(moves, probs):
                nr, nc = get_next_state(r, c, move)
                reward = REWARD_GOAL if grid[nr, nc] == 3 else REWARD_STEP
                q_a += prob * (reward + GAMMA * V[nr, nc])

            q_values.append(q_a)

        subtask2.append(np.argmax(q_values))

# -- Markdown Cell --
# # Submission

# -- Code Cell --
def build_subtask(sid, answers):
    id_ = (
        test_df["cell_index"]
        if sid == 1
        else test_df[(test_df["type"] != 1) & (test_df["type"] != 3)]["cell_index"]
    )
    return pd.DataFrame({
        "subtaskID": sid,
        "datapointID": id_,
        "answer": answers
    })

subtasks = [
    (1, subtask1),
    (2, subtask2)
]

submission = pd.concat([build_subtask(sid, ans) for (sid, ans) in subtasks])

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

# -- Code Cell --
submission.to_csv(f"{root_path}/submission.csv", index=False)