coorl-lost-cities/src/lost_cities_jax/match.py

"""Classic Lost Cities: three rounds, scores summed, highest total wins.

The single-round engine in ``engine.py`` is left alone -- it is the rules oracle
the TypeScript client is differential-tested against. This module wraps it.

Two rules only exist at the match level, both from the Kosmos rulebook:

- "If after three games you have the highest overall score, you win."
- "The player who has more points begins" the next game. Not alternating. The
  rulebook says nothing about an exact tie, so we flip a coin: any deterministic
  tie-break would hand one seat a standing edge in symmetric self-play, and the
  agent would learn to steer for it.

All randomness is drawn in ``match_reset`` and stored in the state, so
``match_step`` stays deterministic and needs no PRNG key. That keeps the step
signature clean through every rollout/eval body, and it means a mirrored match
(same three deals, seats swapped, same coin flips) is just a seat relabel.
"""

from __future__ import annotations

from typing import NamedTuple

import jax
import jax.numpy as jnp

from lost_cities_jax.engine import board_score, legal_action_mask, reset_from_order, step
from lost_cities_jax.types import N_CARDS, State

N_ROUNDS = 3


class MatchState(NamedTuple):
    round: State
    """The round being played right now."""

    deck_orders: jax.Array
    """(N_ROUNDS, N_CARDS) -- every deal of the match, shuffled up front."""

    coin_flips: jax.Array
    """(N_ROUNDS,) tie-breaking starters, used only when the scores are level."""

    round_idx: jax.Array
    carry: jax.Array
    """(2,) points banked by each player in the rounds already finished."""

    done: jax.Array


def match_reset(rng: jax.Array) -> MatchState:
    deck_key, coin_key = jax.random.split(rng)
    deck_keys = jax.random.split(deck_key, N_ROUNDS)
    deck_orders = jax.vmap(
        lambda key: jax.random.permutation(key, jnp.arange(N_CARDS, dtype=jnp.int32), axis=0)
    )(deck_keys).astype(jnp.int8)
    coin_flips = jax.random.bernoulli(coin_key, 0.5, shape=(N_ROUNDS,)).astype(jnp.int8)
    return match_reset_from(deck_orders, coin_flips)


def match_reset_from(deck_orders: jax.Array, coin_flips: jax.Array) -> MatchState:
    """Build a match from explicit deals and coin flips (mirrored eval, fixtures)."""

    deck_orders = jnp.asarray(deck_orders, dtype=jnp.int8)
    coin_flips = jnp.asarray(coin_flips, dtype=jnp.int8)
    carry = jnp.zeros((2,), dtype=jnp.int32)
    round_idx = jnp.asarray(0, dtype=jnp.int32)
    return MatchState(
        round=reset_from_order(deck_orders[0], first_player=starting_player(carry, 0, coin_flips)),
        deck_orders=deck_orders,
        coin_flips=coin_flips,
        round_idx=round_idx,
        carry=carry,
        done=jnp.asarray(False),
    )


def starting_player(carry: jax.Array, round_idx: jax.Array, coin_flips: jax.Array) -> jax.Array:
    """Whoever has banked more points leads; level scores fall back to the coin.

    Round one needs no special case: ``carry`` is (0, 0) there, so the tie branch
    already picks the coin flip, which is exactly the rulebook's arbitrary
    "oldest player begins".
    """

    lead = carry[0] - carry[1]
    coin = coin_flips[round_idx].astype(jnp.int8)
    ahead = jnp.where(lead > 0, jnp.int8(0), jnp.int8(1))
    return jnp.where(lead == 0, coin, ahead)


def match_score(state: MatchState) -> jax.Array:
    """(2,) running totals: rounds already banked plus the board in play."""

    return state.carry + board_score(state.round).astype(jnp.int32)


def match_legal_action_mask(state: MatchState) -> jax.Array:
    return legal_action_mask(state.round) & ~state.done


def match_step(state: MatchState, action: jax.Array) -> tuple[MatchState, jax.Array, jax.Array]:
    """Play one ply. Rolls into the next round when the deck runs out.

    Reward is zero until the third round ends, then it is each player's match
    total. Intermediate rounds pay nothing: only the sum decides the match.
    """

    played, _, _ = step(state.round, action)
    round_over = played.done & ~state.done
    is_final_round = state.round_idx >= N_ROUNDS - 1

    advance = round_over & ~is_final_round
    finished = round_over & is_final_round

    banked = state.carry + board_score(played).astype(jnp.int32)
    next_carry = jnp.where(advance, banked, state.carry)
    next_round_idx = jnp.where(advance, state.round_idx + 1, state.round_idx)

    # Dealt eagerly every ply; only kept on the plies that actually roll over.
    fresh = reset_from_order(
        state.deck_orders[next_round_idx],
        first_player=starting_player(next_carry, next_round_idx, state.coin_flips),
    )
    next_round = jax.tree_util.tree_map(
        lambda new, old: jnp.where(advance, new, old), fresh, played
    )

    next_state = MatchState(
        round=next_round,
        deck_orders=state.deck_orders,
        coin_flips=state.coin_flips,
        round_idx=next_round_idx,
        carry=next_carry,
        done=state.done | finished,
    )

    total = state.carry + board_score(played).astype(jnp.int32)
    reward = jnp.where(finished, total.astype(jnp.float32), jnp.zeros((2,), dtype=jnp.float32))
    return next_state, reward, next_state.done