import jax.numpy as jnp
import jax.random as jrandom

from general_bots_training.network import PolicyValueNetwork, obs_to_tensor
from general_bots_training.ppo import compute_advantages_and_returns


def test_returns_use_raw_advantages_before_policy_normalization():
    rewards = jnp.array([[1.0, 3.0]])
    values = jnp.zeros_like(rewards)
    next_value = jnp.zeros(2)
    dones = jnp.ones_like(rewards, dtype=bool)

    advantages, returns = compute_advantages_and_returns(rewards, values, next_value, dones)

    assert jnp.allclose(advantages, jnp.array([[-1.0, 1.0]]))
    assert jnp.allclose(returns, rewards)


def test_pass_is_one_global_action_and_round_trips():
    network = PolicyValueNetwork(jrandom.PRNGKey(0))
    obs = jnp.zeros((14, 4, 4))
    mask = jnp.zeros((4, 4, 4), dtype=bool)

    action, _, sampled_logprob, entropy = network(obs, mask, jrandom.PRNGKey(1))
    _, _, evaluated_logprob, _ = network(obs, mask, jrandom.PRNGKey(2), action)

    assert action.tolist() == [1, 0, 0, 0, 0]
    assert jnp.allclose(sampled_logprob, evaluated_logprob)
    assert jnp.isclose(entropy, 0.0)


def test_observation_encoder_includes_normalized_timestep():
    class Observation:
        armies = jnp.zeros((2, 3), dtype=jnp.int32)
        generals = jnp.zeros((2, 3), dtype=bool)
        castles = jnp.zeros((2, 3), dtype=bool)
        mountains = jnp.zeros((2, 3), dtype=bool)
        neutral_cells = jnp.ones((2, 3), dtype=bool)
        owned_cells = jnp.zeros((2, 3), dtype=bool)
        opponent_cells = jnp.zeros((2, 3), dtype=bool)
        fog_cells = jnp.zeros((2, 3), dtype=bool)
        structures_in_fog = jnp.zeros((2, 3), dtype=bool)
        owned_land_count = jnp.array(0)
        owned_army_count = jnp.array(0)
        opponent_land_count = jnp.array(0)
        opponent_army_count = jnp.array(0)
        timestep = jnp.array(1200)

    encoded = obs_to_tensor(Observation())

    assert encoded.shape == (14, 2, 3)
    assert jnp.allclose(encoded[-1], 1.0)