"""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])