import jax
import jax.numpy as jnp
import jax.random as jrandom
from generals.core.env import GeneralsEnv

from general_bots_training.network import PolicyValueNetwork
from general_bots_training.opponents import SelfPlayOpponent
from general_bots_training.rollout import make_collect_rollout


def test_self_play_collects_both_player_perspectives():
    key = jrandom.PRNGKey(0)
    key, network_key, pool_key, state_key = jrandom.split(key, 4)
    env = GeneralsEnv(grid_dims=(4, 4), truncation=20, pool_size=8)
    pool, _ = env.reset(pool_key)
    states = jax.vmap(env.init_state)(jrandom.split(state_key, 2))
    states = states._replace(pool_idx=jnp.arange(2, dtype=states.pool_idx.dtype))
    network = PolicyValueNetwork(network_key)

    _, transitions, (_, last_next_obs) = make_collect_rollout(env, 1, SelfPlayOpponent())(
        states, pool, network, key
    )

    assert transitions["obs"].shape == (1, 4, 14, 4, 4)
    assert transitions["action"].shape == (1, 4, 5)
    assert transitions["player"][0].tolist() == [0, 0, 1, 1]
    assert transitions["done"].shape == (1, 4)
    assert last_next_obs.shape == (4, 14, 4, 4)