"""Reinforcement Learning auction bidding agent with Double DQN.

Implements a full RL-based auction strategy for Fantacalcio using a
custom auction environment and a numpy-only Double DQN policy.
No gym/gymnasium required — environment and training loop are self-contained.

Integrates with AuctionConfig from auction_solver.py.
"""

import copy
import logging
import pickle
from collections import deque
from dataclasses import dataclass, field
from math import sqrt
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import numpy as np
import pandas as pd

from .auction_solver import AuctionConfig, PlayerValuation

logger = logging.getLogger(__name__)

ROLES = ["P", "D", "C", "A"]
ROLE_INDEX = {"P": 0, "D": 1, "C": 2, "A": 3}
N_ACTIONS = 11

BID_LEVELS = np.array(
    [0.0, 0.5, 1, 2, 3, 5, 8, 12, 20, 50, 100], dtype=np.float64
)


# ---------------------------------------------------------------------------
# Auction Environment
# ---------------------------------------------------------------------------


@dataclass
class AuctionEnv:
    """Full Fantacalcio auction simulator with N opponents.

    Simulates a sequential auction (player-by-player) where the agent
    bids against rule-based opponents. Observations are flat vectors;
    actions are discrete bid levels.
    """

    player_pool_df: pd.DataFrame
    n_opponents: int = 7
    config: Optional[AuctionConfig] = None

    rng: np.random.RandomState = field(default_factory=lambda: np.random.RandomState(None))

    def __post_init__(self):
        if self.config is None:
            self.config = AuctionConfig()
        self._slot_names = ["P", "D", "C", "A"]
        self._slot_quotas = {
            "P": self.config.n_gk,
            "D": self.config.n_def,
            "C": self.config.n_mid,
            "A": self.config.n_fwd,
        }
        self._total_slots = sum(self._slot_quotas.values())
        self._players = self._build_player_objects()
        self._player_order: list = []
        self._current_step = 0
        self._budget = 0.0
        self._roster: Dict[str, List] = {"P": [], "D": [], "C": [], "A": []}
        self._opponent_budgets: np.ndarray = np.array([])
        self._opponent_rosters: List[Dict[str, int]] = []
        self._opponent_aggression: np.ndarray = np.array([])
        self._purchased: set = set()
        self._role_pool_count: Dict[str, int] = {}
        self._obs_dim = 15

    def _build_player_objects(self) -> list:
        objs = []
        for _, row in self.player_pool_df.iterrows():
            role = str(row.get("role", ""))
            if role not in self._slot_names:
                continue
            objs.append({
                "name": str(row.get("name", "")),
                "role": role,
                "projected_points": float(row.get("projected_points", 6.0)),
                "market_value": float(row.get("market_value", 0.0)),
                "ceiling_price": min(
                    float(row.get("projected_points", 6.0)) * 5,
                    self.config.total_budget * self.config.max_single_bid_pct,
                ),
            })
        logger.info(
            f"AuctionEnv built: {len(objs)} players, "
            f"{self.n_opponents} opponents, budget={self.config.total_budget}"
        )
        return objs

    def _count_role_in_pool(self, role: str) -> int:
        return sum(1 for p in self._players if p["role"] == role)

    # ------------------------------------------------------------------
    # Reset
    # ------------------------------------------------------------------

    def reset(self) -> np.ndarray:
        self.rng.shuffle(self._players)
        self._player_order = list(range(len(self._players)))
        self._current_step = 0
        self._budget = float(self.config.total_budget)
        self._roster = {"P": [], "D": [], "C": [], "A": []}
        self._purchased = set()
        self._role_pool_count = {
            role: self._count_role_in_pool(role) for role in self._slot_names
        }

        self._opponent_budgets = np.full(
            self.n_opponents, float(self.config.total_budget), dtype=np.float64
        )
        self._opponent_rosters = [
            {"P": 0, "D": 0, "C": 0, "A": 0} for _ in range(self.n_opponents)
        ]
        self._opponent_aggression = self.rng.uniform(0.7, 1.3, size=self.n_opponents)

        logger.info(
            f"Auction reset: {len(self._players)} players, "
            f"budget={self._budget}"
        )
        return self._build_observation()

    # ------------------------------------------------------------------
    # Step
    # ------------------------------------------------------------------

    def step(self, action: int) -> Tuple[np.ndarray, float, bool, dict]:
        if self._current_step >= len(self._players):
            return self._build_observation(), 0.0, True, {"reason": "auction_complete"}

        player_idx = self._player_order[self._current_step]
        player = self._players[player_idx]
        role = player["role"]
        info: dict = {
            "player_name": player["name"],
            "player_role": role,
            "player_points": player["projected_points"],
            "action": action,
        }

        valid = self.get_valid_actions()
        if action not in valid:
            action = 0
            info["invalid_action_override"] = True

        # Determine agent bid amount
        agent_bid = BID_LEVELS[action] if action > 0 else 0.0
        agent_bid = min(agent_bid, self._budget)

        if action > 0 and role_slots_filled(self._roster, role) >= self._slot_quotas[role]:
            agent_bid = 0.0

        # Simulate opponents
        opponent_bids = []
        for opp_idx in range(self.n_opponents):
            obid = self._simulate_opponent_decision(player, opp_idx)
            opponent_bids.append(obid)
        opponent_bids = np.array(opponent_bids, dtype=np.float64)
        info["opponent_bids"] = opponent_bids.tolist()

        max_opp_bid = float(np.max(opponent_bids))

        reward = 0.0
        if agent_bid > 0 and agent_bid >= max_opp_bid:
            # Agent wins this player
            self._budget -= agent_bid
            self._roster[role].append(player)
            self._purchased.add(player_idx)
            points = player["projected_points"]
            cost_term = agent_bid / 5.0
            opportunity_cost = self._opportunity_cost_of_slot(role)
            reward = points - cost_term - opportunity_cost
            info["outcome"] = "won"
            info["price_paid"] = agent_bid
        else:
            info["outcome"] = "lost"
            info["price_paid"] = 0.0

        # Determine which opponent won the player
        winning_opp = int(np.argmax(opponent_bids)) if max_opp_bid > 0 else -1
        if winning_opp >= 0:
            opp_bid = opponent_bids[winning_opp]
            self._opponent_budgets[winning_opp] -= opp_bid
            self._opponent_rosters[winning_opp][role] += 1

        self._current_step += 1
        done = self._is_done()
        if done:
            info["reason"] = self._done_reason()

        return self._build_observation(), reward, done, info

    # ------------------------------------------------------------------
    # Observation
    # ------------------------------------------------------------------

    def _build_observation(self) -> np.ndarray:
        if self._current_step >= len(self._players):
            player = {"role": "P", "projected_points": 0.0, "ceiling_price": 0.0}
        else:
            player = self._players[self._player_order[self._current_step]]

        budget_pct = self._budget / max(self.config.total_budget, 1.0)
        role_onehot = np.zeros(4, dtype=np.float64)
        r_idx = ROLE_INDEX.get(player["role"], 0)
        role_onehot[r_idx] = 1.0

        vor = self._compute_vor(player)
        progress = self._current_step / max(len(self._players), 1)
        avg_opp_budget = float(np.mean(self._opponent_budgets))

        gk_slots = self._slot_quotas["P"] - role_slots_filled(self._roster, "P")
        def_slots = self._slot_quotas["D"] - role_slots_filled(self._roster, "D")
        mid_slots = self._slot_quotas["C"] - role_slots_filled(self._roster, "C")
        fwd_slots = self._slot_quotas["A"] - role_slots_filled(self._roster, "A")

        obs = np.array([
            self._budget,
            budget_pct,
            float(gk_slots),
            float(def_slots),
            float(mid_slots),
            float(fwd_slots),
            player["projected_points"],
            role_onehot[0],
            role_onehot[1],
            role_onehot[2],
            role_onehot[3],
            player["ceiling_price"],
            progress,
            vor,
            avg_opp_budget,
        ], dtype=np.float64)

        return obs

    def _compute_vor(self, player: dict) -> float:
        role = player["role"]
        same_role = [
            p for p in self._players
            if p["role"] == role
        ]
        if len(same_role) <= 1:
            return player["projected_points"]
        sorted_pts = sorted([p["projected_points"] for p in same_role], reverse=True)
        replacement_level = (
            sorted_pts[min(max(len(sorted_pts) // 4, 1), len(sorted_pts) - 1)]
        )
        return max(player["projected_points"] - replacement_level, 0.0)

    def _opportunity_cost_of_slot(self, role: str) -> float:
        filled = role_slots_filled(self._roster, role)
        quota = self._slot_quotas[role]
        remaining = quota - filled
        if remaining <= 1:
            return 0.0
        return 1.5 / max(remaining, 1)

    # ------------------------------------------------------------------
    # Opponent simulation
    # ------------------------------------------------------------------

    def _simulate_opponent_decision(self, player: dict, opp_idx: int) -> float:
        budget = self._opponent_budgets[opp_idx]
        if budget <= 0:
            return 0.0

        role = player["role"]
        filled = self._opponent_rosters[opp_idx][role]
        quota = self._slot_quotas[role]
        if filled >= quota:
            return 0.0

        pts = player["projected_points"]
        agg = self._opponent_aggression[opp_idx]
        need_factor = 1.0 + (quota - filled) / max(quota, 1)
        budget_pressure = budget / float(self.config.total_budget)
        noise = self.rng.uniform(0.85, 1.15)

        base_bid = pts * 0.9 * agg * need_factor * noise
        bid = base_bid * budget_pressure

        # Opponents bid more aggressively for scarce roles
        pool_count = self._role_pool_count.get(role, 1)
        slots_needed = quota - filled
        scarcity = slots_needed / max(pool_count, 1)
        bid *= 1.0 + scarcity * 0.5

        bid = max(min(bid, budget), 0.0)
        return bid

    # ------------------------------------------------------------------
    # Terminal conditions
    # ------------------------------------------------------------------

    def _is_done(self) -> bool:
        if self._current_step >= len(self._players):
            return True
        if self._all_slots_filled():
            return True
        if len(self._players) - self._current_step <= 0:
            return True
        for role in self._slot_names:
            remaining_in_pool = 0
            for i in range(self._current_step, len(self._players)):
                if self._players[self._player_order[i]]["role"] == role:
                    remaining_in_pool += 1
            needed = self._slot_quotas[role] - role_slots_filled(self._roster, role)
            if needed > remaining_in_pool:
                return True
        return False

    def _all_slots_filled(self) -> bool:
        for role in self._slot_names:
            filled = role_slots_filled(self._roster, role)
            if filled < self._slot_quotas[role]:
                return False
        return True

    def _done_reason(self) -> str:
        if self._all_slots_filled():
            return "all_slots_filled"
        if self._current_step >= len(self._players):
            return "end_of_players"
        return "auction_complete"

    # ------------------------------------------------------------------
    # Valid actions
    # ------------------------------------------------------------------

    def get_valid_actions(self) -> List[int]:
        if self._current_step >= len(self._players):
            return [0]

        player = self._players[self._player_order[self._current_step]]
        role = player["role"]
        valid = [0]

        if role_slots_filled(self._roster, role) >= self._slot_quotas[role]:
            return valid

        for a in range(1, N_ACTIONS):
            cost = BID_LEVELS[a]
            if cost <= self._budget:
                valid.append(a)

        return valid

    @property
    def state_dim(self) -> int:
        return self._obs_dim

    @property
    def action_dim(self) -> int:
        return N_ACTIONS

    def current_player(self) -> Optional[dict]:
        if self._current_step >= len(self._players):
            return None
        return self._players[self._player_order[self._current_step]]


# ---------------------------------------------------------------------------
# Q-Network (NumPy MLP)
# ---------------------------------------------------------------------------


class QNetwork:
    """Numpy-based MLP for Q-value approximation.

    Architecture: input -> 128 (ReLU) -> 128 (ReLU) -> output (linear).
    Xavier initialization; manual backpropagation with MSE loss.
    """

    def __init__(self, input_dim: int, hidden_dim: int = 128, output_dim: int = 11):
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.output_dim = output_dim

        self.W1 = np.random.randn(input_dim, hidden_dim) * sqrt(2.0 / input_dim)
        self.b1 = np.zeros(hidden_dim, dtype=np.float64)
        self.W2 = np.random.randn(hidden_dim, hidden_dim) * sqrt(2.0 / hidden_dim)
        self.b2 = np.zeros(hidden_dim, dtype=np.float64)
        self.W3 = np.random.randn(hidden_dim, output_dim) * sqrt(2.0 / hidden_dim)
        self.b3 = np.zeros(output_dim, dtype=np.float64)

    def forward(self, x: np.ndarray) -> Tuple[np.ndarray, tuple]:
        if x.ndim == 1:
            x = x.reshape(1, -1)

        z1 = x @ self.W1 + self.b1
        a1 = np.maximum(0, z1)

        z2 = a1 @ self.W2 + self.b2
        a2 = np.maximum(0, z2)

        z3 = a2 @ self.W3 + self.b3

        return z3, (x, z1, a1, z2, a2, z3)

    def predict(self, x: np.ndarray) -> np.ndarray:
        q_vals, _ = self.forward(x)
        if q_vals.shape[0] == 1:
            return q_vals.ravel()
        return q_vals


def _mse_q_loss(
    network: QNetwork,
    states: np.ndarray,
    actions: np.ndarray,
    q_targets: np.ndarray,
) -> float:
    """Compute MSE loss and return gradients for Q-network parameters."""
    batch_size = states.shape[0]

    x, z1, a1, z2, a2, z3 = network.forward(states)[1]

    # Build dL/dz3: only the selected action contributes gradient
    dz3 = np.zeros_like(z3, dtype=np.float64)
    batch_idx = np.arange(batch_size)
    error = z3[batch_idx, actions] - q_targets
    dz3[batch_idx, actions] = error / batch_size

    loss = 0.5 * np.mean(error ** 2)

    # Backprop through W3
    dW3 = a2.T @ dz3
    db3 = np.sum(dz3, axis=0)

    da2 = dz3 @ network.W3.T
    dz2 = da2 * (z2 > 0).astype(np.float64)

    dW2 = a1.T @ dz2
    db2 = np.sum(dz2, axis=0)

    da1 = dz2 @ network.W2.T
    dz1 = da1 * (z1 > 0).astype(np.float64)

    dW1 = x.T @ dz1
    db1 = np.sum(dz1, axis=0)

    # Gradient clipping
    clip_norm = 10.0
    for grad in [dW1, dW2, dW3, db1, db2, db3]:
        norm = np.linalg.norm(grad)
        if norm > clip_norm:
            grad *= clip_norm / norm

    grads = {
        "W1": dW1, "b1": db1,
        "W2": dW2, "b2": db2,
        "W3": dW3, "b3": db3,
    }
    return loss, grads


# ---------------------------------------------------------------------------
# Double DQN Policy
# ---------------------------------------------------------------------------


class RLAuctionPolicy:
    """Double DQN policy for Fantacalcio auction bidding.

    Uses a numpy-based Q-network with target network and experience replay.
    Trains via Double DQN: online network selects max-Q action, target
    network evaluates it.
    """

    def __init__(
        self,
        state_dim: int,
        action_dim: int = 11,
        hidden_dim: int = 128,
        lr: float = 1e-4,
        gamma: float = 0.99,
        buffer_capacity: int = 50000,
    ):
        self.state_dim = state_dim
        self.action_dim = action_dim
        self.gamma = gamma
        self.lr = lr

        self.q_network = QNetwork(state_dim, hidden_dim, action_dim)
        self.target_network = QNetwork(state_dim, hidden_dim, action_dim)
        self._hard_update_target()

        self.replay_buffer = deque(maxlen=buffer_capacity)
        self.rng = np.random.RandomState(42)

        self._obs_mean = np.zeros(state_dim, dtype=np.float64)
        self._obs_std = np.ones(state_dim, dtype=np.float64)
        self._obs_count = 0
        self._trained_steps = 0

    def _hard_update_target(self):
        for src_name, tgt_name in [
            ("W1", "W1"),
            ("b1", "b1"),
            ("W2", "W2"),
            ("b2", "b2"),
            ("W3", "W3"),
            ("b3", "b3"),
        ]:
            src = getattr(self.q_network, src_name)
            setattr(self.target_network, tgt_name, src.copy())

    def _resize_networks(self, new_state_dim: int):
        """Reinitialize networks when state dimension changes."""
        self.q_network = QNetwork(new_state_dim, self.q_network.hidden_dim, self.action_dim)
        self.target_network = QNetwork(new_state_dim, self.target_network.hidden_dim, self.action_dim)
        self._hard_update_target()

    def _normalize_state(self, state: np.ndarray) -> np.ndarray:
        state = np.asarray(state, dtype=np.float64).ravel()
        if len(state) != len(self._obs_mean):
            self._obs_mean = np.zeros(len(state), dtype=np.float64)
            self._obs_std = np.ones(len(state), dtype=np.float64)
            self._obs_count = 0
            self.state_dim = len(state)
            self._resize_networks(len(state))
        self._obs_count += 1
        n = self._obs_count
        old_mean = self._obs_mean.copy()
        self._obs_mean = old_mean + (state - old_mean) / n
        if n > 1:
            self._obs_std = np.sqrt(
                (self._obs_std ** 2 * (n - 2) + (state - old_mean) * (state - self._obs_mean)) / max(n - 1, 1)
            )
        self._obs_std = np.maximum(self._obs_std, 1e-6)
        return (state - self._obs_mean) / self._obs_std

    # ------------------------------------------------------------------
    # Action selection
    # ------------------------------------------------------------------

    def act(self, state: np.ndarray, epsilon: float = 0.05) -> int:
        if self.rng.random() < epsilon:
            return self.rng.randint(0, self.action_dim)
        s = self._normalize_state(state).reshape(1, -1)
        q_vals = self.q_network.predict(s)
        return int(np.argmax(q_vals))

    def act_greedy(self, state: np.ndarray) -> int:
        s = self._normalize_state(state).reshape(1, -1)
        q_vals = self.q_network.predict(s)
        return int(np.argmax(q_vals))

    # ------------------------------------------------------------------
    # Replay & learning
    # ------------------------------------------------------------------

    def remember(
        self,
        state: np.ndarray,
        action: int,
        reward: float,
        next_state: np.ndarray,
        done: bool,
    ):
        self.replay_buffer.append((
            self._normalize_state(state),
            action,
            reward,
            self._normalize_state(next_state),
            done,
        ))

    def replay(self, batch_size: int = 64) -> float:
        if len(self.replay_buffer) < batch_size:
            return 0.0

        indices = self.rng.choice(len(self.replay_buffer), size=batch_size, replace=False)
        batch = [self.replay_buffer[i] for i in indices]

        s_batch = np.vstack([b[0] for b in batch])
        a_batch = np.array([b[1] for b in batch], dtype=int)
        r_batch = np.array([b[2] for b in batch], dtype=np.float64)
        s_next_batch = np.vstack([b[3] for b in batch])
        done_batch = np.array([b[4] for b in batch], dtype=float)

        # Double DQN targets
        q_next_online = self.q_network.predict(s_next_batch)
        q_next_target = self.target_network.predict(s_next_batch)
        best_actions = np.argmax(q_next_online, axis=1)
        max_next_q = q_next_target[np.arange(batch_size), best_actions]
        q_targets = r_batch + self.gamma * max_next_q * (1.0 - done_batch)

        loss, grads = _mse_q_loss(self.q_network, s_batch, a_batch, q_targets)

        # Gradient descent
        for param_name in ["W1", "b1", "W2", "b2", "W3", "b3"]:
            param = getattr(self.q_network, param_name)
            param -= self.lr * grads[param_name]

        self._trained_steps += 1
        return float(loss)

    def update_target_network(self, tau: float = 0.005):
        for param_name in ["W1", "b1", "W2", "b2", "W3", "b3"]:
            online = getattr(self.q_network, param_name)
            target = getattr(self.target_network, param_name)
            setattr(self.target_network, param_name, tau * online + (1.0 - tau) * target)

    def update_target_hard(self):
        self._hard_update_target()

    # ------------------------------------------------------------------
    # Training loop
    # ------------------------------------------------------------------

    def train(
        self,
        env: AuctionEnv,
        episodes: int = 5000,
        eval_interval: int = 100,
        verbose: bool = True,
    ) -> dict:
        """Full Double DQN training loop.

        Args:
            env: AuctionEnv instance.
            episodes: total training episodes.
            eval_interval: evaluate every N episodes.
            verbose: log progress.

        Returns:
            dict with training history: episode_rewards, eval_rewards, losses.
        """
        history = {
            "episode_rewards": [],
            "eval_avg_rewards": [],
            "losses": [],
            "best_eval_reward": -float("inf"),
            "best_episode": 0,
        }

        epsilon_start = 1.0
        epsilon_end = 0.02
        epsilon_decay_episodes = max(1000, episodes // 5)

        for ep in range(episodes):
            state = env.reset()
            done = False
            ep_reward = 0.0
            ep_loss = 0.0
            steps = 0

            # Epsilon decay
            epsilon = epsilon_end + (epsilon_start - epsilon_end) * max(
                0, 1.0 - ep / epsilon_decay_episodes
            )

            while not done:
                valid = env.get_valid_actions()
                if self.rng.random() < epsilon:
                    action = int(self.rng.choice(valid))
                else:
                    s = self._normalize_state(state).reshape(1, -1)
                    q_vals = self.q_network.predict(s)
                    action_mask = np.ones(self.action_dim, dtype=bool)
                    for a in valid:
                        action_mask[a] = True
                    # Select best among valid
                    mask = np.full(self.action_dim, -np.inf)
                    for a in valid:
                        mask[a] = q_vals[a]
                    action = int(np.argmax(mask))

                next_state, reward, done, info = step_in_env(env, action)
                self.remember(state, action, reward, next_state, done)
                loss = self.replay(batch_size=64)
                self.update_target_network(tau=0.005)

                ep_reward += reward
                ep_loss += loss
                state = next_state
                steps += 1

            history["episode_rewards"].append(ep_reward)
            history["losses"].append(ep_loss / max(steps, 1))

            # Periodic evaluation
            if (ep + 1) % eval_interval == 0:
                eval_reward = self._evaluate(env, n_episodes=20)
                history["eval_avg_rewards"].append(eval_reward)

                if eval_reward > history["best_eval_reward"]:
                    history["best_eval_reward"] = eval_reward
                    history["best_episode"] = ep + 1

                if verbose:
                    logger.info(
                        f"Ep {ep + 1:5d} | eps={epsilon:.3f} | "
                        f"train_reward={ep_reward:7.1f} | "
                        f"eval_reward={eval_reward:7.1f} | "
                        f"best={history['best_eval_reward']:7.1f}"
                    )

        logger.info(
            f"Training complete. Best eval reward: {history['best_eval_reward']:.1f} "
            f"at episode {history['best_episode']}"
        )
        return history

    def _evaluate(self, env: AuctionEnv, n_episodes: int = 20) -> float:
        rewards = []
        for _ in range(n_episodes):
            state = env.reset()
            done = False
            total = 0.0
            while not done:
                valid = env.get_valid_actions()
                s = self._normalize_state(state).reshape(1, -1)
                q_vals = self.q_network.predict(s)
                mask = np.full(self.action_dim, -np.inf)
                for a in valid:
                    mask[a] = q_vals[a]
                action = int(np.argmax(mask))
                next_state, reward, done, _ = step_in_env(env, action)
                total += reward
                state = next_state
            rewards.append(total)
        return float(np.mean(rewards))

    # ------------------------------------------------------------------
    # Inference for live use
    # ------------------------------------------------------------------

    def bid(self, player_features: np.ndarray, auction_state: np.ndarray) -> Tuple[int, float]:
        """Return (action, bid_amount) for live auction use.

        Args:
            player_features: array with player-specific feature values.
            auction_state: full observation vector.

        Returns:
            (action_index, bid_amount_in_credits)
        """
        obs = np.concatenate([auction_state, player_features])
        action = self.act_greedy(obs)
        bid_amount = BID_LEVELS[action]
        return action, bid_amount

    # ------------------------------------------------------------------
    # Persistence
    # ------------------------------------------------------------------

    def save(self, path: str):
        p = Path(path)
        p.parent.mkdir(parents=True, exist_ok=True)
        state = {
            "q_network": {
                "W1": self.q_network.W1,
                "b1": self.q_network.b1,
                "W2": self.q_network.W2,
                "b2": self.q_network.b2,
                "W3": self.q_network.W3,
                "b3": self.q_network.b3,
                "input_dim": self.q_network.input_dim,
                "hidden_dim": self.q_network.hidden_dim,
                "output_dim": self.q_network.output_dim,
            },
            "target_network": {
                "W1": self.target_network.W1,
                "b1": self.target_network.b1,
                "W2": self.target_network.W2,
                "b2": self.target_network.b2,
                "W3": self.target_network.W3,
                "b3": self.target_network.b3,
            },
            "obs_mean": self._obs_mean,
            "obs_std": self._obs_std,
            "obs_count": self._obs_count,
            "trained_steps": self._trained_steps,
            "state_dim": self.state_dim,
            "action_dim": self.action_dim,
            "gamma": self.gamma,
            "lr": self.lr,
        }
        with open(p, "wb") as f:
            pickle.dump(state, f)
        logger.info(f"Policy saved to {path}")

    @classmethod
    def load(cls, path: str) -> "RLAuctionPolicy":
        with open(path, "rb") as f:
            state = pickle.load(f)

        inst = cls(
            state_dim=state["state_dim"],
            action_dim=state["action_dim"],
            gamma=state["gamma"],
            lr=state["lr"],
        )

        qn = state["q_network"]
        inst.q_network = QNetwork(qn["input_dim"], qn["hidden_dim"], qn["output_dim"])
        inst.q_network.W1 = qn["W1"]
        inst.q_network.b1 = qn["b1"]
        inst.q_network.W2 = qn["W2"]
        inst.q_network.b2 = qn["b2"]
        inst.q_network.W3 = qn["W3"]
        inst.q_network.b3 = qn["b3"]

        tn = state["target_network"]
        inst.target_network = QNetwork(qn["input_dim"], qn["hidden_dim"], qn["output_dim"])
        inst.target_network.W1 = tn["W1"]
        inst.target_network.b1 = tn["b1"]
        inst.target_network.W2 = tn["W2"]
        inst.target_network.b2 = tn["b2"]
        inst.target_network.W3 = tn["W3"]
        inst.target_network.b3 = tn["b3"]

        inst._obs_mean = state["obs_mean"]
        inst._obs_std = state["obs_std"]
        inst._obs_count = state["obs_count"]
        inst._trained_steps = state["trained_steps"]

        logger.info(f"Policy loaded from {path} ({inst._trained_steps} training steps)")
        return inst


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def role_slots_filled(roster: dict, role: str) -> int:
    return len(roster.get(role, []))


def step_in_env(env: AuctionEnv, action: int) -> Tuple[np.ndarray, float, bool, dict]:
    """Safe step wrapper that falls back to action 0 on invalid action."""
    valid = env.get_valid_actions()
    if action not in valid:
        action = 0
    return env.step(action)


# ---------------------------------------------------------------------------
# Trainer helper
# ---------------------------------------------------------------------------


class RLAuctionTrainer:
    """Convenience class for training and evaluating the RL auction agent."""

    def __init__(
        self,
        player_pool: Optional[pd.DataFrame] = None,
        n_opponents: int = 7,
        config: Optional[AuctionConfig] = None,
        model_dir: str = "models_trained",
    ):
        self.player_pool = player_pool
        self.n_opponents = n_opponents
        self.config = config or AuctionConfig()
        self.model_dir = Path(model_dir)
        self.model_dir.mkdir(parents=True, exist_ok=True)

    def train(
        self,
        n_episodes: int = 5000,
        verbose: bool = True,
    ) -> RLAuctionPolicy:
        """Train RL policy on stored player pool."""
        if self.player_pool is None:
            raise ValueError("No player_pool provided to trainer")
        env = self.prepare_training_data(self.player_pool, n_opponents=self.n_opponents, config=self.config)
        policy, _ = self.train_agent(env, episodes=n_episodes, eval_interval=100)
        return policy

    def prepare_training_data(
        self,
        player_pool_df: pd.DataFrame,
        n_opponents: int = 7,
        config: Optional[AuctionConfig] = None,
    ) -> AuctionEnv:
        return AuctionEnv(player_pool_df, n_opponents=n_opponents, config=config)

    def train_agent(
        self,
        env: AuctionEnv,
        episodes: int = 5000,
        eval_interval: int = 100,
        hidden_dim: int = 128,
        lr: float = 1e-4,
        gamma: float = 0.99,
    ) -> Tuple[RLAuctionPolicy, dict]:
        policy = RLAuctionPolicy(
            state_dim=env.state_dim,
            action_dim=env.action_dim,
            hidden_dim=hidden_dim,
            lr=lr,
            gamma=gamma,
        )
        history = policy.train(env, episodes=episodes, eval_interval=eval_interval)
        return policy, history

    def evaluate_vs_baselines(
        self,
        agent: RLAuctionPolicy,
        player_pool: pd.DataFrame,
        n_sims: int = 100,
        n_opponents: int = 7,
        config: Optional[AuctionConfig] = None,
    ) -> dict:
        """Compare RL agent against greedy and MILP baselines.

        Runs n_sims independent auction simulations with each strategy
        and reports average total value, cost, and efficiency.

        Returns:
            dict with per-strategy metrics.
        """
        from .auction_solver import AuctionSolver

        results: Dict[str, List[dict]] = {
            "rl_agent": [],
            "greedy_baseline": [],
            "milp_baseline": [],
        }

        for sim in range(n_sims):
            # RL agent
            env = AuctionEnv(player_pool.copy(), n_opponents=n_opponents, config=config)
            state = env.reset()
            done = False
            total_value_rl = 0.0
            total_cost_rl = 0.0
            while not done:
                obs = agent._normalize_state(state).reshape(1, -1)
                valid = env.get_valid_actions()
                q_vals = agent.q_network.predict(obs)
                mask = np.full(env.action_dim, -np.inf)
                for a in valid:
                    mask[a] = q_vals[a]
                action = int(np.argmax(mask))
                state, reward, done, info = step_in_env(env, action)
                if info.get("outcome") == "won":
                    total_value_rl += info.get("player_points", 0)
                    total_cost_rl += info.get("price_paid", 0)
            results["rl_agent"].append({
                "total_value": total_value_rl,
                "total_cost": total_cost_rl,
                "efficiency": total_value_rl / max(total_cost_rl, 1),
            })

            # Greedy baseline
            env2 = AuctionEnv(player_pool.copy(), n_opponents=n_opponents, config=config)
            results["greedy_baseline"].append(
                self._run_greedy_sim(env2)
            )

            # MILP baseline (on fresh env without opponents)
            try:
                solver = AuctionSolver(config=config)
                solver.add_players(player_pool)
                sol = solver.solve()
                results["milp_baseline"].append({
                    "total_value": sol["total_value"],
                    "total_cost": sol["total_cost"],
                    "efficiency": sol["total_value"] / max(sol["total_cost"], 1),
                    "status": sol["status"],
                })
            except Exception as e:
                logger.warning(f"MILP baseline failed in sim {sim}: {e}")
                results["milp_baseline"].append({
                    "total_value": 0, "total_cost": 0, "efficiency": 0, "status": "error",
                })

        # Aggregate
        summary = {}
        for strategy, sims in results.items():
            if not sims:
                continue
            values = [s["total_value"] for s in sims]
            costs = [s["total_cost"] for s in sims]
            effs = [s["efficiency"] for s in sims]
            summary[strategy] = {
                "avg_value": float(np.mean(values)),
                "std_value": float(np.std(values)),
                "avg_cost": float(np.mean(costs)),
                "avg_efficiency": float(np.mean(effs)),
                "n_successful": len(values),
            }

        for strategy, metrics in list(summary.items()):
            summary[f"{strategy}_total_value"] = metrics["avg_value"]

        summary["rl_total_value"] = summary.get("rl_agent_total_value", 0)
        summary["greedy_total_value"] = summary.get("greedy_baseline_total_value", 0)

        logger.info(
            f"Benchmark complete: RL={summary.get('rl_agent', {}).get('avg_value', 0):.1f} pts "
            f"vs Greedy={summary.get('greedy_baseline', {}).get('avg_value', 0):.1f} "
            f"vs MILP={summary.get('milp_baseline', {}).get('avg_value', 0):.1f}"
        )
        return summary

    def _run_greedy_sim(self, env: AuctionEnv) -> dict:
        """Run a greedy-bidding simulation against opponents."""
        state = env.reset()
        done = False
        total_value = 0.0
        total_cost = 0.0

        while not done:
            player = env.current_player()
            if player is None:
                break
            role = player["role"]
            filled = role_slots_filled(env._roster, role)
            quota = env._slot_quotas[role]
            if filled >= quota:
                action = 0
            else:
                pts = player["projected_points"]
                vor = env._compute_vor(player)
                bid = min(pts * 0.6, env._budget * 0.5)
                action = max(
                    range(1, N_ACTIONS),
                    key=lambda a: 1.0 if BID_LEVELS[a] <= bid else -1.0,
                )
                if BID_LEVELS[action] > bid and action > 1:
                    action -= 1

            state, reward, done, info = step_in_env(env, action)
            if info.get("outcome") == "won":
                total_value += info.get("player_points", 0)
                total_cost += info.get("price_paid", 0)

        return {
            "total_value": total_value,
            "total_cost": total_cost,
            "efficiency": total_value / max(total_cost, 1),
        }