"""Jitted vmap wrappers for the Lost Cities JAX engine."""

from __future__ import annotations

import jax

from lost_cities_jax.engine import legal_action_mask, reset, reset_from_order, step
from lost_cities_jax.obs import observation

batched_reset = jax.jit(jax.vmap(reset))
batched_reset_from_order = jax.jit(jax.vmap(reset_from_order))
batched_step = jax.jit(jax.vmap(step, in_axes=(0, 0)))
batched_legal_mask = jax.jit(jax.vmap(legal_action_mask))
batched_obs = jax.jit(jax.vmap(observation, in_axes=(0, 0)))

__all__ = [
    "batched_reset",
    "batched_reset_from_order",
    "batched_step",
    "batched_legal_mask",
    "batched_obs",
]
