"""
norvig.rl – AIMA Chapter 21
=============================
Tabular reinforcement learning.  Stable-Baselines3 covers deep RL;
this module fills the gap for classical tabular methods.

Classes
-------
QLearning        – off-policy TD control (AIMA §21.3)
SARSA            – on-policy TD control (AIMA §21.3)
TDLambda         – TD(λ) with eligibility traces (AIMA §21.4)
REINFORCE        – Monte-Carlo policy gradient (Williams 1992)
EpsilonGreedy    – ε-greedy exploration strategy
gymnasium wrappers exposed at bottom for convenience.

Example – CartPole with Q-Learning (discretised)
-------------------------------------------------
>>> import gymnasium as gym, numpy as np
>>> from ml_ai_library.norvig.rl import QLearning
>>>
>>> env = gym.make("CartPole-v1")
>>> agent = QLearning(n_actions=2, state_fn=my_discretise, alpha=0.1, gamma=0.99)
>>> agent.train(env, episodes=500)
>>> rewards = agent.evaluate(env, episodes=10)
"""

from __future__ import annotations
from typing import Any, Callable, Dict, List, Optional, Tuple
import numpy as np
from collections import defaultdict


# ---------------------------------------------------------------------------
# Exploration strategy
# ---------------------------------------------------------------------------

class EpsilonGreedy:
    """ε-greedy with optional linear decay."""

    def __init__(self, epsilon: float = 1.0, min_epsilon: float = 0.01,
                 decay: float = 0.995):
        self.epsilon     = epsilon
        self.min_epsilon = min_epsilon
        self.decay       = decay

    def choose(self, q_values: np.ndarray) -> int:
        if np.random.rand() < self.epsilon:
            return int(np.random.randint(len(q_values)))
        return int(np.argmax(q_values))

    def step(self):
        self.epsilon = max(self.min_epsilon, self.epsilon * self.decay)


# ---------------------------------------------------------------------------
# Tabular Q(s,a) storage
# ---------------------------------------------------------------------------

def _make_q():
    """Default Q-table: nested defaultdict → 0."""
    return defaultdict(lambda: defaultdict(float))


# ---------------------------------------------------------------------------
# Q-Learning (AIMA §21.3, off-policy)
# ---------------------------------------------------------------------------

class QLearning:
    """
    Tabular Q-Learning: Q(s,a) ← Q(s,a) + α[r + γ max_a' Q(s',a') - Q(s,a)]

    Parameters
    ----------
    n_actions  : number of discrete actions
    state_fn   : callable(obs) → hashable state key (default: identity)
    alpha      : learning rate
    gamma      : discount factor
    epsilon    : initial ε for ε-greedy exploration
    """

    def __init__(
        self,
        n_actions: int,
        state_fn: Callable = None,
        alpha: float = 0.1,
        gamma: float = 0.99,
        epsilon: float = 1.0,
        min_epsilon: float = 0.01,
        decay: float = 0.995,
    ) -> None:
        self.n_actions = n_actions
        self.state_fn  = state_fn or (lambda x: x)
        self.alpha     = alpha
        self.gamma     = gamma
        self.Q         = _make_q()
        self.explorer  = EpsilonGreedy(epsilon, min_epsilon, decay)
        self.rewards_history_: List[float] = []

    def _q_row(self, s) -> np.ndarray:
        return np.array([self.Q[s][a] for a in range(self.n_actions)])

    def act(self, obs) -> int:
        return self.explorer.choose(self._q_row(self.state_fn(obs)))

    def update(self, obs, action: int, reward: float, next_obs, done: bool):
        s  = self.state_fn(obs)
        s2 = self.state_fn(next_obs)
        target = reward + (0 if done else self.gamma * max(
            self.Q[s2][a] for a in range(self.n_actions)
        ))
        self.Q[s][action] += self.alpha * (target - self.Q[s][action])

    def train(self, env, episodes: int = 500, max_steps: int = 500) -> List[float]:
        for ep in range(episodes):
            obs, _ = env.reset()
            total  = 0.0
            for _ in range(max_steps):
                action = self.act(obs)
                next_obs, reward, terminated, truncated, _ = env.step(action)
                done = terminated or truncated
                self.update(obs, action, reward, next_obs, done)
                obs    = next_obs
                total += reward
                if done:
                    break
            self.explorer.step()
            self.rewards_history_.append(total)
        return self.rewards_history_

    def evaluate(self, env, episodes: int = 10, max_steps: int = 500) -> List[float]:
        saved_eps = self.explorer.epsilon
        self.explorer.epsilon = 0.0          # greedy
        rewards = []
        for _ in range(episodes):
            obs, _ = env.reset()
            total  = 0.0
            for _ in range(max_steps):
                action = self.act(obs)
                obs, reward, terminated, truncated, _ = env.step(action)
                total += reward
                if terminated or truncated:
                    break
            rewards.append(total)
        self.explorer.epsilon = saved_eps
        return rewards


# ---------------------------------------------------------------------------
# SARSA (on-policy TD) – AIMA §21.3
# ---------------------------------------------------------------------------

class SARSA:
    """
    SARSA: Q(s,a) ← Q(s,a) + α[r + γ Q(s',a') - Q(s,a)]
    On-policy: the next action a' is chosen by the same policy.
    """

    def __init__(
        self,
        n_actions: int,
        state_fn: Callable = None,
        alpha: float = 0.1,
        gamma: float = 0.99,
        epsilon: float = 1.0,
        min_epsilon: float = 0.01,
        decay: float = 0.995,
    ) -> None:
        self.n_actions = n_actions
        self.state_fn  = state_fn or (lambda x: x)
        self.alpha     = alpha
        self.gamma     = gamma
        self.Q         = _make_q()
        self.explorer  = EpsilonGreedy(epsilon, min_epsilon, decay)
        self.rewards_history_: List[float] = []

    def _q_row(self, s) -> np.ndarray:
        return np.array([self.Q[s][a] for a in range(self.n_actions)])

    def act(self, obs) -> int:
        return self.explorer.choose(self._q_row(self.state_fn(obs)))

    def train(self, env, episodes: int = 500, max_steps: int = 500) -> List[float]:
        for ep in range(episodes):
            obs, _ = env.reset()
            s      = self.state_fn(obs)
            a      = self.explorer.choose(self._q_row(s))
            total  = 0.0
            for _ in range(max_steps):
                next_obs, reward, terminated, truncated, _ = env.step(a)
                done = terminated or truncated
                s2   = self.state_fn(next_obs)
                a2   = self.explorer.choose(self._q_row(s2))
                target = reward + (0 if done else self.gamma * self.Q[s2][a2])
                self.Q[s][a] += self.alpha * (target - self.Q[s][a])
                s, a   = s2, a2
                total += reward
                if done:
                    break
            self.explorer.step()
            self.rewards_history_.append(total)
        return self.rewards_history_


# ---------------------------------------------------------------------------
# TD(λ) with eligibility traces (AIMA §21.4)
# ---------------------------------------------------------------------------

class TDLambda:
    """
    TD(λ) prediction for a fixed policy.
    Updates V(s) ← V(s) + α δ e(s) where e(s) are eligibility traces.

    Parameters
    ----------
    lam   : λ trace-decay parameter (0 = TD(0), 1 = Monte-Carlo)
    alpha : learning rate
    gamma : discount factor
    """

    def __init__(
        self,
        state_fn: Callable = None,
        lam: float = 0.9,
        alpha: float = 0.1,
        gamma: float = 0.99,
    ) -> None:
        self.state_fn = state_fn or (lambda x: x)
        self.lam   = lam
        self.alpha = alpha
        self.gamma = gamma
        self.V: Dict[Any, float] = defaultdict(float)

    def train(self, env, policy: Callable, episodes: int = 200,
              max_steps: int = 500) -> Dict:
        """
        Parameters
        ----------
        policy : callable(obs) → action
        """
        for _ in range(episodes):
            obs, _ = env.reset()
            e: Dict[Any, float] = defaultdict(float)     # eligibility traces
            for _ in range(max_steps):
                s      = self.state_fn(obs)
                action = policy(obs)
                next_obs, reward, terminated, truncated, _ = env.step(action)
                done   = terminated or truncated
                s2     = self.state_fn(next_obs)
                delta  = reward + (0 if done else self.gamma * self.V[s2]) - self.V[s]
                e[s]  += 1.0
                for state in list(e.keys()):
                    self.V[state] += self.alpha * delta * e[state]
                    e[state]      *= self.gamma * self.lam
                    if abs(e[state]) < 1e-6:
                        del e[state]
                obs = next_obs
                if done:
                    break
        return dict(self.V)


# ---------------------------------------------------------------------------
# REINFORCE – Monte-Carlo policy gradient (Williams 1992 / AIMA §21.5)
# ---------------------------------------------------------------------------

class REINFORCE:
    """
    REINFORCE with a PyTorch softmax policy network.

    Parameters
    ----------
    policy_net : torch.nn.Module  – outputs logits for each action
    optimizer  : torch.optim.Optimizer
    gamma      : discount factor
    """

    def __init__(self, policy_net, optimizer, gamma: float = 0.99):
        self.net   = policy_net
        self.opt   = optimizer
        self.gamma = gamma
        self.rewards_history_: List[float] = []

    def train(self, env, episodes: int = 500, max_steps: int = 500) -> List[float]:
        import torch
        import torch.nn.functional as F

        for ep in range(episodes):
            obs, _ = env.reset()
            log_probs, rewards = [], []

            for _ in range(max_steps):
                x    = torch.FloatTensor(obs).unsqueeze(0)
                logits = self.net(x)
                probs  = F.softmax(logits, dim=-1)
                dist   = torch.distributions.Categorical(probs)
                action = dist.sample()
                log_probs.append(dist.log_prob(action))

                obs, reward, terminated, truncated, _ = env.step(action.item())
                rewards.append(reward)
                if terminated or truncated:
                    break

            # Discounted returns
            G, returns = 0.0, []
            for r in reversed(rewards):
                G = r + self.gamma * G
                returns.insert(0, G)
            returns = torch.FloatTensor(returns)
            returns = (returns - returns.mean()) / (returns.std() + 1e-8)

            # Policy gradient loss
            loss = -sum(lp * G for lp, G in zip(log_probs, returns))
            self.opt.zero_grad()
            loss.backward()
            self.opt.step()
            self.rewards_history_.append(sum(rewards))

        return self.rewards_history_
