import numpy as np


class MDP:
    def __init__(self, n_states, n_actions, P, R, gamma=0.95):
        """P[s, a] -> dict {s_prime: probability}
        R[s, a] -> dict {s_prime: reward}"""
        self.n_states = n_states
        self.n_actions = n_actions
        self.P = P
        self.R = R
        self.gamma = gamma

    def bellman_backup_state_value(self, V, s, a):
        return sum(
            prob * (self.R[s, a].get(sp, 0.0) + self.gamma * V[sp])
            for sp, prob in self.P[s, a].items()
        )


def value_iteration(mdp: MDP, theta=1e-8, max_iters=10_000):
    V = np.zeros(mdp.n_states)
    for _ in range(max_iters):
        delta = 0.0
        new_V = V.copy()
        for s in range(mdp.n_states):
            action_values = [
                mdp.bellman_backup_state_value(V, s, a)
                for a in range(mdp.n_actions)
            ]
            new_V[s] = max(action_values)
            delta = max(delta, abs(new_V[s] - V[s]))
        V = new_V
        if delta < theta:
            break

    policy = np.zeros(mdp.n_states, dtype=int)
    for s in range(mdp.n_states):
        action_values = [
            mdp.bellman_backup_state_value(V, s, a)
            for a in range(mdp.n_actions)
        ]
        policy[s] = int(np.argmax(action_values))
    return V, policy


def policy_evaluation(mdp: MDP, policy, theta=1e-8, max_iters=10_000):
    V = np.zeros(mdp.n_states)
    for _ in range(max_iters):
        delta = 0.0
        new_V = V.copy()
        for s in range(mdp.n_states):
            a = policy[s]
            new_V[s] = mdp.bellman_backup_state_value(V, s, a)
            delta = max(delta, abs(new_V[s] - V[s]))
        V = new_V
        if delta < theta:
            break
    return V


def policy_iteration(mdp: MDP, max_iters=1000):
    policy = np.zeros(mdp.n_states, dtype=int)
    for _ in range(max_iters):
        V = policy_evaluation(mdp, policy)
        stable = True
        for s in range(mdp.n_states):
            old_action = policy[s]
            action_values = [
                mdp.bellman_backup_state_value(V, s, a)
                for a in range(mdp.n_actions)
            ]
            policy[s] = int(np.argmax(action_values))
            if policy[s] != old_action:
                stable = False
        if stable:
            break
    return V, policy
