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

"""Phase 1: the classic three-round match layer."""

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

from lost_cities_jax.engine import board_score
from lost_cities_jax.match import (
    N_ROUNDS,
    MatchState,
    match_legal_action_mask,
    match_reset,
    match_reset_from,
    match_score,
    match_step,
    starting_player,
)
from lost_cities_jax.opponents import random_legal_action
from lost_cities_jax.types import DECK_DRAWS, DRAW_DECK, N_CARDS

MATCH_STEP_CAP = 1400  # 3 rounds x the engine's 400-ply safety cap, with slack


def _decks(seed: int) -> jnp.ndarray:
    rng = np.random.default_rng(seed)
    return jnp.asarray(
        np.stack([rng.permutation(N_CARDS) for _ in range(N_ROUNDS)]), dtype=jnp.int8
    )


def _play_match(state: MatchState, key: jax.Array):
    """Drive a match to completion with random legal play, recording each ply."""
    plies = []
    while not bool(state.done):
        key, step_key = jax.random.split(key)
        action = random_legal_action(state.round, state.round.to_move, step_key)
        assert bool(match_legal_action_mask(state)[action])
        prev = state
        state, reward, _ = match_step(state, action)
        plies.append((prev, int(action), state, np.asarray(reward)))
        assert len(plies) < MATCH_STEP_CAP, "match failed to terminate"
    return state, plies


# --- match structure -------------------------------------------------------


def test_match_plays_exactly_three_rounds_and_then_ends():
    state = match_reset(jax.random.PRNGKey(0))
    final, plies = _play_match(state, jax.random.PRNGKey(1))

    assert bool(final.done)
    assert int(final.round_idx) == N_ROUNDS - 1
    # round_idx advances 0 -> 1 -> 2 and stops.
    seen = {int(before.round_idx) for before, _, _, _ in plies}
    assert seen == {0, 1, 2}


def test_reward_is_paid_only_when_the_third_round_ends():
    state = match_reset(jax.random.PRNGKey(2))
    final, plies = _play_match(state, jax.random.PRNGKey(3))

    rewards = np.stack([reward for _, _, _, reward in plies])
    paid = np.flatnonzero(np.any(rewards != 0.0, axis=1))
    # Intermediate rounds bank points into carry but pay nothing.
    assert paid.tolist() == [len(plies) - 1]
    assert np.allclose(rewards[-1], np.asarray(match_score(final), dtype=np.float32))


def test_each_round_runs_44_deck_draws_plus_one_ply_per_discard_draw():
    """The deck clock: only deck draws end a round, so discard draws extend it."""
    state = match_reset(jax.random.PRNGKey(4))
    _, plies = _play_match(state, jax.random.PRNGKey(5))

    per_round: dict[int, list[int]] = {0: [], 1: [], 2: []}
    for before, action, _, _ in plies:
        per_round[int(before.round_idx)].append(action % 6)

    for round_idx, draws in per_round.items():
        deck_draws = sum(1 for src in draws if src == DRAW_DECK)
        discard_draws = len(draws) - deck_draws
        assert deck_draws == DECK_DRAWS, f"round {round_idx} drew {deck_draws} deck cards"
        assert len(draws) == DECK_DRAWS + discard_draws


# --- carry -----------------------------------------------------------------


def test_carry_banks_each_finished_round_exactly_once():
    state = match_reset(jax.random.PRNGKey(6))
    final, plies = _play_match(state, jax.random.PRNGKey(7))

    deltas = [
        np.asarray(after.carry) - np.asarray(before.carry)
        for before, _, after, _ in plies
        if int(after.round_idx) != int(before.round_idx)
    ]

    # Two roll-overs for three rounds, and carry is exactly their sum.
    assert len(deltas) == N_ROUNDS - 1
    assert np.array_equal(np.asarray(final.carry), sum(deltas))

    # The third round is still on the board, not yet banked.
    board = np.asarray(board_score(final.round)).astype(np.int32)
    assert np.array_equal(np.asarray(match_score(final)), np.asarray(final.carry) + board)


def test_match_score_does_not_jump_across_a_round_boundary():
    """Phi = carry + board diff must be continuous, or PBRS gets a free kick."""
    state = match_reset(jax.random.PRNGKey(8))
    _, plies = _play_match(state, jax.random.PRNGKey(9))

    for before, _, after, _ in plies:
        rolled_over = int(after.round_idx) != int(before.round_idx)
        if not rolled_over:
            continue
        # The finished round's board is folded into carry and the new board is
        # empty, so the running total is unchanged by the roll-over itself.
        assert np.array_equal(np.asarray(after.carry), np.asarray(match_score(after)))
        assert int(np.asarray(board_score(after.round)).sum()) == 0


# --- the start-player rule -------------------------------------------------


@pytest.mark.parametrize(
    ("carry", "coin", "expected"),
    [
        ((60, 10), 1, 0),  # p0 ahead -> p0 leads, coin ignored
        ((10, 60), 0, 1),  # p1 ahead -> p1 leads, coin ignored
        ((30, 30), 0, 0),  # level -> coin
        ((30, 30), 1, 1),  # level -> coin
    ],
)
def test_the_player_with_more_points_begins(carry, coin, expected):
    coin_flips = jnp.asarray([coin, coin, coin], dtype=jnp.int8)
    starter = starting_player(jnp.asarray(carry, dtype=jnp.int32), jnp.int32(1), coin_flips)
    assert int(starter) == expected


def test_round_two_is_actually_led_by_whoever_is_ahead():
    state = match_reset(jax.random.PRNGKey(10))
    _, plies = _play_match(state, jax.random.PRNGKey(11))

    for before, _, after, _ in plies:
        if int(after.round_idx) == int(before.round_idx):
            continue
        carry = np.asarray(after.carry)
        if carry[0] == carry[1]:
            continue  # coin flip, covered above
        leader = 0 if carry[0] > carry[1] else 1
        assert int(after.round.to_move) == leader


def test_round_one_start_player_is_a_fair_coin():
    starts = [int(match_reset(jax.random.PRNGKey(seed)).round.to_move) for seed in range(400)]
    share = sum(starts) / len(starts)
    assert 0.4 < share < 0.6, f"round-1 starter is skewed: {share:.2f}"


# --- mirrored matches ------------------------------------------------------


def test_mirroring_the_seats_negates_the_match_score():
    """Same deals, same coins, seats swapped: the result must flip sign exactly."""
    decks = _decks(21)
    coins = jnp.asarray([0, 1, 0], dtype=jnp.int8)

    # Swapping seats == swapping the two dealt hands and the coin bits.
    swapped = decks.at[:, :16].set(jnp.concatenate([decks[:, 8:16], decks[:, :8]], axis=1))
    flipped_coins = (1 - coins).astype(jnp.int8)

    a, _ = _play_match(match_reset_from(decks, coins), jax.random.PRNGKey(30))
    b, _ = _play_match(match_reset_from(swapped, flipped_coins), jax.random.PRNGKey(30))

    # Both matches see identical information; only the seat labels differ, so the
    # per-seat totals must be each other's mirror image.
    assert np.array_equal(np.asarray(match_score(a)), np.asarray(match_score(b))[::-1])