"""
norvig.mdp – AIMA Chapter 17
==============================
Uses numpy as the numerical engine for value/policy iteration.
Wraps pymdptoolbox when available for richer MDP support.

    from ml_ai_library.norvig.mdp import (
        MDP,            # abstract interface
        ValueIteration, # numpy-based AIMA Fig 17.4
        PolicyIteration,# numpy-based AIMA Fig 17.7
        GridWorldMDP,   # classic 4x3 grid example
    )
"""

from __future__ import annotations
from typing import Any, Dict, List, Optional, Tuple

import numpy as np


# ---------------------------------------------------------------------------
# Abstract MDP interface (AIMA §17.1)
# ---------------------------------------------------------------------------

class MDP:
    """Abstract MDP.  Override: states, actions, T, R, terminal."""

    def states(self) -> List[Any]:
        raise NotImplementedError

    def actions(self, state: Any) -> List[Any]:
        raise NotImplementedError

    def T(self, state: Any, action: Any) -> List[Tuple[float, Any]]:
        """Return [(prob, next_state), ...] transition distribution."""
        raise NotImplementedError

    def R(self, state: Any, action: Any, next_state: Any) -> float:
        raise NotImplementedError

    def terminal(self, state: Any) -> bool:
        return False

    def to_numpy(self):
        """
        Convert to dense numpy arrays (S x A x S transition, S x A reward).
        Used by pymdptoolbox wrappers.
        Returns (T_array shape (A,S,S), R_array shape (S,A)).
        """
        states  = self.states()
        idx     = {s: i for i, s in enumerate(states)}
        S       = len(states)
        actions_all = sorted({a for s in states for a in self.actions(s)},
                             key=str)
        A       = len(actions_all)
        aidx    = {a: i for i, a in enumerate(actions_all)}

        T = np.zeros((A, S, S))
        R = np.zeros((S, A))
        for s in states:
            si = idx[s]
            for a in self.actions(s):
                ai = aidx[a]
                for prob, s2 in self.T(s, a):
                    T[ai, si, idx[s2]] += prob
                    R[si, ai] += prob * self.R(s, a, s2)
        # Rows must sum to 1; fill missing with self-loop
        for ai in range(A):
            for si in range(S):
                row_sum = T[ai, si].sum()
                if row_sum == 0:
                    T[ai, si, si] = 1.0
        return T, R, actions_all, idx


# ---------------------------------------------------------------------------
# Value Iteration – numpy matrix operations (AIMA Fig 17.4)
# ---------------------------------------------------------------------------

class ValueIteration:
    """
    V_{k+1}(s) = max_a Σ_s' T(s,a,s')[R(s,a,s') + γ V_k(s')]

    Uses numpy vectorised operations for the Bellman update.
    """

    def __init__(self, mdp: MDP, gamma: float = 0.9):
        self.mdp   = mdp
        self.gamma = gamma

    def solve(self, epsilon: float = 1e-6,
              max_iter: int = 10_000) -> Tuple[Dict, Dict]:
        """Returns (V dict, pi dict)."""
        mdp    = self.mdp
        states = mdp.states()
        T_arr, R_arr, actions_all, idx = mdp.to_numpy()

        S     = len(states)
        V     = np.zeros(S)
        gamma = self.gamma

        for _ in range(max_iter):
            # Q(s,a) = R(s,a) + gamma * sum_{s'} T(a,s,s') * V(s')
            # Shape: (A, S) for Q; then max over A → new V
            Q     = R_arr.T + gamma * np.einsum("aij,j->ai", T_arr, V)
            V_new = Q.max(axis=0)
            # Terminal states stay at 0
            for s in states:
                if mdp.terminal(s):
                    V_new[idx[s]] = 0.0
            if np.max(np.abs(V_new - V)) < epsilon:
                V = V_new
                break
            V = V_new

        # Extract policy
        Q_final = R_arr.T + gamma * np.einsum("aij,j->ai", T_arr, V)
        pi_idx  = Q_final.argmax(axis=0)
        V_dict  = {s: float(V[idx[s]]) for s in states}
        pi_dict = {s: None if mdp.terminal(s) else actions_all[pi_idx[idx[s]]]
                   for s in states}
        return V_dict, pi_dict


# ---------------------------------------------------------------------------
# Policy Iteration – numpy (AIMA Fig 17.7)
# ---------------------------------------------------------------------------

class PolicyIteration:
    """
    Alternates policy evaluation (linear-system solve via numpy) and
    policy improvement.
    """

    def __init__(self, mdp: MDP, gamma: float = 0.9):
        self.mdp   = mdp
        self.gamma = gamma

    def solve(self) -> Tuple[Dict, Dict]:
        mdp    = self.mdp
        states = mdp.states()
        T_arr, R_arr, actions_all, idx = mdp.to_numpy()
        S     = len(states)
        gamma = self.gamma

        # Initialise arbitrary policy (first action for each state)
        pi = np.zeros(S, dtype=int)

        while True:
            # Policy evaluation: solve (I - gamma * T_pi) V = R_pi
            T_pi = T_arr[pi, np.arange(S), :]           # (S, S)
            R_pi = R_arr[np.arange(S), pi]               # (S,)
            A_mat = np.eye(S) - gamma * T_pi
            V    = np.linalg.solve(A_mat, R_pi)

            # Policy improvement
            Q      = R_arr.T + gamma * np.einsum("aij,j->ai", T_arr, V)
            pi_new = Q.argmax(axis=0)

            # Terminal states: clear policy
            for s in states:
                if mdp.terminal(s):
                    pi_new[idx[s]] = 0

            if np.array_equal(pi_new, pi):
                break
            pi = pi_new

        V_dict  = {s: float(V[idx[s]]) for s in states}
        pi_dict = {s: None if mdp.terminal(s) else actions_all[pi[idx[s]]]
                   for s in states}
        return V_dict, pi_dict


# ---------------------------------------------------------------------------
# Classic 4×3 Grid World (AIMA Fig 17.1)
# ---------------------------------------------------------------------------

class GridWorldMDP(MDP):
    """
    4×3 grid world from AIMA Fig 17.1.
    Goal=(0,3) +1, Pit=(1,3) -1, Wall=(1,1) blocked.
    Stochastic: 80% intended direction, 10% each perpendicular.
    """

    ACTS = ["UP", "DOWN", "LEFT", "RIGHT"]
    ROWS, COLS = 3, 4
    GOAL, PIT, WALL = (0, 3), (1, 3), (1, 1)
    STEP_R = -0.04

    def states(self):
        return [(r, c) for r in range(self.ROWS)
                for c in range(self.COLS) if (r, c) != self.WALL]

    def actions(self, state):
        return [] if self.terminal(state) else self.ACTS

    def terminal(self, state):
        return state in (self.GOAL, self.PIT)

    def T(self, state, action):
        if self.terminal(state):
            return [(1.0, state)]
        perp = {"UP": ("LEFT","RIGHT"), "DOWN": ("LEFT","RIGHT"),
                "LEFT": ("UP","DOWN"), "RIGHT": ("UP","DOWN")}[action]
        outcomes: Dict[Any, float] = {}
        for a, p in [(action, 0.8), (perp[0], 0.1), (perp[1], 0.1)]:
            s2 = self._move(state, a)
            outcomes[s2] = outcomes.get(s2, 0.0) + p
        return list(outcomes.items())

    def R(self, state, action, next_state):
        if next_state == self.GOAL: return 1.0
        if next_state == self.PIT:  return -1.0
        return self.STEP_R

    def _move(self, state, action):
        r, c = state
        dr, dc = {"UP":(-1,0),"DOWN":(1,0),"LEFT":(0,-1),"RIGHT":(0,1)}[action]
        nr, nc = r + dr, c + dc
        if 0 <= nr < self.ROWS and 0 <= nc < self.COLS and (nr,nc) != self.WALL:
            return (nr, nc)
        return state
