"""Compact PPO stack for the JAX Lost Cities engine."""

from __future__ import annotations

import argparse
import json
import math
import shutil
import time
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import NamedTuple

import flax.linen as nn
import jax
import jax.numpy as jnp
import optax
import orbax.checkpoint as ocp
import yaml
from flax.training import train_state

from lost_cities_jax.engine import board_score, legal_action_mask, reset, reset_from_order, step
from lost_cities_jax.obs import observation
from lost_cities_jax.opponents import policy_by_name, random_legal_action
from lost_cities_jax.types import (
    CARDS_PER_COLOR,
    LOC_P0_BOARD,
    MAX_STEPS,
    N_ACTIONS,
    N_CARDS,
    N_COLORS,
    OBS_DIM,
    PLAY,
    State,
)

ARTIFACT_ROOT = "/mnt/2tbhdd/coolrl-lost-cities-artifacts/jax-ppo-static-opponents"
NEG_INF = -1.0e9


@dataclass
class RunConfig:
    experiment_name: str = "jax-ppo-discard-only"
    seed: int = 20260704
    learner_seat: int = 0
    total_updates: int = 200
    log_every: int = 1
    checkpoint_every: int = 10
    artifact_root: str = ARTIFACT_ROOT


@dataclass
class OpponentConfig:
    name: str = "discard_only"


@dataclass
class NetworkConfig:
    hidden_size: int = 512
    num_layers: int = 3


@dataclass
class PPOHyperConfig:
    batch_games: int = 8192
    rollout_steps: int = MAX_STEPS
    gamma: float = 1.0
    gae_lambda: float = 0.95
    clip_epsilon: float = 0.2
    entropy_coef: float = 0.01
    value_coef: float = 0.5
    max_grad_norm: float = 0.5
    learning_rate: float = 3.0e-4
    epochs: int = 4
    minibatches: int = 8


@dataclass
class RewardConfig:
    """Terminal reward plus potential-based shaping on the score difference.

    ``potential_shaping_anneal_steps`` counts **learner actions**, not scan
    steps. It used to be fed padded scan steps, which made a 5M-step anneal
    expire inside the first two updates; any earlier conclusion that shaping
    does not help was drawn with shaping effectively off.
    """

    terminal_scale: float = 50.0
    potential_shaping_initial: float = 1.0
    potential_shaping_final: float = 0.0
    potential_shaping_anneal_steps: int = 5_000_000


@dataclass
class EvalConfig:
    games: int = 10_000
    duplicate: bool = True
    shuffle_bank_seed: int = 20260704
    batch_games: int = 8192


@dataclass
class JaxPPOConfig:
    run: RunConfig = field(default_factory=RunConfig)
    opponent: OpponentConfig = field(default_factory=OpponentConfig)
    network: NetworkConfig = field(default_factory=NetworkConfig)
    ppo: PPOHyperConfig = field(default_factory=PPOHyperConfig)
    reward: RewardConfig = field(default_factory=RewardConfig)
    evaluation: EvalConfig = field(default_factory=EvalConfig)


class ActorCritic(nn.Module):
    hidden_size: int = 512
    num_layers: int = 3

    @nn.compact
    def __call__(self, obs: jax.Array) -> tuple[jax.Array, jax.Array]:
        x = obs
        for _ in range(self.num_layers):
            x = nn.Dense(self.hidden_size)(x)
            x = nn.relu(x)
        logits = nn.Dense(N_ACTIONS)(x)
        value = nn.Dense(1)(x)
        return logits, jnp.squeeze(value, axis=-1)


class TrainState(train_state.TrainState):
    pass


class Transition(NamedTuple):
    obs: jax.Array
    legal_mask: jax.Array
    action: jax.Array
    log_prob: jax.Array
    value: jax.Array
    reward: jax.Array
    done: jax.Array
    active: jax.Array
    actor_mask: jax.Array
    entropy: jax.Array
    play_action: jax.Array


class EpisodeEnd(NamedTuple):
    """Per-step episode-completion records emitted by the rollout scan.

    With in-scan auto-reset an env slot hosts several episodes per rollout, so
    ``final_env`` is mid-episode and summing rewards over the scan axis mixes
    episodes. Every field here is masked by ``done``: read it only where
    ``done`` is set.
    """

    done: jax.Array
    episode_return: jax.Array
    length: jax.Array
    opened_colors: jax.Array
    positive_expeditions: jax.Array
    hit_max_steps: jax.Array


class EvalBatch(NamedTuple):
    wins: jax.Array
    losses: jax.Array
    ties: jax.Array
    score_diff_sum: jax.Array
    game_length_sum: jax.Array
    play_action_sum: jax.Array
    learner_action_sum: jax.Array
    opened_colors_sum: jax.Array
    positive_expeditions_sum: jax.Array
    games: jax.Array


class MatchBatch(NamedTuple):
    score_diff: jax.Array
    game_length: jax.Array
    play_action_sum: jax.Array
    learner_action_sum: jax.Array
    opened_colors: jax.Array
    positive_expeditions: jax.Array
    max_steps: jax.Array


class LeagueAssignments(NamedTuple):
    learner_seat: jax.Array
    opponent_index: jax.Array
    use_mirror: jax.Array


def load_config(path: str | Path | None = None, overrides: dict | None = None) -> JaxPPOConfig:
    data: dict = {}
    if path is not None:
        with Path(path).open("r", encoding="utf-8") as handle:
            loaded = yaml.safe_load(handle) or {}
        if not isinstance(loaded, dict):
            msg = f"config root must be a mapping: {path}"
            raise ValueError(msg)
        data = loaded
    if overrides:
        data = _deep_merge(data, overrides)
    cfg = JaxPPOConfig()
    return JaxPPOConfig(
        run=RunConfig(**{**asdict(cfg.run), **data.get("run", {})}),
        opponent=OpponentConfig(**{**asdict(cfg.opponent), **data.get("opponent", {})}),
        network=NetworkConfig(**{**asdict(cfg.network), **data.get("network", {})}),
        ppo=PPOHyperConfig(**{**asdict(cfg.ppo), **data.get("ppo", {})}),
        reward=RewardConfig(**{**asdict(cfg.reward), **data.get("reward", {})}),
        evaluation=EvalConfig(**{**asdict(cfg.evaluation), **data.get("evaluation", {})}),
    )


def create_train_state(cfg: JaxPPOConfig, rng: jax.Array) -> TrainState:
    model = ActorCritic(cfg.network.hidden_size, cfg.network.num_layers)
    params = model.init(rng, jnp.zeros((1, OBS_DIM), dtype=jnp.float32))
    tx = optax.chain(
        optax.clip_by_global_norm(cfg.ppo.max_grad_norm),
        optax.adam(cfg.ppo.learning_rate),
    )
    return TrainState.create(apply_fn=model.apply, params=params, tx=tx)


def train(
    cfg: JaxPPOConfig,
    *,
    resume: str | None = None,
    opponent_policy=None,
) -> Path:
    run_dir = _create_run_dir(cfg)
    _write_json(run_dir / "config.json", asdict(cfg))
    metrics_path = run_dir / "metrics.jsonl"

    rng = jax.random.PRNGKey(cfg.run.seed)
    rng, init_key, reset_key = jax.random.split(rng, 3)
    state = create_train_state(cfg, init_key)
    env_state = jax.jit(jax.vmap(reset))(jax.random.split(reset_key, cfg.ppo.batch_games))

    if resume:
        state = restore_checkpoint(Path(resume), state)

    if opponent_policy is None:
        opponent_policy = policy_by_name(cfg.opponent.name)
    train_iteration = make_train_iteration(cfg, opponent_policy)
    start = time.perf_counter()

    # Anneal against learner actions taken, not scan steps issued: most scan
    # steps used to be spent stepping already-done envs, which made the anneal
    # finish within the first couple of updates. The count lags by one update
    # because it comes back as a device metric.
    learner_actions_seen = 0

    for update in range(cfg.run.total_updates):
        shaping_coef = shaping_coefficient(cfg, learner_actions_seen)
        iter_start = time.perf_counter()
        state, env_state, rng, metrics = train_iteration(state, env_state, rng, shaping_coef)
        jax.tree_util.tree_leaves(metrics)[0].block_until_ready()
        row = _metrics_to_row(metrics, update, shaping_coef, cfg)
        learner_actions_seen += int(row["learner_actions"])
        row["learner_actions_total"] = learner_actions_seen
        row["iteration_seconds"] = time.perf_counter() - iter_start
        row["elapsed_seconds"] = time.perf_counter() - start
        _append_jsonl(metrics_path, row)

        if update % cfg.run.log_every == 0:
            print(
                json.dumps(
                    {
                        "update": update,
                        "return_mean": row["return_mean"],
                        "game_length_mean": row["game_length_mean"],
                        "play_action_rate": row["play_action_rate"],
                        "entropy_mean": row["entropy_mean"],
                    },
                    sort_keys=True,
                ),
                flush=True,
            )

        if cfg.run.checkpoint_every and (update + 1) % cfg.run.checkpoint_every == 0:
            save_checkpoint(run_dir / f"checkpoint_{update + 1:06d}", state, cfg)
            save_checkpoint(run_dir / "latest", state, cfg)

    save_checkpoint(run_dir / "latest", state, cfg)
    return run_dir


def make_train_iteration(cfg: JaxPPOConfig, opponent_policy):
    rollout_fn = make_rollout_fn(cfg, opponent_policy)

    @jax.jit
    def train_iteration(
        state: TrainState, env_state: State, rng: jax.Array, shaping_coef: jax.Array
    ):
        rng, rollout_key, update_key = jax.random.split(rng, 3)
        # The rollout resets finished envs in-scan, so env_state never comes
        # back done and needs no reset here.
        env_state, transitions, last_value, rollout_stats = rollout_fn(
            state, env_state, rollout_key, shaping_coef
        )
        advantages, returns = compute_gae(
            transitions.reward,
            transitions.value,
            transitions.done,
            cfg.ppo.gamma,
            cfg.ppo.gae_lambda,
            last_value,
        )
        state, update_metrics = ppo_update(state, transitions, advantages, returns, update_key, cfg)
        return state, env_state, rng, {**rollout_stats, **update_metrics}

    return train_iteration


def make_league_train_iteration(
    cfg: JaxPPOConfig,
    opponent_policies,
    opponent_probs: jax.Array,
    mirror_probability: float,
):
    rollout_fn = make_league_rollout_fn(
        cfg,
        opponent_policies,
        jnp.asarray(opponent_probs, dtype=jnp.float32),
        mirror_probability,
    )

    @jax.jit
    def train_iteration(
        state: TrainState,
        env_state: State,
        assignments: LeagueAssignments,
        rng: jax.Array,
        shaping_coef: jax.Array,
    ):
        rng, rollout_key, update_key = jax.random.split(rng, 3)
        # The rollout resets finished envs and redraws their league assignments
        # in-scan, so neither needs to be refreshed here.
        env_state, assignments, transitions, last_value, rollout_stats = rollout_fn(
            state, env_state, assignments, rollout_key, shaping_coef
        )
        advantages, returns = compute_gae(
            transitions.reward,
            transitions.value,
            transitions.done,
            cfg.ppo.gamma,
            cfg.ppo.gae_lambda,
            last_value,
        )
        state, update_metrics = ppo_update(state, transitions, advantages, returns, update_key, cfg)
        return state, env_state, assignments, rng, {**rollout_stats, **update_metrics}

    return train_iteration


def make_rollout_fn(cfg: JaxPPOConfig, opponent_policy):
    learner = jnp.asarray(cfg.run.learner_seat, dtype=jnp.int32)
    opponent = jnp.asarray(1 - cfg.run.learner_seat, dtype=jnp.int32)

    learners = jnp.full((cfg.ppo.batch_games,), cfg.run.learner_seat, dtype=jnp.int32)

    @jax.jit
    def rollout_fn(state: TrainState, env_state: State, rng: jax.Array, shaping_coef: jax.Array):
        def body(carry, _):
            env, episode_return, key = carry
            key, learner_key, opponent_key, reset_key = jax.random.split(key, 4)
            obs = jax.vmap(observation, in_axes=(0, None))(env, learner)
            legal = jax.vmap(legal_action_mask)(env)
            logits, value = state.apply_fn(state.params, obs)
            masked_logits = mask_logits(logits, legal)
            learner_actions = jax.random.categorical(learner_key, masked_logits, axis=-1).astype(
                jnp.int32
            )
            log_prob = action_log_prob(masked_logits, learner_actions)
            entropy = categorical_entropy(masked_logits)

            opponent_keys = jax.random.split(opponent_key, cfg.ppo.batch_games)
            opponent_actions = jax.vmap(opponent_policy, in_axes=(0, None, 0))(
                env, opponent, opponent_keys
            )
            learner_turn = env.to_move.astype(jnp.int32) == learner
            active = ~env.done
            actions = jnp.where(learner_turn, learner_actions, opponent_actions)

            before_diff = batch_score_diff(env, learner)
            next_env, _, _ = jax.vmap(step, in_axes=(0, 0))(env, actions)
            after_diff = batch_score_diff(next_env, learner)
            terminal_reward = jnp.tanh(after_diff / cfg.reward.terminal_scale)
            done = next_env.done
            terminal = active & done
            reward = jnp.where(terminal, terminal_reward, 0.0)
            reward = reward + shaping_coef * (after_diff - before_diff)
            reward = jnp.where(active, reward, 0.0)

            actor_mask = active & learner_turn
            place_type = (actions % 12) // 6
            transition = Transition(
                obs=obs,
                legal_mask=legal,
                action=actions,
                log_prob=log_prob,
                value=value,
                reward=reward,
                done=done,
                active=active,
                actor_mask=actor_mask,
                entropy=entropy,
                play_action=(place_type == PLAY) & actor_mask,
            )

            episode_return = episode_return + reward
            episode = _episode_end(next_env, learners, terminal, episode_return)
            next_env = reset_done_envs(next_env, reset_key, cfg.ppo.batch_games)
            episode_return = jnp.where(done, 0.0, episode_return)
            return (next_env, episode_return, key), (transition, episode)

        init_return = jnp.zeros((cfg.ppo.batch_games,), dtype=jnp.float32)
        (next_env, _, rng), (transitions, episodes) = jax.lax.scan(
            body, (env_state, init_return, rng), xs=None, length=cfg.ppo.rollout_steps
        )
        final_obs = jax.vmap(observation, in_axes=(0, None))(next_env, learner)
        _, last_value = state.apply_fn(state.params, final_obs)
        metrics = episode_metrics(transitions, episodes)
        return next_env, transitions, last_value, metrics

    return rollout_fn


def make_league_rollout_fn(
    cfg: JaxPPOConfig,
    opponent_policies,
    opponent_probs: jax.Array,
    mirror_probability: float,
):
    if not opponent_policies:
        msg = "league rollout requires at least one pool opponent"
        raise ValueError(msg)

    @jax.jit
    def rollout_fn(
        state: TrainState,
        env_state: State,
        assignments: LeagueAssignments,
        rng: jax.Array,
        shaping_coef: jax.Array,
    ):
        def body(carry, _):
            env, assignments, episode_return, key = carry
            key, learner_key, mirror_key, pool_key, reset_key = jax.random.split(key, 5)
            learner = assignments.learner_seat.astype(jnp.int32)
            opponent = 1 - learner

            obs = jax.vmap(observation, in_axes=(0, 0))(env, learner)
            legal = jax.vmap(legal_action_mask)(env)
            logits, value = state.apply_fn(state.params, obs)
            masked_logits = mask_logits(logits, legal)
            learner_actions = jax.random.categorical(learner_key, masked_logits, axis=-1).astype(
                jnp.int32
            )
            log_prob = action_log_prob(masked_logits, learner_actions)
            entropy = categorical_entropy(masked_logits)

            mirror_obs = jax.vmap(observation, in_axes=(0, 0))(env, opponent)
            mirror_logits, _ = state.apply_fn(state.params, mirror_obs)
            mirror_logits = jax.lax.stop_gradient(mirror_logits)
            mirror_actions = jax.random.categorical(
                mirror_key, mask_logits(mirror_logits, legal), axis=-1
            ).astype(jnp.int32)

            pool_keys = jax.random.split(pool_key, cfg.ppo.batch_games)
            pool_actions = []
            for policy in opponent_policies:
                pool_actions.append(jax.vmap(policy, in_axes=(0, 0, 0))(env, opponent, pool_keys))
            stacked_pool_actions = jnp.stack(pool_actions, axis=0).astype(jnp.int32)
            selected_pool_actions = jnp.take_along_axis(
                stacked_pool_actions,
                assignments.opponent_index.astype(jnp.int32)[None, :],
                axis=0,
            )[0]
            opponent_actions = jnp.where(
                assignments.use_mirror, mirror_actions, selected_pool_actions
            )

            learner_turn = env.to_move.astype(jnp.int32) == learner
            active = ~env.done
            actions = jnp.where(learner_turn, learner_actions, opponent_actions)

            before_diff = batch_score_diff_for_players(env, learner)
            next_env, _, _ = jax.vmap(step, in_axes=(0, 0))(env, actions)
            after_diff = batch_score_diff_for_players(next_env, learner)
            terminal_reward = jnp.tanh(after_diff / cfg.reward.terminal_scale)
            done = next_env.done
            terminal = active & done
            reward = jnp.where(terminal, terminal_reward, 0.0)
            reward = reward + shaping_coef * (after_diff - before_diff)
            reward = jnp.where(active, reward, 0.0)

            actor_mask = active & learner_turn
            place_type = (actions % 12) // 6
            transition = Transition(
                obs=obs,
                legal_mask=legal,
                action=actions,
                log_prob=log_prob,
                value=value,
                reward=reward,
                done=done,
                active=active,
                actor_mask=actor_mask,
                entropy=entropy,
                play_action=(place_type == PLAY) & actor_mask,
            )

            episode_return = episode_return + reward
            # Scored against the seat that just played the episode out, not the
            # freshly drawn one.
            episode = _episode_end(next_env, learner, terminal, episode_return)
            next_env, assignments = reset_done_envs_with_assignments(
                next_env,
                assignments,
                reset_key,
                cfg.ppo.batch_games,
                opponent_probs,
                mirror_probability,
            )
            episode_return = jnp.where(done, 0.0, episode_return)
            return (next_env, assignments, episode_return, key), (transition, episode)

        init_return = jnp.zeros((cfg.ppo.batch_games,), dtype=jnp.float32)
        (next_env, assignments, _, rng), (transitions, episodes) = jax.lax.scan(
            body, (env_state, assignments, init_return, rng), xs=None, length=cfg.ppo.rollout_steps
        )
        final_learner = assignments.learner_seat.astype(jnp.int32)
        final_obs = jax.vmap(observation, in_axes=(0, 0))(next_env, final_learner)
        _, last_value = state.apply_fn(state.params, final_obs)
        metrics = episode_metrics(transitions, episodes)
        return next_env, assignments, transitions, last_value, metrics

    return rollout_fn


def random_rollout(cfg: JaxPPOConfig) -> dict:
    opponent_policy = policy_by_name(cfg.opponent.name)
    rng = jax.random.PRNGKey(cfg.run.seed)
    rng, reset_key = jax.random.split(rng)
    env_state = jax.jit(jax.vmap(reset))(jax.random.split(reset_key, cfg.ppo.batch_games))
    learner = jnp.asarray(cfg.run.learner_seat, dtype=jnp.int32)
    opponent = jnp.asarray(1 - cfg.run.learner_seat, dtype=jnp.int32)
    learners = jnp.full((cfg.ppo.batch_games,), cfg.run.learner_seat, dtype=jnp.int32)

    @jax.jit
    def rollout(env_state: State, rng: jax.Array):
        def body(carry, _):
            env, episode_return, key = carry
            key, learner_key, opponent_key = jax.random.split(key, 3)
            learner_keys = jax.random.split(learner_key, cfg.ppo.batch_games)
            opponent_keys = jax.random.split(opponent_key, cfg.ppo.batch_games)
            learner_actions = jax.vmap(random_legal_action, in_axes=(0, None, 0))(
                env, learner, learner_keys
            )
            opponent_actions = jax.vmap(opponent_policy, in_axes=(0, None, 0))(
                env, opponent, opponent_keys
            )
            learner_turn = env.to_move.astype(jnp.int32) == learner
            active = ~env.done
            actions = jnp.where(learner_turn, learner_actions, opponent_actions)
            before_diff = batch_score_diff(env, learner)
            next_env, _, _ = jax.vmap(step, in_axes=(0, 0))(env, actions)
            after_diff = batch_score_diff(next_env, learner)
            reward = jnp.where(active, after_diff - before_diff, 0.0)
            terminal = active & next_env.done
            actor_mask = active & learner_turn
            place_type = (actions % 12) // 6
            dummy_obs = jnp.zeros((cfg.ppo.batch_games, OBS_DIM), dtype=jnp.float32)
            dummy_mask = jnp.zeros((cfg.ppo.batch_games, N_ACTIONS), dtype=jnp.bool_)
            transition = Transition(
                obs=dummy_obs,
                legal_mask=dummy_mask,
                action=actions,
                log_prob=jnp.zeros((cfg.ppo.batch_games,), dtype=jnp.float32),
                value=jnp.zeros((cfg.ppo.batch_games,), dtype=jnp.float32),
                reward=reward,
                done=next_env.done,
                active=active,
                actor_mask=actor_mask,
                entropy=jnp.zeros((cfg.ppo.batch_games,), dtype=jnp.float32),
                play_action=(place_type == PLAY) & actor_mask,
            )
            episode_return = episode_return + reward
            episode = _episode_end(next_env, learners, terminal, episode_return)
            return (next_env, episode_return, key), (transition, episode)

        init_return = jnp.zeros((cfg.ppo.batch_games,), dtype=jnp.float32)
        (next_env, _, rng), (transitions, episodes) = jax.lax.scan(
            body, (env_state, init_return, rng), xs=None, length=cfg.ppo.rollout_steps
        )
        return rng, episode_metrics(transitions, episodes)

    rng, metrics = rollout(env_state, rng)
    jax.tree_util.tree_leaves(metrics)[0].block_until_ready()
    return _metrics_to_row(metrics, update=0, shaping_coef=0.0, cfg=cfg)


def reset_done_envs(env_state: State, rng: jax.Array, batch_games: int) -> State:
    fresh = jax.vmap(reset)(jax.random.split(rng, batch_games))
    done = env_state.done
    return jax.tree_util.tree_map(
        lambda old, new: jnp.where(_broadcast_done(done, old), new, old), env_state, fresh
    )


def sample_league_assignments(
    rng: jax.Array,
    batch_games: int,
    opponent_probs: jax.Array,
    mirror_probability: float,
) -> LeagueAssignments:
    seat_key, opponent_key, mirror_key = jax.random.split(rng, 3)
    learner_seat = jax.random.bernoulli(seat_key, 0.5, shape=(batch_games,)).astype(jnp.int32)
    logits = jnp.log(jnp.maximum(jnp.asarray(opponent_probs, dtype=jnp.float32), 1.0e-12))
    opponent_index = jax.random.categorical(opponent_key, logits, shape=(batch_games,)).astype(
        jnp.int32
    )
    use_mirror = jax.random.bernoulli(mirror_key, mirror_probability, shape=(batch_games,))
    return LeagueAssignments(
        learner_seat=learner_seat,
        opponent_index=opponent_index,
        use_mirror=use_mirror,
    )


def reset_done_envs_with_assignments(
    env_state: State,
    assignments: LeagueAssignments,
    rng: jax.Array,
    batch_games: int,
    opponent_probs: jax.Array,
    mirror_probability: float,
) -> tuple[State, LeagueAssignments]:
    env_key, assignment_key = jax.random.split(rng)
    fresh_env = jax.vmap(reset)(jax.random.split(env_key, batch_games))
    fresh_assignments = sample_league_assignments(
        assignment_key, batch_games, opponent_probs, mirror_probability
    )
    done = env_state.done
    next_env = jax.tree_util.tree_map(
        lambda old, new: jnp.where(_broadcast_done(done, old), new, old), env_state, fresh_env
    )
    next_assignments = LeagueAssignments(
        learner_seat=jnp.where(done, fresh_assignments.learner_seat, assignments.learner_seat),
        opponent_index=jnp.where(
            done, fresh_assignments.opponent_index, assignments.opponent_index
        ),
        use_mirror=jnp.where(done, fresh_assignments.use_mirror, assignments.use_mirror),
    )
    return next_env, next_assignments


def ppo_update(
    state: TrainState,
    transitions: Transition,
    advantages: jax.Array,
    returns: jax.Array,
    rng: jax.Array,
    cfg: JaxPPOConfig,
) -> tuple[TrainState, dict[str, jax.Array]]:
    batch_size = cfg.ppo.rollout_steps * cfg.ppo.batch_games
    minibatch_size = batch_size // cfg.ppo.minibatches
    flat = _flatten_transitions(transitions)
    flat_advantages = advantages.reshape((batch_size,))
    flat_returns = returns.reshape((batch_size,))
    actor_mask = flat.actor_mask.astype(jnp.float32)
    adv_mean = masked_mean(flat_advantages, actor_mask)
    adv_std = jnp.sqrt(masked_mean((flat_advantages - adv_mean) ** 2, actor_mask) + 1.0e-8)
    flat_advantages = (flat_advantages - adv_mean) / adv_std

    def loss_fn(params, mb, mb_adv, mb_returns):
        logits, value = state.apply_fn(params, mb.obs)
        masked_logits = mask_logits(logits, mb.legal_mask)
        log_prob = action_log_prob(masked_logits, mb.action)
        entropy = categorical_entropy(masked_logits)
        ratio = jnp.exp(log_prob - mb.log_prob)
        actor_weight = mb.actor_mask.astype(jnp.float32)
        active_weight = mb.active.astype(jnp.float32)
        unclipped = ratio * mb_adv
        clipped = jnp.clip(ratio, 1.0 - cfg.ppo.clip_epsilon, 1.0 + cfg.ppo.clip_epsilon) * mb_adv
        actor_loss = -masked_mean(jnp.minimum(unclipped, clipped), actor_weight)
        value_loss = masked_mean((mb_returns - value) ** 2, active_weight)
        entropy_loss = masked_mean(entropy, actor_weight)
        loss = actor_loss + cfg.ppo.value_coef * value_loss - cfg.ppo.entropy_coef * entropy_loss
        approx_kl = masked_mean(mb.log_prob - log_prob, actor_weight)
        return loss, {
            "loss": loss,
            "actor_loss": actor_loss,
            "value_loss": value_loss,
            "entropy_loss": entropy_loss,
            "approx_kl": approx_kl,
        }

    def epoch_update(carry, key):
        train_state = carry
        metrics_init = {
            "loss": jnp.array(0.0),
            "actor_loss": jnp.array(0.0),
            "value_loss": jnp.array(0.0),
            "entropy_loss": jnp.array(0.0),
            "approx_kl": jnp.array(0.0),
        }
        permutation = jax.random.permutation(key, batch_size)

        def minibatch_update(carry, idx):
            train_state, metrics_accum = carry
            start = idx * minibatch_size
            mb_idx = jax.lax.dynamic_slice(permutation, (start,), (minibatch_size,))
            mb = jax.tree_util.tree_map(lambda x: x[mb_idx], flat)
            mb_adv = flat_advantages[mb_idx]
            mb_returns = flat_returns[mb_idx]
            (loss, metrics), grads = jax.value_and_grad(loss_fn, has_aux=True)(
                train_state.params, mb, mb_adv, mb_returns
            )
            del loss
            train_state = train_state.apply_gradients(grads=grads)
            metrics_accum = jax.tree_util.tree_map(lambda a, b: a + b, metrics_accum, metrics)
            return (train_state, metrics_accum), None

        (train_state, metrics), _ = jax.lax.scan(
            minibatch_update,
            (train_state, metrics_init),
            jnp.arange(cfg.ppo.minibatches),
        )
        metrics = jax.tree_util.tree_map(lambda x: x / cfg.ppo.minibatches, metrics)
        return train_state, metrics

    keys = jax.random.split(rng, cfg.ppo.epochs)
    state, metrics = jax.lax.scan(epoch_update, state, keys)
    metrics = jax.tree_util.tree_map(lambda x: jnp.mean(x, axis=0), metrics)
    return state, metrics


def compute_gae(
    rewards: jax.Array,
    values: jax.Array,
    dones: jax.Array,
    gamma: float,
    gae_lambda: float,
    last_value: jax.Array | None = None,
) -> tuple[jax.Array, jax.Array]:
    """GAE over a rollout that may end mid-episode.

    ``last_value`` is V(s_T) for the state the scan stopped on. Episodes that
    are cut by the rollout boundary bootstrap from it; leaving it at zero would
    tell the critic every truncated episode is worth nothing.
    """

    if last_value is None:
        last_value = jnp.zeros_like(values[-1])

    def body(carry, x):
        next_value, next_advantage = carry
        reward, value, done = x
        nonterminal = 1.0 - done.astype(jnp.float32)
        delta = reward + gamma * next_value * nonterminal - value
        advantage = delta + gamma * gae_lambda * nonterminal * next_advantage
        return (value, advantage), advantage

    init = (last_value, jnp.zeros_like(values[-1]))
    _, advantages_rev = jax.lax.scan(body, init, (rewards[::-1], values[::-1], dones[::-1]))
    advantages = advantages_rev[::-1]
    return advantages, advantages + values


def evaluate(
    cfg: JaxPPOConfig,
    checkpoint: str | Path,
    *,
    opponent_name: str | None = None,
    games: int | None = None,
    duplicate: bool | None = None,
    output: str | Path | None = None,
) -> dict:
    opponent_name = opponent_name or cfg.opponent.name
    games = games or cfg.evaluation.games
    duplicate = cfg.evaluation.duplicate if duplicate is None else duplicate
    eval_cfg = cfg.evaluation
    state = create_train_state(cfg, jax.random.PRNGKey(0))
    state = restore_checkpoint(Path(checkpoint), state)
    opponent_policy = policy_by_name(opponent_name)
    eval_batch = make_eval_batch_fn(cfg, opponent_policy)

    total = None
    for seat in [0, 1] if duplicate else [cfg.run.learner_seat]:
        for start in range(0, games, eval_cfg.batch_games):
            count = min(eval_cfg.batch_games, games - start)
            orders = make_shuffle_bank(eval_cfg.shuffle_bank_seed + start, count)
            metrics = eval_batch(state.params, jnp.asarray(orders, dtype=jnp.int8), seat)
            metrics = jax.tree_util.tree_map(lambda x: float(x), metrics)
            total = metrics if total is None else _sum_dicts(total, metrics)

    assert total is not None
    total_games = int(total["games"])
    wins = int(total["wins"])
    win_rate = wins / total_games
    wilson_low, wilson_high = wilson_interval(wins, total_games)
    result = {
        **total,
        "games": total_games,
        "win_rate": win_rate,
        "wilson_low": wilson_low,
        "wilson_high": wilson_high,
        "mean_score_diff": total["score_diff_sum"] / total_games,
        "mean_game_length": total["game_length_sum"] / total_games,
        "play_action_rate": total["play_action_sum"] / max(total["learner_action_sum"], 1.0),
        "opened_colors_per_game": total["opened_colors_sum"] / total_games,
        "positive_expeditions_per_game": total["positive_expeditions_sum"] / total_games,
        "opponent": opponent_name,
        "duplicate": duplicate,
    }
    if output:
        _write_json(Path(output), result)
    return result


def make_eval_batch_fn(cfg: JaxPPOConfig, opponent_policy):
    opponent_fns = opponent_policy

    @jax.jit
    def eval_batch(params, deck_orders: jax.Array, learner_seat: int):
        learner = jnp.asarray(learner_seat, dtype=jnp.int32)
        opponent = 1 - learner
        env_state = jax.vmap(reset_from_order)(deck_orders)

        def body(env, _):
            obs = jax.vmap(observation, in_axes=(0, None))(env, learner)
            legal = jax.vmap(legal_action_mask)(env)
            model = ActorCritic(cfg.network.hidden_size, cfg.network.num_layers)
            logits, _ = model.apply(params, obs)
            learner_actions = jnp.argmax(mask_logits(logits, legal), axis=-1).astype(jnp.int32)
            keys = jax.random.split(jax.random.PRNGKey(0), deck_orders.shape[0])
            opponent_actions = jax.vmap(opponent_fns, in_axes=(0, None, 0))(env, opponent, keys)
            learner_turn = env.to_move.astype(jnp.int32) == learner
            active = ~env.done
            actions = jnp.where(learner_turn, learner_actions, opponent_actions)
            next_env, _, _ = jax.vmap(step, in_axes=(0, 0))(env, actions)
            place_type = (actions % 12) // 6
            return next_env, ((place_type == PLAY) & learner_turn & active, learner_turn & active)

        final_env, (play_actions, learner_actions) = jax.lax.scan(
            body, env_state, xs=None, length=MAX_STEPS
        )
        scores = jax.vmap(board_score)(final_env)
        diff = scores[:, learner] - scores[:, opponent]
        color_scores = jax.vmap(color_scores_for_player, in_axes=(0, None))(final_env, learner)
        opened = jnp.sum(final_env.col_len[:, learner, :] > 0, axis=-1)
        positive = jnp.sum(color_scores > 0, axis=-1)
        return EvalBatch(
            wins=jnp.sum(diff > 0),
            losses=jnp.sum(diff < 0),
            ties=jnp.sum(diff == 0),
            score_diff_sum=jnp.sum(diff),
            game_length_sum=jnp.sum(final_env.step_count),
            play_action_sum=jnp.sum(play_actions),
            learner_action_sum=jnp.sum(learner_actions),
            opened_colors_sum=jnp.sum(opened),
            positive_expeditions_sum=jnp.sum(positive),
            games=deck_orders.shape[0],
        )._asdict()

    return eval_batch


def checkpoint_policy(cfg: JaxPPOConfig, checkpoint: str | Path):
    state = create_train_state(cfg, jax.random.PRNGKey(0))
    state = restore_checkpoint(Path(checkpoint), state)
    return checkpoint_policy_from_params(cfg, state.params)


def checkpoint_policy_from_params(cfg: JaxPPOConfig, params):
    model = ActorCritic(cfg.network.hidden_size, cfg.network.num_layers)

    def policy(state: State, player: jax.Array, rng: jax.Array) -> jax.Array:
        del rng
        obs = observation(state, player)[None, :]
        legal = legal_action_mask(state)[None, :]
        logits, _ = model.apply(params, obs)
        return jnp.argmax(mask_logits(logits, legal), axis=-1)[0].astype(jnp.int32)

    return policy


def train_against_checkpoint(
    cfg: JaxPPOConfig,
    opponent_cfg: JaxPPOConfig,
    opponent_checkpoint: str | Path,
    *,
    resume: str | None = None,
) -> Path:
    policy = checkpoint_policy(opponent_cfg, opponent_checkpoint)
    return train(cfg, resume=resume, opponent_policy=policy)


def evaluate_checkpoint_match(
    learner_cfg: JaxPPOConfig,
    learner_checkpoint: str | Path,
    opponent_cfg: JaxPPOConfig,
    opponent_checkpoint: str | Path,
    *,
    games: int | None = None,
    duplicate: bool | None = None,
    output: str | Path | None = None,
) -> dict:
    learner_policy = checkpoint_policy(learner_cfg, learner_checkpoint)
    opponent_policy = checkpoint_policy(opponent_cfg, opponent_checkpoint)
    games = games or learner_cfg.evaluation.games
    duplicate = learner_cfg.evaluation.duplicate if duplicate is None else duplicate
    return evaluate_policy_match(
        learner_policy,
        opponent_policy,
        games=games,
        duplicate=duplicate,
        shuffle_bank_seed=learner_cfg.evaluation.shuffle_bank_seed,
        batch_games=learner_cfg.evaluation.batch_games,
        output=output,
    )


def evaluate_checkpoint_vs_static(
    learner_cfg: JaxPPOConfig,
    learner_checkpoint: str | Path,
    opponent_name: str,
    *,
    games: int | None = None,
    duplicate: bool | None = None,
    output: str | Path | None = None,
) -> dict:
    learner_policy = checkpoint_policy(learner_cfg, learner_checkpoint)
    opponent_policy = policy_by_name(opponent_name)
    games = games or learner_cfg.evaluation.games
    duplicate = learner_cfg.evaluation.duplicate if duplicate is None else duplicate
    result = evaluate_policy_match(
        learner_policy,
        opponent_policy,
        games=games,
        duplicate=duplicate,
        shuffle_bank_seed=learner_cfg.evaluation.shuffle_bank_seed,
        batch_games=learner_cfg.evaluation.batch_games,
        output=None,
    )
    result["opponent"] = opponent_name
    if output:
        _write_json(Path(output), result)
    return result


def evaluate_static_mirror(
    policy_name: str,
    *,
    games: int,
    duplicate: bool = True,
    shuffle_bank_seed: int = 20260704,
    batch_games: int = 8192,
    output: str | Path | None = None,
) -> dict:
    policy = policy_by_name(policy_name)
    return evaluate_policy_match(
        policy,
        policy,
        games=games,
        duplicate=duplicate,
        shuffle_bank_seed=shuffle_bank_seed,
        batch_games=batch_games,
        output=output,
    )


def evaluate_static_match(
    learner_name: str,
    opponent_name: str,
    *,
    games: int,
    duplicate: bool = True,
    shuffle_bank_seed: int = 20260704,
    batch_games: int = 8192,
    output: str | Path | None = None,
) -> dict:
    result = evaluate_policy_match(
        policy_by_name(learner_name),
        policy_by_name(opponent_name),
        games=games,
        duplicate=duplicate,
        shuffle_bank_seed=shuffle_bank_seed,
        batch_games=batch_games,
        output=None,
    )
    result["learner"] = learner_name
    result["opponent"] = opponent_name
    if output:
        _write_json(Path(output), result)
    return result


def evaluate_policy_match(
    learner_policy,
    opponent_policy,
    *,
    games: int,
    duplicate: bool,
    shuffle_bank_seed: int,
    batch_games: int,
    output: str | Path | None = None,
) -> dict:
    eval_batch = make_policy_match_batch_fn(learner_policy, opponent_policy)
    diffs = []
    lengths = []
    play_actions = 0.0
    learner_actions = 0.0
    opened_sum = 0.0
    positive_sum = 0.0
    max_steps_sum = 0.0

    for seat in [0, 1] if duplicate else [0]:
        for start in range(0, games, batch_games):
            count = min(batch_games, games - start)
            orders = make_shuffle_bank(shuffle_bank_seed + start, count)
            batch = eval_batch(jnp.asarray(orders, dtype=jnp.int8), seat)
            batch = jax.tree_util.tree_map(lambda x: x.block_until_ready(), batch)
            diffs.append(jax.device_get(batch["score_diff"]))
            lengths.append(jax.device_get(batch["game_length"]))
            play_actions += float(jnp.sum(batch["play_action_sum"]))
            learner_actions += float(jnp.sum(batch["learner_action_sum"]))
            opened_sum += float(jnp.sum(batch["opened_colors"]))
            positive_sum += float(jnp.sum(batch["positive_expeditions"]))
            max_steps_sum += float(jnp.sum(batch["max_steps"]))

    import numpy as np

    diff_arr = np.concatenate(diffs).astype(np.float64)
    length_arr = np.concatenate(lengths).astype(np.float64)
    total_games = int(diff_arr.shape[0])
    wins = int(np.sum(diff_arr > 0))
    losses = int(np.sum(diff_arr < 0))
    ties = int(np.sum(diff_arr == 0))
    wilson_low, wilson_high = wilson_interval(wins, total_games)
    quantiles = np.quantile(diff_arr, [0.0, 0.05, 0.25, 0.5, 0.75, 0.95, 1.0])
    diff_std = float(np.std(diff_arr, ddof=1)) if total_games > 1 else 0.0
    diff_se = diff_std / math.sqrt(total_games) if total_games > 0 else 0.0
    hist_counts, hist_edges = np.histogram(
        diff_arr,
        bins=[-400, -300, -200, -100, -50, 0, 50, 100, 200, 300, 400, 600],
    )
    result = {
        "duplicate": duplicate,
        "games": total_games,
        "wins": float(wins),
        "losses": float(losses),
        "ties": float(ties),
        "win_rate": wins / total_games,
        "wilson_low": wilson_low,
        "wilson_high": wilson_high,
        "mean_score_diff": float(np.mean(diff_arr)),
        "score_diff_std": diff_std,
        "score_diff_ci95_low": float(np.mean(diff_arr) - 1.959963984540054 * diff_se),
        "score_diff_ci95_high": float(np.mean(diff_arr) + 1.959963984540054 * diff_se),
        "score_diff_sum": float(np.sum(diff_arr)),
        "score_diff_quantiles": {
            "min": float(quantiles[0]),
            "p05": float(quantiles[1]),
            "p25": float(quantiles[2]),
            "p50": float(quantiles[3]),
            "p75": float(quantiles[4]),
            "p95": float(quantiles[5]),
            "max": float(quantiles[6]),
        },
        "score_diff_histogram": {
            "edges": hist_edges.astype(float).tolist(),
            "counts": hist_counts.astype(int).tolist(),
        },
        "mean_game_length": float(np.mean(length_arr)),
        "game_length_quantiles": {
            "min": float(np.min(length_arr)),
            "p50": float(np.quantile(length_arr, 0.50)),
            "p95": float(np.quantile(length_arr, 0.95)),
            "max": float(np.max(length_arr)),
        },
        "max_steps_rate": max_steps_sum / total_games,
        "play_action_rate": play_actions / max(learner_actions, 1.0),
        "opened_colors_per_game": opened_sum / total_games,
        "positive_expeditions_per_game": positive_sum / total_games,
    }
    if output:
        _write_json(Path(output), result)
    return result


def make_policy_match_batch_fn(learner_policy, opponent_policy):
    @jax.jit
    def eval_batch(deck_orders: jax.Array, learner_seat: int):
        learner = jnp.asarray(learner_seat, dtype=jnp.int32)
        opponent = 1 - learner
        env_state = jax.vmap(reset_from_order)(deck_orders)
        keys = jax.random.split(jax.random.PRNGKey(0), deck_orders.shape[0])

        def body(carry, _):
            env, key = carry
            key, learner_key, opponent_key = jax.random.split(key, 3)
            learner_keys = jax.random.split(learner_key, deck_orders.shape[0])
            opponent_keys = jax.random.split(opponent_key, deck_orders.shape[0])
            learner_actions = jax.vmap(learner_policy, in_axes=(0, None, 0))(
                env, learner, learner_keys
            )
            opponent_actions = jax.vmap(opponent_policy, in_axes=(0, None, 0))(
                env, opponent, opponent_keys
            )
            learner_turn = env.to_move.astype(jnp.int32) == learner
            active = ~env.done
            actions = jnp.where(learner_turn, learner_actions, opponent_actions)
            next_env, _, _ = jax.vmap(step, in_axes=(0, 0))(env, actions)
            place_type = (actions % 12) // 6
            return (next_env, key), (
                (place_type == PLAY) & learner_turn & active,
                learner_turn & active,
            )

        (final_env, _), (play_actions, learner_actions) = jax.lax.scan(
            body, (env_state, keys[0]), xs=None, length=MAX_STEPS
        )
        scores = jax.vmap(board_score)(final_env)
        diff = scores[:, learner] - scores[:, opponent]
        color_scores = jax.vmap(color_scores_for_player, in_axes=(0, None))(final_env, learner)
        opened = jnp.sum(final_env.col_len[:, learner, :] > 0, axis=-1)
        positive = jnp.sum(color_scores > 0, axis=-1)
        return MatchBatch(
            score_diff=diff,
            game_length=final_env.step_count,
            play_action_sum=jnp.sum(play_actions, axis=0),
            learner_action_sum=jnp.sum(learner_actions, axis=0),
            opened_colors=opened,
            positive_expeditions=positive,
            max_steps=final_env.step_count >= MAX_STEPS,
        )._asdict()

    return eval_batch


def _episode_end(
    next_env: State,
    learners: jax.Array,
    done: jax.Array,
    episode_return: jax.Array,
) -> EpisodeEnd:
    """Snapshot the finished-episode stats of ``next_env`` before it is reset."""

    learners = jnp.broadcast_to(learners.astype(jnp.int32), done.shape)
    color_scores = jax.vmap(color_scores_for_player, in_axes=(0, 0))(next_env, learners)
    batch_idx = jnp.arange(done.shape[0])
    return EpisodeEnd(
        done=done,
        episode_return=episode_return,
        length=next_env.step_count,
        opened_colors=jnp.sum(next_env.col_len[batch_idx, learners, :] > 0, axis=-1),
        positive_expeditions=jnp.sum(color_scores > 0, axis=-1),
        hit_max_steps=next_env.step_count >= MAX_STEPS,
    )


def episode_metrics(transitions: Transition, episodes: EpisodeEnd) -> dict[str, jax.Array]:
    """Aggregate over episodes that actually finished inside the rollout.

    Averaging over ``final_env`` instead would sample envs mid-episode, and
    summing rewards along the scan axis would add up several episodes per slot.
    """

    done = episodes.done.astype(jnp.float32)
    completed = jnp.sum(done)
    denom = jnp.maximum(completed, 1.0)
    actor_count = jnp.sum(transitions.actor_mask)

    returns = episodes.episode_return
    return_mean = jnp.sum(returns * done) / denom
    return_var = jnp.sum(((returns - return_mean) ** 2) * done) / denom

    return {
        "return_mean": return_mean,
        "return_std": jnp.sqrt(jnp.maximum(return_var, 0.0)),
        "game_length_mean": jnp.sum(episodes.length.astype(jnp.float32) * done) / denom,
        "game_length_max": jnp.max(jnp.where(episodes.done, episodes.length, 0)),
        "max_steps_rate": jnp.sum(episodes.hit_max_steps.astype(jnp.float32) * done) / denom,
        "play_action_rate": jnp.sum(transitions.play_action) / jnp.maximum(actor_count, 1),
        "opened_colors_mean": jnp.sum(episodes.opened_colors.astype(jnp.float32) * done) / denom,
        "positive_expeditions_mean": (
            jnp.sum(episodes.positive_expeditions.astype(jnp.float32) * done) / denom
        ),
        "entropy_mean": masked_mean(
            transitions.entropy, transitions.actor_mask.astype(jnp.float32)
        ),
        "active_steps": jnp.sum(transitions.active),
        "learner_actions": actor_count,
        "episodes_completed": completed,
    }


def color_scores_for_player(state: State, player: jax.Array) -> jax.Array:
    card_ids = jnp.arange(N_CARDS, dtype=jnp.int32)
    colors = card_ids // CARDS_PER_COLOR
    slots = card_ids % CARDS_PER_COLOR
    ranks = jnp.where(slots >= 3, slots - 1, 0)
    board_loc = LOC_P0_BOARD + player.astype(jnp.int32)
    on_board = state.card_loc.astype(jnp.int32) == board_loc
    rank_sum = []
    for color in range(N_COLORS):
        rank_sum.append(jnp.sum(jnp.where(on_board & (colors == color), ranks, 0)))
    rank_sum = jnp.stack(rank_sum).astype(jnp.int32)
    length = state.col_len[player.astype(jnp.int32)].astype(jnp.int32)
    handshakes = state.col_hs[player.astype(jnp.int32)].astype(jnp.int32)
    base = (rank_sum - 20) * (1 + handshakes)
    bonus = jnp.where(length >= 8, 20, 0)
    return jnp.where(length > 0, base + bonus, 0).astype(jnp.float32)


def batch_score_diff(state: State, learner: jax.Array) -> jax.Array:
    scores = jax.vmap(board_score)(state)
    opponent = 1 - learner.astype(jnp.int32)
    return scores[:, learner.astype(jnp.int32)] - scores[:, opponent]


def batch_score_diff_for_players(state: State, learners: jax.Array) -> jax.Array:
    scores = jax.vmap(board_score)(state)
    learners = learners.astype(jnp.int32)
    opponents = 1 - learners
    batch_idx = jnp.arange(learners.shape[0])
    return scores[batch_idx, learners] - scores[batch_idx, opponents]


def mask_logits(logits: jax.Array, legal_mask: jax.Array) -> jax.Array:
    return jnp.where(legal_mask, logits, NEG_INF)


def action_log_prob(masked_logits: jax.Array, action: jax.Array) -> jax.Array:
    log_probs = jax.nn.log_softmax(masked_logits, axis=-1)
    return jnp.take_along_axis(log_probs, action[:, None], axis=-1)[:, 0]


def categorical_entropy(masked_logits: jax.Array) -> jax.Array:
    log_probs = jax.nn.log_softmax(masked_logits, axis=-1)
    probs = jnp.exp(log_probs)
    return -jnp.sum(jnp.where(probs > 0, probs * log_probs, 0.0), axis=-1)


def masked_mean(values: jax.Array, weights: jax.Array) -> jax.Array:
    return jnp.sum(values * weights) / jnp.maximum(jnp.sum(weights), 1.0)


def _broadcast_done(done: jax.Array, leaf: jax.Array) -> jax.Array:
    shape = done.shape + (1,) * (leaf.ndim - 1)
    return done.reshape(shape)


def _flatten_transitions(transitions: Transition) -> Transition:
    return jax.tree_util.tree_map(lambda x: x.reshape((-1, *x.shape[2:])), transitions)


def shaping_coefficient(cfg: JaxPPOConfig, learner_actions: int) -> float:
    reward_cfg = cfg.reward
    if reward_cfg.potential_shaping_anneal_steps <= 0:
        return reward_cfg.potential_shaping_final
    progress = min(learner_actions / reward_cfg.potential_shaping_anneal_steps, 1.0)
    return reward_cfg.potential_shaping_initial + progress * (
        reward_cfg.potential_shaping_final - reward_cfg.potential_shaping_initial
    )


def make_shuffle_bank(seed: int, games: int) -> list[list[int]]:
    import numpy as np

    rng = np.random.default_rng(seed)
    return [rng.permutation(N_CARDS).astype("int8").tolist() for _ in range(games)]


def wilson_interval(wins: int, games: int, z: float = 1.959963984540054) -> tuple[float, float]:
    if games == 0:
        return 0.0, 0.0
    p = wins / games
    denom = 1.0 + z * z / games
    center = (p + z * z / (2 * games)) / denom
    margin = z * math.sqrt((p * (1 - p) + z * z / (4 * games)) / games) / denom
    return max(0.0, center - margin), min(1.0, center + margin)


def save_checkpoint(path: Path, state: TrainState, cfg: JaxPPOConfig) -> None:
    del cfg
    path = path.resolve()
    if path.exists():
        shutil.rmtree(path)
    checkpointer = ocp.PyTreeCheckpointer()
    checkpointer.save(path, state, force=True)


def restore_checkpoint(path: Path, state: TrainState) -> TrainState:
    # Explicit restore_args keep per-leaf shardings; the bare item= form loses
    # them on this orbax version and fails with "sharding ... Got None".
    restore_args = ocp.checkpoint_utils.construct_restore_args(state)
    return ocp.PyTreeCheckpointer().restore(path.resolve(), item=state, restore_args=restore_args)


def cli_main(argv: list[str] | None = None) -> None:
    parser = argparse.ArgumentParser(prog="lost-cities-jax-ppo")
    sub = parser.add_subparsers(dest="command", required=True)

    smoke = sub.add_parser("rollout-smoke")
    smoke.add_argument("--config", required=True)
    smoke.add_argument("--output")
    smoke.add_argument("--set", action="append", default=[])

    train_parser = sub.add_parser("train")
    train_parser.add_argument("--config", required=True)
    train_parser.add_argument("--resume")
    train_parser.add_argument("--set", action="append", default=[])

    exploiter_parser = sub.add_parser("train-exploiter")
    exploiter_parser.add_argument("--config", required=True)
    exploiter_parser.add_argument("--opponent-config", required=True)
    exploiter_parser.add_argument("--opponent-checkpoint", required=True)
    exploiter_parser.add_argument("--resume")
    exploiter_parser.add_argument("--set", action="append", default=[])

    eval_parser = sub.add_parser("eval")
    eval_parser.add_argument("--config", required=True)
    eval_parser.add_argument("--checkpoint", required=True)
    eval_parser.add_argument("--opponent")
    eval_parser.add_argument("--games", type=int)
    eval_parser.add_argument("--duplicate", action="store_true")
    eval_parser.add_argument("--output")
    eval_parser.add_argument("--set", action="append", default=[])

    checkpoint_eval_parser = sub.add_parser("eval-checkpoint")
    checkpoint_eval_parser.add_argument("--config", required=True)
    checkpoint_eval_parser.add_argument("--checkpoint", required=True)
    checkpoint_eval_parser.add_argument("--opponent-config", required=True)
    checkpoint_eval_parser.add_argument("--opponent-checkpoint", required=True)
    checkpoint_eval_parser.add_argument("--games", type=int)
    checkpoint_eval_parser.add_argument("--duplicate", action="store_true")
    checkpoint_eval_parser.add_argument("--output")
    checkpoint_eval_parser.add_argument("--set", action="append", default=[])

    checkpoint_static_eval_parser = sub.add_parser("eval-checkpoint-static")
    checkpoint_static_eval_parser.add_argument("--config", required=True)
    checkpoint_static_eval_parser.add_argument("--checkpoint", required=True)
    checkpoint_static_eval_parser.add_argument("--opponent", required=True)
    checkpoint_static_eval_parser.add_argument("--games", type=int)
    checkpoint_static_eval_parser.add_argument("--duplicate", action="store_true")
    checkpoint_static_eval_parser.add_argument("--output")
    checkpoint_static_eval_parser.add_argument("--set", action="append", default=[])

    static_mirror_parser = sub.add_parser("eval-static-mirror")
    static_mirror_parser.add_argument("--policy", required=True)
    static_mirror_parser.add_argument("--games", type=int, default=10_000)
    static_mirror_parser.add_argument("--batch-games", type=int, default=8192)
    static_mirror_parser.add_argument("--shuffle-bank-seed", type=int, default=20260704)
    static_mirror_parser.add_argument("--duplicate", action="store_true")
    static_mirror_parser.add_argument("--output")

    static_match_parser = sub.add_parser("eval-static-match")
    static_match_parser.add_argument("--learner", required=True)
    static_match_parser.add_argument("--opponent", required=True)
    static_match_parser.add_argument("--games", type=int, default=10_000)
    static_match_parser.add_argument("--batch-games", type=int, default=8192)
    static_match_parser.add_argument("--shuffle-bank-seed", type=int, default=20260704)
    static_match_parser.add_argument("--duplicate", action="store_true")
    static_match_parser.add_argument("--output")

    play_parser = sub.add_parser("play")
    play_parser.add_argument("--checkpoint", required=True)
    play_parser.add_argument("--config")
    play_parser.add_argument("--seat", type=int, choices=[0, 1], default=0)
    play_parser.add_argument("--seed", type=int, default=20260706)
    play_parser.add_argument("--duplicate", action="store_true")
    play_parser.add_argument("--log-dir")
    play_parser.add_argument("--bank-state")

    human_play_parser = sub.add_parser("human-play")
    human_play_sub = human_play_parser.add_subparsers(dest="human_play_command", required=True)
    human_summary_parser = human_play_sub.add_parser("summarize")
    human_summary_parser.add_argument("--log-dir")

    league_parser = sub.add_parser("league")
    league_sub = league_parser.add_subparsers(dest="league_command", required=True)
    league_run_parser = league_sub.add_parser("run")
    league_run_parser.add_argument("--config", required=True)

    gates_parser = sub.add_parser("gates")
    gates_sub = gates_parser.add_subparsers(dest="gates_command", required=True)
    gates_run_parser = gates_sub.add_parser("run")
    gates_run_parser.add_argument("--config", required=True)
    gates_repair_parser = gates_sub.add_parser("repair")
    gates_repair_parser.add_argument("--config", required=True)
    gates_repair_parser.add_argument("--run-dir", required=True)

    args = parser.parse_args(argv)
    cfg = (
        load_config(args.config, overrides=parse_overrides(args.set))
        if hasattr(args, "config") and args.command not in {"league", "gates", "play", "human-play"}
        else None
    )
    if args.command == "rollout-smoke":
        assert cfg is not None
        row = random_rollout(cfg)
        print(json.dumps(row, indent=2, sort_keys=True))
        if args.output:
            _write_json(Path(args.output), row)
    elif args.command == "train":
        assert cfg is not None
        run_dir = train(cfg, resume=args.resume)
        print(run_dir)
    elif args.command == "train-exploiter":
        assert cfg is not None
        opponent_cfg = load_config(args.opponent_config)
        run_dir = train_against_checkpoint(
            cfg,
            opponent_cfg,
            args.opponent_checkpoint,
            resume=args.resume,
        )
        print(run_dir)
    elif args.command == "eval":
        assert cfg is not None
        result = evaluate(
            cfg,
            args.checkpoint,
            opponent_name=args.opponent,
            games=args.games,
            duplicate=args.duplicate,
            output=args.output,
        )
        print(json.dumps(result, indent=2, sort_keys=True))
    elif args.command == "eval-checkpoint":
        assert cfg is not None
        opponent_cfg = load_config(args.opponent_config)
        result = evaluate_checkpoint_match(
            cfg,
            args.checkpoint,
            opponent_cfg,
            args.opponent_checkpoint,
            games=args.games,
            duplicate=args.duplicate,
            output=args.output,
        )
        print(json.dumps(result, indent=2, sort_keys=True))
    elif args.command == "eval-checkpoint-static":
        assert cfg is not None
        result = evaluate_checkpoint_vs_static(
            cfg,
            args.checkpoint,
            args.opponent,
            games=args.games,
            duplicate=args.duplicate,
            output=args.output,
        )
        print(json.dumps(result, indent=2, sort_keys=True))
    elif args.command == "eval-static-mirror":
        result = evaluate_static_mirror(
            args.policy,
            games=args.games,
            duplicate=args.duplicate,
            shuffle_bank_seed=args.shuffle_bank_seed,
            batch_games=args.batch_games,
            output=args.output,
        )
        print(json.dumps(result, indent=2, sort_keys=True))
    elif args.command == "eval-static-match":
        result = evaluate_static_match(
            args.learner,
            args.opponent,
            games=args.games,
            duplicate=args.duplicate,
            shuffle_bank_seed=args.shuffle_bank_seed,
            batch_games=args.batch_games,
            output=args.output,
        )
        print(json.dumps(result, indent=2, sort_keys=True))
    elif args.command == "play":
        from lost_cities_jax.human_play import (
            DEFAULT_HUMAN_BANK_STATE,
            DEFAULT_HUMAN_LOG_DIR,
            play_cli,
        )

        play_cli(
            checkpoint=args.checkpoint,
            config=args.config,
            seat=args.seat,
            seed=args.seed,
            duplicate=args.duplicate,
            log_dir=args.log_dir or DEFAULT_HUMAN_LOG_DIR,
            bank_state=args.bank_state or DEFAULT_HUMAN_BANK_STATE,
        )
    elif args.command == "human-play":
        if args.human_play_command == "summarize":
            from lost_cities_jax.human_play import (
                DEFAULT_HUMAN_LOG_DIR,
                print_summary,
                summarize_logs,
            )

            print_summary(summarize_logs(args.log_dir or DEFAULT_HUMAN_LOG_DIR))
    elif args.command == "league":
        if args.league_command == "run":
            from lost_cities_jax.league import run_league

            run_dir = run_league(args.config)
            print(run_dir)
    elif args.command == "gates":
        if args.gates_command == "run":
            from lost_cities_jax.gates import run_gates

            run_dir = run_gates(args.config)
            print(run_dir)
        elif args.gates_command == "repair":
            from lost_cities_jax.gates import repair_gates

            run_dir = repair_gates(args.config, args.run_dir)
            print(run_dir)


def _create_run_dir(cfg: JaxPPOConfig) -> Path:
    timestamp = time.strftime("%Y-%m-%d_%H%M%S")
    root = Path(cfg.run.artifact_root).resolve()
    run_dir = root / f"{timestamp}_{_slugify(cfg.run.experiment_name)}"
    run_dir.mkdir(parents=True, exist_ok=False)
    return run_dir


def _metrics_to_row(
    metrics: dict[str, jax.Array], update: int, shaping_coef: float, cfg: JaxPPOConfig
) -> dict:
    row = {
        key: float(value) if getattr(value, "shape", ()) == () else value.tolist()
        for key, value in metrics.items()
    }
    row["update"] = update
    row["shaping_coef"] = float(shaping_coef)
    row["env_steps"] = int((update + 1) * cfg.ppo.batch_games * cfg.ppo.rollout_steps)
    row["opponent"] = cfg.opponent.name
    return row


def _write_json(path: Path, data: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(data, indent=2, sort_keys=True) + "\n", encoding="utf-8")


def _append_jsonl(path: Path, data: dict) -> None:
    with path.open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(data, sort_keys=True) + "\n")


def _sum_dicts(left: dict, right: dict) -> dict:
    return {key: left.get(key, 0.0) + right.get(key, 0.0) for key in set(left) | set(right)}


def _deep_merge(base: dict, override: dict) -> dict:
    merged = dict(base)
    for key, value in override.items():
        if isinstance(value, dict) and isinstance(merged.get(key), dict):
            merged[key] = _deep_merge(merged[key], value)
        else:
            merged[key] = value
    return merged


def parse_overrides(items: list[str]) -> dict:
    root: dict = {}
    for item in items:
        if "=" not in item:
            msg = f"--set must be PATH=VALUE, got: {item}"
            raise ValueError(msg)
        path, raw_value = item.split("=", 1)
        value = yaml.safe_load(raw_value)
        cursor = root
        parts = path.split(".")
        for part in parts[:-1]:
            cursor = cursor.setdefault(part, {})
        cursor[parts[-1]] = value
    return root


def _slugify(value: str) -> str:
    return "".join(ch.lower() if ch.isalnum() else "-" for ch in value).strip("-")


__all__ = [
    "ActorCritic",
    "JaxPPOConfig",
    "LeagueAssignments",
    "TrainState",
    "checkpoint_policy",
    "checkpoint_policy_from_params",
    "cli_main",
    "create_train_state",
    "evaluate",
    "evaluate_checkpoint_match",
    "evaluate_checkpoint_vs_static",
    "evaluate_policy_match",
    "evaluate_static_match",
    "evaluate_static_mirror",
    "load_config",
    "make_league_train_iteration",
    "random_rollout",
    "sample_league_assignments",
    "train",
    "train_against_checkpoint",
]
