coorl-lost-cities/tests/lost_cities_jax/test_ppo_stack.py

from __future__ import annotations

from pathlib import Path

import jax
import jax.numpy as jnp
import numpy as np
import pytest

from lost_cities_jax import legal_action_mask, reset
from lost_cities_jax.engine import current_hand_sorted, decode_action, reset_from_order
from lost_cities_jax.opponents import (
    HeuristicExpertConfig,
    discard_only_action,
    heuristic_balanced_action,
    heuristic_cautious_action,
    heuristic_expert_action,
    make_heuristic_expert_policy,
)
from lost_cities_jax.ppo import (
    JaxPPOConfig,
    NetworkConfig,
    OpponentConfig,
    PPOHyperConfig,
    RunConfig,
    create_train_state,
    evaluate,
    evaluate_checkpoint_match,
    evaluate_static_mirror,
    load_config,
    make_train_iteration,
    policy_by_name,
    random_rollout,
    train,
)
from lost_cities_jax.types import (
    CARDS_PER_COLOR,
    HAND_SIZE,
    LOC_DECK,
    LOC_P0_BOARD,
    LOC_P0_HAND,
    N_CARDS,
    N_COLORS,
    PLAY,
)

GATE3_CHECKPOINT = (
    "/mnt/2tbhdd/coolrl-lost-cities-artifacts/jax-ppo-static-opponents/"
    "2026-07-04_230150_jax-ppo-cautious/latest"
)


def tiny_config(tmp_path) -> JaxPPOConfig:
    return JaxPPOConfig(
        run=RunConfig(
            experiment_name="pytest-jax-ppo",
            seed=123,
            total_updates=1,
            checkpoint_every=1,
            artifact_root=str(tmp_path),
        ),
        opponent=OpponentConfig(name="discard_only"),
        network=NetworkConfig(hidden_size=32, num_layers=1),
        # A round runs at least 44 plies (one per deck draw), so a shorter
        # rollout would finish no episodes and leave the episode metrics empty.
        ppo=PPOHyperConfig(batch_games=8, rollout_steps=80, epochs=1, minibatches=2),
    )


def test_static_opponents_return_legal_actions():
    state = reset(jax.random.PRNGKey(0))
    mask = np.asarray(legal_action_mask(state), dtype=bool)
    for fn in [
        discard_only_action,
        heuristic_balanced_action,
        heuristic_cautious_action,
        heuristic_expert_action,
    ]:
        action = int(fn(state, jnp.int32(0), jax.random.PRNGKey(1)))
        assert mask[action]


def test_cautious_discards_weak_unopened_hand():
    low_cards = [_rank_card(0, rank) for rank in range(2, 7)]
    state = _manual_state(p0_hand=low_cards, to_move=0)
    action = int(heuristic_cautious_action(state, jnp.int32(0), jax.random.PRNGKey(1)))
    _, place_type, _ = [int(x) for x in decode_action(jnp.asarray(action))]
    assert place_type != PLAY


def test_balanced_opens_strong_unopened_rank():
    strong = _rank_card(0, 10)
    state = _manual_state(p0_hand=[strong, _rank_card(1, 2), _rank_card(2, 3)], to_move=0)
    action = int(heuristic_balanced_action(state, jnp.int32(0), jax.random.PRNGKey(1)))
    hand_slot, place_type, _ = [int(x) for x in decode_action(jnp.asarray(action))]
    hand = [int(x) for x in current_hand_sorted(state, 0)]
    assert place_type == PLAY
    assert hand[hand_slot] == strong


def test_cautious_uses_own_board_when_playing_as_p1():
    existing = _rank_card(0, 5)
    playable = _rank_card(0, 6)
    state = _manual_state(
        p1_hand=[
            playable,
            _rank_card(1, 2),
            _rank_card(1, 3),
            _rank_card(2, 2),
            _rank_card(2, 3),
            _rank_card(3, 2),
            _rank_card(3, 3),
            _rank_card(4, 2),
        ],
        p1_board=[existing],
        to_move=1,
    )
    action = int(heuristic_cautious_action(state, jnp.int32(1), jax.random.PRNGKey(1)))
    hand_slot, place_type, _ = [int(x) for x in decode_action(jnp.asarray(action))]
    hand = [int(x) for x in current_hand_sorted(state, 1)]
    assert place_type == PLAY
    assert hand[hand_slot] == playable


def test_expert_opens_strong_ev_hand():
    strong_cards = [
        _hs_card(0, 0),
        _rank_card(0, 7),
        _rank_card(0, 8),
        _rank_card(0, 9),
        _rank_card(0, 10),
        _rank_card(1, 2),
        _rank_card(2, 2),
        _rank_card(3, 2),
    ]
    state = _manual_state(p0_hand=strong_cards, to_move=0)
    action = int(heuristic_expert_action(state, jnp.int32(0), jax.random.PRNGKey(1)))
    hand_slot, place_type, _ = [int(x) for x in decode_action(jnp.asarray(action))]
    hand = [int(x) for x in current_hand_sorted(state, 0)]
    assert place_type == PLAY
    assert hand[hand_slot] // CARDS_PER_COLOR == 0


def test_expert_rejects_weak_new_open():
    weak_cards = [
        _rank_card(0, 2),
        _rank_card(0, 3),
        _rank_card(0, 4),
        _rank_card(1, 2),
        _rank_card(2, 2),
        _rank_card(3, 2),
        _rank_card(4, 2),
        _rank_card(4, 3),
    ]
    state = _manual_state(p0_hand=weak_cards, to_move=0)
    action = int(heuristic_expert_action(state, jnp.int32(0), jax.random.PRNGKey(1)))
    _, place_type, _ = [int(x) for x in decode_action(jnp.asarray(action))]
    assert place_type != PLAY


def test_expert_avoids_discarding_immediately_useful_opponent_card():
    dangerous = _rank_card(0, 8)
    safe = _rank_card(4, 2)
    state = _manual_state(
        p0_hand=[
            dangerous,
            safe,
            _rank_card(1, 2),
            _rank_card(1, 3),
            _rank_card(2, 2),
            _rank_card(2, 3),
            _rank_card(3, 2),
            _rank_card(3, 3),
        ],
        p1_board=[_rank_card(0, 7)],
        to_move=0,
    )
    action = int(heuristic_expert_action(state, jnp.int32(0), jax.random.PRNGKey(1)))
    hand_slot, place_type, _ = [int(x) for x in decode_action(jnp.asarray(action))]
    hand = [int(x) for x in current_hand_sorted(state, 0)]
    if place_type != PLAY:
        assert hand[hand_slot] != dangerous


def test_expert_cap2_blocks_third_open_color():
    state = _manual_state(
        p0_hand=[
            _hs_card(2),
            _rank_card(2, 7),
            _rank_card(2, 8),
            _rank_card(2, 9),
            _rank_card(2, 10),
            _rank_card(3, 2),
            _rank_card(3, 3),
            _rank_card(4, 2),
        ],
        p0_board=[_rank_card(0, 2), _rank_card(1, 2)],
        to_move=0,
    )
    policy = make_heuristic_expert_policy(HeuristicExpertConfig(max_open_colors=2))
    action = int(policy(state, jnp.int32(0), jax.random.PRNGKey(1)))
    hand_slot, place_type, _ = [int(x) for x in decode_action(jnp.asarray(action))]
    hand = [int(x) for x in current_hand_sorted(state, 0)]
    opened_new_color = (
        place_type == PLAY and state.col_len[0, hand[hand_slot] // CARDS_PER_COLOR] == 0
    )
    assert not bool(opened_new_color)


def test_expert_cap3_allows_third_but_blocks_fourth_open_color():
    cap3 = make_heuristic_expert_policy(HeuristicExpertConfig(max_open_colors=3))
    third_state = _manual_state(
        p0_hand=[
            _hs_card(2),
            _rank_card(2, 7),
            _rank_card(2, 8),
            _rank_card(2, 9),
            _rank_card(2, 10),
            _rank_card(3, 2),
            _rank_card(3, 3),
            _rank_card(4, 2),
        ],
        p0_board=[_rank_card(0, 2), _rank_card(1, 2)],
        to_move=0,
    )
    action = int(cap3(third_state, jnp.int32(0), jax.random.PRNGKey(1)))
    hand_slot, place_type, _ = [int(x) for x in decode_action(jnp.asarray(action))]
    hand = [int(x) for x in current_hand_sorted(third_state, 0)]
    assert place_type == PLAY
    assert third_state.col_len[0, hand[hand_slot] // CARDS_PER_COLOR] == 0

    fourth_state = _manual_state(
        p0_hand=[
            _hs_card(3),
            _rank_card(3, 7),
            _rank_card(3, 8),
            _rank_card(3, 9),
            _rank_card(3, 10),
            _rank_card(4, 2),
            _rank_card(4, 3),
            _rank_card(4, 4),
        ],
        p0_board=[_rank_card(0, 2), _rank_card(1, 2), _rank_card(2, 2)],
        to_move=0,
    )
    action = int(cap3(fourth_state, jnp.int32(0), jax.random.PRNGKey(1)))
    hand_slot, place_type, _ = [int(x) for x in decode_action(jnp.asarray(action))]
    hand = [int(x) for x in current_hand_sorted(fourth_state, 0)]
    opened_new_color = (
        place_type == PLAY and fourth_state.col_len[0, hand[hand_slot] // CARDS_PER_COLOR] == 0
    )
    assert not bool(opened_new_color)


@pytest.mark.parametrize(
    "policy_name", ["discard_only", "heuristic_balanced", "heuristic_cautious"]
)
def test_static_policy_duplicate_mirror_score_diff_is_zero(policy_name):
    result = evaluate_static_mirror(
        policy_name,
        games=16,
        duplicate=True,
        shuffle_bank_seed=20260704,
        batch_games=16,
    )
    assert result["games"] == 32
    assert result["mean_score_diff"] == 0.0
    assert result["wins"] == result["losses"]


def test_expert_duplicate_mirror_regression_stats():
    result = evaluate_static_mirror(
        "heuristic_expert",
        games=1000,
        duplicate=True,
        shuffle_bank_seed=20260704,
        batch_games=1000,
    )
    assert result["games"] == 2000
    assert abs(result["mean_score_diff"]) <= 1.0e-6
    assert result["wins"] == result["losses"]
    assert result["max_steps_rate"] < 0.05
    assert 2.0 <= result["opened_colors_per_game"] <= 3.5
    assert 0.30 <= result["play_action_rate"] <= 0.60
    assert result["mean_game_length"] <= 70.0


def test_gate3_checkpoint_duplicate_self_mirror_score_diff_is_zero():
    checkpoint = Path(GATE3_CHECKPOINT)
    if not checkpoint.exists():
        pytest.skip(f"gate-3 checkpoint not available: {checkpoint}")
    if not any(device.platform == "gpu" for device in jax.local_devices()):
        pytest.skip("gate-3 checkpoint was saved with CUDA sharding; run under CUDA JAX")
    cfg = load_config("configs/jax_ppo/cautious.yaml")
    result = evaluate_checkpoint_match(
        cfg,
        checkpoint,
        cfg,
        checkpoint,
        games=4,
        duplicate=True,
    )
    assert result["games"] == 8
    assert result["mean_score_diff"] == 0.0
    assert result["wins"] == result["losses"]


def test_random_rollout_smoke(tmp_path):
    row = random_rollout(tiny_config(tmp_path))
    assert row["env_steps"] == 8 * 80
    assert 0.0 <= row["play_action_rate"] <= 1.0
    # Averaged over episodes that actually finished, so it must clear the
    # 44-ply floor rather than report a mid-episode step count.
    assert row["episodes_completed"] > 0
    assert row["game_length_mean"] >= 44.0


def test_train_checkpoint_and_eval_smoke(tmp_path):
    cfg = tiny_config(tmp_path)
    run_dir = train(cfg)
    assert (run_dir / "latest").exists()
    result = evaluate(cfg, run_dir / "latest", games=8, duplicate=True)
    assert result["games"] == 16
    assert 0.0 <= result["win_rate"] <= 1.0


def test_one_jitted_train_iteration_shapes(tmp_path):
    cfg = tiny_config(tmp_path)
    train_state = create_train_state(cfg, jax.random.PRNGKey(0))
    env_state = jax.jit(jax.vmap(reset))(jax.random.split(jax.random.PRNGKey(1), 8))
    train_iteration = make_train_iteration(cfg, policy_by_name("discard_only"))

    train_state, env_state, rng, metrics = train_iteration(
        train_state, env_state, jax.random.PRNGKey(2), jnp.asarray(1.0)
    )

    assert env_state.to_move.shape == (8,)
    assert "loss" in metrics
    assert jnp.isfinite(metrics["loss"])


def test_load_config_file():
    cfg = load_config("configs/jax_ppo/smoke.yaml")
    assert cfg.opponent.name == "discard_only"
    assert cfg.ppo.batch_games == 64


def _rank_card(color: int, rank: int) -> int:
    return color * CARDS_PER_COLOR + rank + 1


def _hs_card(color: int, slot: int = 0) -> int:
    return color * CARDS_PER_COLOR + slot


def _manual_state(
    *,
    p0_hand: list[int] | None = None,
    p1_hand: list[int] | None = None,
    p0_board: list[int] | None = None,
    p1_board: list[int] | None = None,
    to_move: int = 0,
):
    p0_hand = _fill_hand(p0_hand or [], set((p0_board or []) + (p1_board or [])))
    p1_hand = _fill_hand(p1_hand or [], set(p0_hand + (p0_board or []) + (p1_board or [])))
    rest = [card for card in range(N_CARDS) if card not in set(p0_hand + p1_hand)]
    deck_order = jnp.asarray(p0_hand + p1_hand + rest, dtype=jnp.int8)
    state = reset_from_order(deck_order)
    card_loc = jnp.full((N_CARDS,), LOC_DECK, dtype=jnp.int8)
    card_loc = card_loc.at[jnp.asarray(p0_hand)].set(LOC_P0_HAND)
    card_loc = card_loc.at[jnp.asarray(p1_hand)].set(LOC_P0_HAND + 1)
    if p0_board:
        card_loc = card_loc.at[jnp.asarray(p0_board)].set(LOC_P0_BOARD)
    if p1_board:
        card_loc = card_loc.at[jnp.asarray(p1_board)].set(LOC_P0_BOARD + 1)
    col_top, col_hs, col_len = _columns(p0_board or [], p1_board or [])
    return state._replace(
        card_loc=card_loc,
        col_top=col_top,
        col_hs=col_hs,
        col_len=col_len,
        to_move=jnp.asarray(to_move, dtype=jnp.int8),
    )


def _fill_hand(cards: list[int], reserved: set[int]) -> list[int]:
    result = list(cards)
    blocked = set(result) | reserved
    for card in range(N_CARDS):
        if len(result) == HAND_SIZE:
            return sorted(result)
        if card not in blocked:
            result.append(card)
            blocked.add(card)
    raise AssertionError("could not fill hand")


def _columns(p0_board: list[int], p1_board: list[int]):
    col_top = np.zeros((2, N_COLORS), dtype=np.int8)
    col_hs = np.zeros((2, N_COLORS), dtype=np.int8)
    col_len = np.zeros((2, N_COLORS), dtype=np.int8)
    for player, cards in enumerate([p0_board, p1_board]):
        for card in cards:
            color = card // CARDS_PER_COLOR
            slot = card % CARDS_PER_COLOR
            col_len[player, color] += 1
            if slot < 3:
                col_hs[player, color] += 1
            else:
                col_top[player, color] = max(col_top[player, color], slot - 1)
    return jnp.asarray(col_top), jnp.asarray(col_hs), jnp.asarray(col_len)