import jax.numpy as jnp
import numpy as np
from generals.core import game

from general_bots_training.mcts import (
    build_cost_grid,
    competition_step,
    observation_from_wire,
    sample_determinization,
    valid_build_mask,
)


def make_observation(armies=60):
    types = np.ones((5, 5), dtype=np.int32)
    owners = np.zeros((5, 5), dtype=np.int32)
    army_grid = np.zeros((5, 5), dtype=np.int32)
    types[2, 2] = 4
    owners[2, 2] = 1
    army_grid[2, 2] = armies
    types[0, 0] = 0
    return observation_from_wire(100, 1, armies, 1, 20, types, owners, army_grid)


def test_wire_observation_does_not_treat_fog_as_neutral():
    observation = make_observation()

    assert bool(observation.fog_cells[0, 0])
    assert not bool(observation.neutral_cells[0, 0])
    assert bool(observation.neutral_cells[0, 1])


def test_observation_build_cost_and_legality_match_rules():
    observation = make_observation()
    costs = build_cost_grid(observation)

    assert int(costs[2, 2]) == 49
    assert int(costs[2, 3]) == 47
    assert not bool(valid_build_mask(observation)[2, 2])

    armies = observation.armies.at[2, 3].set(47)
    owned = observation.owned_cells.at[2, 3].set(True)
    observation = observation._replace(armies=armies, owned_cells=owned)
    assert bool(valid_build_mask(observation)[2, 3])


def test_determinization_matches_observed_global_totals():
    observation = make_observation()
    state = sample_determinization(observation, 0, np.random.default_rng(0))

    assert int(state.ownership[0].sum()) == 1
    assert int(state.ownership[1].sum()) == 1
    assert int((state.armies * state.ownership[1]).sum()) == 20
    assert bool(state.ownership[0, 2, 2])


def test_competition_step_applies_build_action():
    grid = jnp.zeros((5, 5), dtype=jnp.int32).at[2, 2].set(1).at[4, 4].set(2)
    state = game.create_initial_state(grid)
    state = state._replace(
        armies=state.armies.at[2, 3].set(60).at[4, 4].set(10),
        ownership=state.ownership.at[0, 2, 3].set(True),
        ownership_neutral=state.ownership_neutral.at[2, 3].set(False),
    )
    actions = jnp.array([[2, 2, 3, 0, 0], [1, 0, 0, 0, 0]], dtype=jnp.int32)

    new_state, _ = competition_step(state, actions)

    assert bool(new_state.castles[2, 3])
    assert int(new_state.armies[2, 3]) == 13