"""
Proximal Policy Optimization (PPO) for Essence Wars.

This module provides a complete PPO implementation with:
- Action masking for illegal action handling
- Mixed opponent strategy (self, greedy, random)
- Vectorized environment support for high throughput
- TensorBoard logging

Based on CleanRL's PPO implementation with modifications for card games.
"""

from __future__ import annotations

import time
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Literal

import numpy as np
import numpy.typing as npt
import torch
import torch.nn as nn
import torch.optim as optim

from essence_wars._core import STATE_TENSOR_SIZE
from essence_wars.agents.embeddings import create_network
from essence_wars.agents.utils import RunningMeanStd
from essence_wars.env import (
    EssenceWarsEnv,
    VectorizedEssenceWars,
    VectorizedEssenceWarsWithShaping,
)

if TYPE_CHECKING:
    from collections.abc import Callable

    from essence_wars.agents.networks import EssenceWarsNetwork

# Type alias for numpy arrays
NDArrayFloat = npt.NDArray[np.floating[Any]]
NDArrayBool = npt.NDArray[np.bool_]


# Faction to deck mapping
FACTION_DECKS: dict[str, list[str]] = {
    "argentum": ["architect_fortify", "artificer_tokens", "sanctum_healer", "vex_piercing"],
    "obsidion": ["archon_burst", "deathmaster_assassin", "shadow_weaver", "sovereign_lifesteal"],
    "symbiote": ["alpha_frenzy", "broodmother_pack", "grove_regenerate", "plague_volatile"],
}

ALL_DECKS: list[str] = [deck for decks in FACTION_DECKS.values() for deck in decks]

# Playstyle to deck mapping (aligned with Rust bot tuning system)
PLAYSTYLE_DECKS: dict[str, list[str]] = {
    "aggro": ["vex_piercing", "archon_burst", "shadow_weaver", "alpha_frenzy", "broodmother_pack"],
    "control": ["architect_fortify", "sanctum_healer", "plague_volatile"],
    "tempo": ["artificer_tokens", "deathmaster_assassin", "sovereign_lifesteal"],
    "midrange": ["grove_regenerate"],
}
ALL_PLAYSTYLES: list[str] = list(PLAYSTYLE_DECKS.keys())


@dataclass
class PPOConfig:
    """Configuration for PPO training."""

    # Environment
    num_envs: int = 64
    max_episode_steps: int = 500

    # Deck selection
    # Priority order: playstyle > faction > deck > default
    # If player_playstyle is set, cycles through that playstyle's decks
    # If player_faction is set, cycles through faction's decks during training
    # If player_deck is set, uses that specific deck
    # If none set, uses default decks
    player_playstyle: str | None = None  # "aggro", "control", "tempo", or "midrange"
    player_faction: str | None = None  # "argentum", "obsidion", or "symbiote"
    player_deck: str | None = None     # Specific deck name
    opponent_playstyle: str | None = None  # If set, opponent uses this playstyle's decks
    opponent_faction: str | None = None  # If set, opponent uses random faction deck
    opponent_deck: str | None = None   # Specific opponent deck (default: random from all)
    deck_cycle_interval: int = 25_000  # Steps between deck changes (if using playstyle/faction)

    # Training
    total_timesteps: int = 1_000_000
    learning_rate: float = 3e-4
    num_steps: int = 128  # Steps per rollout per env
    num_epochs: int = 4  # PPO epochs per update
    num_minibatches: int = 4
    gamma: float = 0.99
    gae_lambda: float = 0.95

    # PPO specific
    clip_coef: float = 0.2
    clip_vloss: bool = True
    ent_coef: float = 0.02  # Increased from 0.01 to reduce policy collapse
    vf_coef: float = 0.5
    max_grad_norm: float = 0.5
    target_kl: float | None = None  # Early stopping KL threshold

    # Early stopping and best checkpoint
    save_best: bool = True  # Save checkpoint when eval improves
    early_stopping_patience: int | None = None  # Stop if no improvement for N evals (None = disabled)

    # Network
    hidden_dim: int = 256

    # Observation mode for card embeddings
    # "flat" = original 328-float tensor
    # "embedded" = learned card embeddings (end-to-end)
    # "embedded_pretrained" = pre-trained card embeddings
    observation_mode: str = "flat"
    embed_dim: int = 64  # Card embedding dimension
    pretrained_embeds_path: str | None = None  # Path to pre-trained embeddings
    freeze_embeds: bool = False  # Freeze embeddings during training
    include_embed_section: bool = False  # Include trailing card ID section

    # Observation normalization
    normalize_obs: bool = True  # Use running mean/std normalization

    # Reward shaping (dense intermediate rewards)
    use_reward_shaping: bool = False  # Enable dense reward shaping
    shaping_scale: float = 0.01  # Scale factor for shaped rewards
    life_weight: float = 1.0  # Weight for life differential reward
    board_weight: float = 0.5  # Weight for board control reward

    # Opponent strategy (for single-agent training)
    # Options: "greedy", "random", "mixed", "mcts<N>" (e.g. "mcts25"), "alphabeta<N>" (e.g. "alphabeta3")
    opponent_type: str = "greedy"

    # Evaluation
    eval_interval: int = 10_000
    eval_episodes: int = 100

    # Logging
    log_interval: int = 1000
    save_interval: int = 50_000

    # Device
    device: str = "cuda" if torch.cuda.is_available() else "cpu"

    @property
    def batch_size(self) -> int:
        return self.num_envs * self.num_steps

    @property
    def minibatch_size(self) -> int:
        return self.batch_size // self.num_minibatches


class RolloutBuffer:
    """Buffer for storing rollout experience."""

    def __init__(
        self,
        num_steps: int,
        num_envs: int,
        obs_dim: int,
        action_dim: int,
        device: str,
    ) -> None:
        self.num_steps = num_steps
        self.num_envs = num_envs
        self.device = device

        # Pre-allocate tensors
        self.obs = torch.zeros((num_steps, num_envs, obs_dim), device=device)
        self.actions = torch.zeros((num_steps, num_envs), dtype=torch.long, device=device)
        self.action_masks = torch.zeros((num_steps, num_envs, action_dim), dtype=torch.bool, device=device)
        self.log_probs = torch.zeros((num_steps, num_envs), device=device)
        self.rewards = torch.zeros((num_steps, num_envs), device=device)
        self.dones = torch.zeros((num_steps, num_envs), device=device)
        self.values = torch.zeros((num_steps, num_envs), device=device)

        # Computed during finalization
        self.advantages = torch.zeros((num_steps, num_envs), device=device)
        self.returns = torch.zeros((num_steps, num_envs), device=device)

        self.step = 0

    def add(
        self,
        obs: torch.Tensor,
        action: torch.Tensor,
        action_mask: torch.Tensor,
        log_prob: torch.Tensor,
        reward: torch.Tensor,
        done: torch.Tensor,
        value: torch.Tensor,
    ) -> None:
        """Add a step to the buffer."""
        self.obs[self.step] = obs
        self.actions[self.step] = action
        self.action_masks[self.step] = action_mask
        self.log_probs[self.step] = log_prob
        self.rewards[self.step] = reward
        self.dones[self.step] = done
        self.values[self.step] = value
        self.step += 1

    def compute_returns_and_advantages(
        self,
        last_value: torch.Tensor,
        gamma: float,
        gae_lambda: float,
    ) -> None:
        """Compute GAE advantages and returns."""
        last_gae: torch.Tensor | float = 0.0
        for t in reversed(range(self.num_steps)):
            if t == self.num_steps - 1:
                next_non_terminal = 1.0 - self.dones[t]
                next_value = last_value
            else:
                next_non_terminal = 1.0 - self.dones[t]
                next_value = self.values[t + 1]

            delta = self.rewards[t] + gamma * next_value * next_non_terminal - self.values[t]
            self.advantages[t] = last_gae = delta + gamma * gae_lambda * next_non_terminal * last_gae

        self.returns = self.advantages + self.values

    def get_batches(self, minibatch_size: int):
        """Yield minibatches of experience."""
        batch_size = self.num_steps * self.num_envs
        indices = np.random.permutation(batch_size)

        # Flatten tensors
        obs = self.obs.reshape(-1, self.obs.shape[-1])
        actions = self.actions.reshape(-1)
        action_masks = self.action_masks.reshape(-1, self.action_masks.shape[-1])
        log_probs = self.log_probs.reshape(-1)
        advantages = self.advantages.reshape(-1)
        returns = self.returns.reshape(-1)
        values = self.values.reshape(-1)

        for start in range(0, batch_size, minibatch_size):
            end = start + minibatch_size
            batch_indices = indices[start:end]

            yield (
                obs[batch_indices],
                actions[batch_indices],
                action_masks[batch_indices],
                log_probs[batch_indices],
                advantages[batch_indices],
                returns[batch_indices],
                values[batch_indices],
            )

    def reset(self) -> None:
        """Reset buffer for next rollout."""
        self.step = 0


class PPOTrainer:
    """
    PPO trainer for Essence Wars.

    Trains a policy network using PPO with action masking and
    mixed opponent strategy.

    Args:
        config: PPO configuration
        network: Optional pre-built network
        writer: Optional TensorBoard SummaryWriter

    Example:
        trainer = PPOTrainer(PPOConfig(num_envs=64))
        trainer.train(total_timesteps=1_000_000)

        # Save trained model
        trainer.save("checkpoints/ppo_model.pt")

        # Evaluate
        win_rate = trainer.evaluate_vs_greedy(100)
    """

    def __init__(
        self,
        config: PPOConfig | None = None,
        network: EssenceWarsNetwork | None = None,
        writer=None,
    ) -> None:
        self.config = config or PPOConfig()
        self.device = torch.device(self.config.device)

        # Create network
        if network is not None:
            self.network: EssenceWarsNetwork = network.to(self.device)
        else:
            observation_mode = self.config.observation_mode
            # Cast to Literal type for create_network
            assert observation_mode in ("flat", "embedded", "embedded_pretrained")
            obs_mode: Literal["flat", "embedded", "embedded_pretrained"] = observation_mode  # type: ignore[assignment]
            # create_network returns EssenceWarsNetwork (PPONetwork/EmbeddedPPONetwork)
            created_network: EssenceWarsNetwork = create_network(  # type: ignore[assignment]
                observation_mode=obs_mode,
                network_type="ppo",
                embed_dim=self.config.embed_dim,
                hidden_dim=self.config.hidden_dim,
                pretrained_path=self.config.pretrained_embeds_path,
                freeze_embeds=self.config.freeze_embeds,
                include_embed_section=self.config.include_embed_section,
            )
            self.network = created_network.to(self.device)

        # Optimizer
        self.optimizer = optim.Adam(
            self.network.parameters(),
            lr=self.config.learning_rate,
            eps=1e-5,
        )

        # Deck selection for faction specialist training
        self._setup_deck_cycling()

        # Vectorized environment
        self.envs = self._create_envs()

        # Rollout buffer
        device_str = str(self.device)
        self.buffer = RolloutBuffer(
            num_steps=self.config.num_steps,
            num_envs=self.config.num_envs,
            obs_dim=STATE_TENSOR_SIZE,
            action_dim=422,
            device=device_str,
        )

        # TensorBoard writer
        self.writer = writer

        # Observation normalizer
        self.obs_normalizer = RunningMeanStd((STATE_TENSOR_SIZE,)) if self.config.normalize_obs else None

        # Training state
        self.global_step = 0
        self.start_time: float | None = None
        self._last_obs: NDArrayFloat | None = None
        self._last_masks: NDArrayBool | None = None

        # Episode tracking
        self.episode_rewards: list[float] = []
        self.episode_lengths: list[int] = []

        # Best checkpoint tracking
        self.best_win_rate = 0.0
        self.best_checkpoint_path: str | None = None
        self.evals_without_improvement = 0

    def _setup_deck_cycling(self) -> None:
        """Setup deck lists for playstyle/faction specialist training.

        Priority order: playstyle > faction > deck > default
        """
        import random as py_random

        # Determine player decks (priority: playstyle > faction > deck > default)
        if self.config.player_playstyle:
            if self.config.player_playstyle not in PLAYSTYLE_DECKS:
                raise ValueError(f"Unknown playstyle: {self.config.player_playstyle}. "
                                f"Choose from: {ALL_PLAYSTYLES}")
            self._player_decks = PLAYSTYLE_DECKS[self.config.player_playstyle].copy()
            if len(self._player_decks) == 1:
                print(f"  [Info] Playstyle '{self.config.player_playstyle}' has only 1 deck "
                      f"({self._player_decks[0]}), no player deck cycling will occur")
        elif self.config.player_faction:
            if self.config.player_faction not in FACTION_DECKS:
                raise ValueError(f"Unknown faction: {self.config.player_faction}. "
                                f"Choose from: {list(FACTION_DECKS.keys())}")
            self._player_decks = FACTION_DECKS[self.config.player_faction].copy()
        elif self.config.player_deck:
            self._player_decks = [self.config.player_deck]
        else:
            self._player_decks = ALL_DECKS.copy()  # Generalist: train on all decks

        # Determine opponent decks (priority: opponent_playstyle > opponent_faction > opponent_deck > inferred)
        if self.config.opponent_playstyle:
            if self.config.opponent_playstyle not in PLAYSTYLE_DECKS:
                raise ValueError(f"Unknown playstyle: {self.config.opponent_playstyle}. "
                                f"Choose from: {ALL_PLAYSTYLES}")
            self._opponent_decks = PLAYSTYLE_DECKS[self.config.opponent_playstyle].copy()
        elif self.config.opponent_faction:
            if self.config.opponent_faction not in FACTION_DECKS:
                raise ValueError(f"Unknown faction: {self.config.opponent_faction}")
            self._opponent_decks = FACTION_DECKS[self.config.opponent_faction].copy()
        elif self.config.opponent_deck:
            self._opponent_decks = [self.config.opponent_deck]
        elif self.config.player_playstyle:
            # Playstyle specialist: opponent uses OTHER playstyles' decks
            self._opponent_decks = [
                deck for playstyle, decks in PLAYSTYLE_DECKS.items()
                if playstyle != self.config.player_playstyle
                for deck in decks
            ]
        elif self.config.player_faction:
            # Faction specialist: opponent uses OTHER factions only (no mirror matches)
            self._opponent_decks = [
                deck for faction, decks in FACTION_DECKS.items()
                if faction != self.config.player_faction
                for deck in decks
            ]
        else:
            self._opponent_decks = ALL_DECKS.copy()  # All decks for opponent variety

        # Shuffle for variety
        py_random.shuffle(self._player_decks)
        py_random.shuffle(self._opponent_decks)

        # Current deck indices
        self._player_deck_idx = 0
        self._opponent_deck_idx = 0
        self._last_deck_cycle_step = 0

    def _get_current_decks(self) -> tuple[str, str]:
        """Get current deck names for player and opponent."""
        deck1 = self._player_decks[self._player_deck_idx % len(self._player_decks)]
        deck2 = self._opponent_decks[self._opponent_deck_idx % len(self._opponent_decks)]
        return deck1, deck2

    def _create_envs(self) -> VectorizedEssenceWars:
        """Create vectorized environments with current deck configuration."""
        deck1, deck2 = self._get_current_decks()

        if self.config.use_reward_shaping:
            return VectorizedEssenceWarsWithShaping(
                num_envs=self.config.num_envs,
                deck1=deck1,
                deck2=deck2,
                opponent=self.config.opponent_type,
                shaping_scale=self.config.shaping_scale,
                life_weight=self.config.life_weight,
                board_weight=self.config.board_weight,
            )
        else:
            return VectorizedEssenceWars(
                num_envs=self.config.num_envs,
                deck1=deck1,
                deck2=deck2,
                opponent=self.config.opponent_type,
            )

    def _maybe_cycle_decks(self) -> None:
        """Cycle to next deck pair if interval has elapsed."""
        # Skip cycling if both player and opponent have only one deck
        if len(self._player_decks) <= 1 and len(self._opponent_decks) <= 1:
            return

        steps_since_cycle = self.global_step - self._last_deck_cycle_step
        if steps_since_cycle >= self.config.deck_cycle_interval:
            # Move to next deck combination
            self._player_deck_idx = (self._player_deck_idx + 1) % len(self._player_decks)
            self._opponent_deck_idx = (self._opponent_deck_idx + 1) % len(self._opponent_decks)
            self._last_deck_cycle_step = self.global_step

            # Recreate environments with new decks
            _old_deck1, _ = self._get_current_decks()
            deck1, deck2 = self._get_current_decks()
            print(f"  [Deck cycle] Now training with: {deck1} vs {deck2}")
            self.envs = self._create_envs()

            # Force reset on next rollout
            self._last_obs = None
            self._last_masks = None

    def collect_rollout(self) -> dict[str, Any]:
        """Collect rollout experience from vectorized environments."""
        self.network.eval()
        self.buffer.reset()

        # Reset if first step or after deck cycle (when _last_obs is None)
        if self.global_step == 0 or self._last_obs is None or self._last_masks is None:
            obs, masks = self.envs.reset()
        else:
            obs, masks = self._last_obs, self._last_masks

        episode_infos = []

        for _step in range(self.config.num_steps):
            # Normalize observations if enabled
            if self.obs_normalizer is not None:
                self.obs_normalizer.update(obs)
                obs_normalized = self.obs_normalizer.normalize(obs).astype(np.float32)
            else:
                obs_normalized = obs

            # Convert to tensors
            obs_t = torch.tensor(obs_normalized, dtype=torch.float32, device=self.device)
            masks_t = torch.tensor(masks, dtype=torch.bool, device=self.device)

            # Get action from policy
            with torch.no_grad():
                action, log_prob, _ = self.network.get_action(obs_t, masks_t)
                value = self.network.get_value(obs_t)

            # Step environment
            next_obs, rewards, dones, next_masks = self.envs.step(action.cpu().numpy())

            # Track episodes that finished
            for i, done in enumerate(dones):
                if done:
                    episode_infos.append({
                        "reward": self.envs.get_episode_rewards()[i],
                        "length": self.envs.get_episode_lengths()[i],
                    })

            # Store in buffer
            self.buffer.add(
                obs=obs_t,
                action=action,
                action_mask=masks_t,
                log_prob=log_prob,
                reward=torch.tensor(rewards, dtype=torch.float32, device=self.device),
                done=torch.tensor(dones, dtype=torch.float32, device=self.device),
                value=value,
            )

            obs, masks = next_obs, next_masks
            self.global_step += self.config.num_envs

        # Store last observation for next rollout
        self._last_obs = obs
        self._last_masks = masks

        # Compute final value for GAE
        with torch.no_grad():
            if self.obs_normalizer is not None:
                obs_normalized = self.obs_normalizer.normalize(obs).astype(np.float32)
            else:
                obs_normalized = obs
            obs_t = torch.tensor(obs_normalized, dtype=torch.float32, device=self.device)
            last_value = self.network.get_value(obs_t)

        # Compute advantages
        self.buffer.compute_returns_and_advantages(
            last_value=last_value,
            gamma=self.config.gamma,
            gae_lambda=self.config.gae_lambda,
        )

        return {
            "episode_infos": episode_infos,
        }

    def update(self) -> dict[str, float]:
        """Perform PPO update on collected rollout."""
        self.network.train()

        # Normalize advantages
        advantages = self.buffer.advantages.reshape(-1)
        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
        self.buffer.advantages = advantages.reshape(self.config.num_steps, self.config.num_envs)

        # Track metrics
        pg_losses = []
        value_losses = []
        entropy_losses = []
        approx_kls = []
        clipfracs = []

        for _epoch in range(self.config.num_epochs):
            for batch in self.buffer.get_batches(self.config.minibatch_size):
                (
                    obs_batch,
                    actions_batch,
                    masks_batch,
                    old_log_probs_batch,
                    advantages_batch,
                    returns_batch,
                    old_values_batch,
                ) = batch

                # Get new log probs, entropy, and values
                new_log_probs, entropy, new_values = self.network.evaluate_actions(
                    obs_batch, masks_batch, actions_batch
                )

                # Compute ratio
                log_ratio = new_log_probs - old_log_probs_batch
                ratio = log_ratio.exp()

                # Approximate KL divergence
                with torch.no_grad():
                    approx_kl = ((ratio - 1) - log_ratio).mean()
                    approx_kls.append(approx_kl.item())
                    clipfracs.append(((ratio - 1.0).abs() > self.config.clip_coef).float().mean().item())

                # Policy loss
                pg_loss1 = -advantages_batch * ratio
                pg_loss2 = -advantages_batch * torch.clamp(
                    ratio, 1 - self.config.clip_coef, 1 + self.config.clip_coef
                )
                pg_loss = torch.max(pg_loss1, pg_loss2).mean()

                # Value loss
                if self.config.clip_vloss:
                    v_loss_unclipped = (new_values - returns_batch) ** 2
                    v_clipped = old_values_batch + torch.clamp(
                        new_values - old_values_batch,
                        -self.config.clip_coef,
                        self.config.clip_coef,
                    )
                    v_loss_clipped = (v_clipped - returns_batch) ** 2
                    v_loss = 0.5 * torch.max(v_loss_unclipped, v_loss_clipped).mean()
                else:
                    v_loss = 0.5 * ((new_values - returns_batch) ** 2).mean()

                # Entropy loss
                entropy_loss = entropy.mean()

                # Total loss
                loss = pg_loss - self.config.ent_coef * entropy_loss + self.config.vf_coef * v_loss

                # Optimize
                self.optimizer.zero_grad()
                loss.backward()
                nn.utils.clip_grad_norm_(self.network.parameters(), self.config.max_grad_norm)
                self.optimizer.step()

                pg_losses.append(pg_loss.item())
                value_losses.append(v_loss.item())
                entropy_losses.append(entropy_loss.item())

            # Early stopping on KL
            if self.config.target_kl is not None and approx_kl > self.config.target_kl:
                break

        return {
            "policy_loss": np.mean(pg_losses),
            "value_loss": np.mean(value_losses),
            "entropy": np.mean(entropy_losses),
            "approx_kl": np.mean(approx_kls),
            "clipfrac": np.mean(clipfracs),
        }

    def train(
        self,
        total_timesteps: int | None = None,
        callback: Callable[[int, dict[str, Any]], bool] | None = None,
        save_path: str | None = None,
    ) -> dict[str, Any]:
        """
        Train the agent.

        Args:
            total_timesteps: Override config total_timesteps
            callback: Optional callback(step, info) -> should_stop
            save_path: Directory to save best checkpoints

        Returns:
            Training statistics
        """
        total_timesteps = total_timesteps or self.config.total_timesteps
        self.start_time = time.time()

        # Setup best checkpoint path
        if save_path and self.config.save_best:
            from pathlib import Path
            self.best_checkpoint_path = str(Path(save_path) / "best_model.pt")

        # Initialize first observation
        self._last_obs, self._last_masks = self.envs.reset(seed=42)

        num_updates = total_timesteps // self.config.batch_size
        all_episode_rewards = []

        print(f"Starting PPO training for {total_timesteps:,} timesteps...")
        print(f"  Batch size: {self.config.batch_size:,}")
        print(f"  Updates: {num_updates:,}")
        print(f"  Device: {self.device}")

        # Print initial deck configuration
        deck1, deck2 = self._get_current_decks()
        if self.config.player_playstyle:
            print(f"  Playstyle: {self.config.player_playstyle} (cycling through {len(self._player_decks)} decks)")
        elif self.config.player_faction:
            print(f"  Faction: {self.config.player_faction} (cycling through {len(self._player_decks)} decks)")
        print(f"  Initial decks: {deck1} vs {deck2}")

        for update in range(1, num_updates + 1):
            # Check if we should cycle to different decks (playstyle/faction training)
            self._maybe_cycle_decks()

            # Collect rollout
            rollout_info = self.collect_rollout()

            # PPO update
            update_info = self.update()

            # Track episode rewards
            for ep_info in rollout_info["episode_infos"]:
                all_episode_rewards.append(ep_info["reward"])

            # Logging
            if update % (self.config.log_interval // self.config.batch_size + 1) == 0:
                assert self.start_time is not None  # Set at start of train()
                elapsed = time.time() - self.start_time
                sps = self.global_step / elapsed

                recent_rewards = all_episode_rewards[-100:] if all_episode_rewards else [0]
                mean_reward = np.mean(recent_rewards)

                print(
                    f"Step {self.global_step:>10,} | "
                    f"SPS: {sps:>6,.0f} | "
                    f"Reward: {mean_reward:>6.2f} | "
                    f"PG Loss: {update_info['policy_loss']:>7.4f} | "
                    f"V Loss: {update_info['value_loss']:>7.4f} | "
                    f"Entropy: {update_info['entropy']:>5.3f}"
                )

                if self.writer is not None:
                    self.writer.add_scalar("charts/SPS", sps, self.global_step)
                    self.writer.add_scalar("charts/mean_reward", mean_reward, self.global_step)
                    self.writer.add_scalar("losses/policy_loss", update_info["policy_loss"], self.global_step)
                    self.writer.add_scalar("losses/value_loss", update_info["value_loss"], self.global_step)
                    self.writer.add_scalar("losses/entropy", update_info["entropy"], self.global_step)

            # Evaluation
            if update % (self.config.eval_interval // self.config.batch_size + 1) == 0:
                win_rate = self.evaluate_vs_greedy(self.config.eval_episodes)
                print(f"  -> Eval vs Greedy: {win_rate:.1%} win rate", end="")

                # Best checkpoint saving
                if win_rate > self.best_win_rate:
                    self.best_win_rate = win_rate
                    self.evals_without_improvement = 0
                    if self.best_checkpoint_path:
                        self.save(self.best_checkpoint_path)
                        print(" [NEW BEST]", end="")
                else:
                    self.evals_without_improvement += 1

                print()  # End the line

                if self.writer is not None:
                    self.writer.add_scalar("eval/win_rate_vs_greedy", win_rate, self.global_step)
                    self.writer.add_scalar("eval/best_win_rate", self.best_win_rate, self.global_step)

                # Early stopping check
                if (self.config.early_stopping_patience is not None and
                    self.evals_without_improvement >= self.config.early_stopping_patience):
                    print(f"  -> Early stopping: no improvement for {self.evals_without_improvement} evals")
                    break

            # Callback
            if callback is not None:
                info = {
                    "step": self.global_step,
                    "update": update,
                    **update_info,
                }
                if callback(self.global_step, info):
                    break

        return {
            "total_timesteps": self.global_step,
            "mean_reward": np.mean(all_episode_rewards) if all_episode_rewards else 0,
            "final_win_rate": self.evaluate_vs_greedy(100),
            "best_win_rate": self.best_win_rate,
        }

    def _evaluate_vs_opponent(self, opponent: str, num_games: int, seed_offset: int) -> float:
        """Evaluate trained policy against an opponent with random deck matchups.

        Uses random deck selection for both player and opponent to measure
        true generalization performance (not just one deck matchup).

        Args:
            opponent: Opponent type ("greedy" or "random")
            num_games: Number of evaluation games
            seed_offset: Seed offset for reproducibility

        Returns:
            Win rate (0.0 to 1.0)
        """
        self.network.eval()
        wins = 0
        rng = np.random.default_rng(seed_offset)

        for game_idx in range(num_games):
            deck1 = rng.choice(ALL_DECKS)
            deck2 = rng.choice(ALL_DECKS)
            env = EssenceWarsEnv(deck1=deck1, deck2=deck2, opponent=opponent)
            obs, info = env.reset(seed=game_idx + seed_offset)
            done = False

            while not done:
                # Normalize observation if enabled
                if self.obs_normalizer is not None:
                    obs_normalized = self.obs_normalizer.normalize(obs).astype(np.float32)
                else:
                    obs_normalized = obs

                obs_t = torch.tensor(obs_normalized, dtype=torch.float32, device=self.device).unsqueeze(0)
                mask_t = torch.tensor(info["action_mask"], dtype=torch.bool, device=self.device).unsqueeze(0)

                with torch.no_grad():
                    action, _, _ = self.network.get_action(obs_t, mask_t, deterministic=True)

                obs, reward, terminated, truncated, info = env.step(int(action.item()))
                done = terminated or truncated

                if done and float(reward) > 0:
                    wins += 1

        return wins / num_games

    def evaluate_vs_greedy(self, num_games: int = 100) -> float:
        """Evaluate trained policy against GreedyBot with random decks.

        Args:
            num_games: Number of evaluation games

        Returns:
            Win rate (0.0 to 1.0)
        """
        return self._evaluate_vs_opponent("greedy", num_games, seed_offset=10000)

    def evaluate_vs_random(self, num_games: int = 100) -> float:
        """Evaluate trained policy against RandomBot with random decks."""
        return self._evaluate_vs_opponent("random", num_games, seed_offset=20000)

    def evaluate_vs_mcts(self, num_games: int = 50, simulations: int = 100) -> float:
        """Evaluate trained policy against MCTS bot with random decks."""
        return self._evaluate_vs_opponent(f"mcts{simulations}", num_games, seed_offset=30000)

    def save(self, path: str) -> None:
        """Save model checkpoint."""
        checkpoint = {
            "network_state_dict": self.network.state_dict(),
            "optimizer_state_dict": self.optimizer.state_dict(),
            "global_step": self.global_step,
            "config": self.config,
        }
        # Save normalizer state if enabled
        if self.obs_normalizer is not None:
            checkpoint["obs_normalizer"] = {
                "mean": self.obs_normalizer.mean,
                "var": self.obs_normalizer.var,
                "count": self.obs_normalizer.count,
            }
        torch.save(checkpoint, path)
        print(f"Saved checkpoint to {path}")

    def load(self, path: str) -> None:
        """Load model checkpoint."""
        checkpoint = torch.load(path, map_location=self.device, weights_only=False)
        self.network.load_state_dict(checkpoint["network_state_dict"])
        self.optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
        self.global_step = checkpoint["global_step"]
        # Restore normalizer state if present
        if "obs_normalizer" in checkpoint and self.obs_normalizer is not None:
            self.obs_normalizer.mean = checkpoint["obs_normalizer"]["mean"]
            self.obs_normalizer.var = checkpoint["obs_normalizer"]["var"]
            self.obs_normalizer.count = checkpoint["obs_normalizer"]["count"]
        print(f"Loaded checkpoint from {path}")

    def get_policy(self) -> EssenceWarsNetwork:
        """Get the trained policy network."""
        return self.network
