from __future__ import annotations

import jax.numpy as jnp
import numpy as np

from lost_cities_jax import legal_action_mask, reset_from_order, score, step
from lost_cities_jax.engine import current_hand_sorted
from lost_cities_jax.types import CARDS_PER_COLOR, N_CARDS
from tests.lost_cities_jax.helpers import assert_states_equal, first_deck_draw_action


def test_discarded_card_cannot_be_drawn_from_same_pile():
    state = reset_from_order(jnp.arange(N_CARDS, dtype=jnp.int8))
    hand = np.asarray(current_hand_sorted(state))
    first_card_color = int(hand[0]) // CARDS_PER_COLOR
    discard_then_same_pile = 0 * 12 + 1 * 6 + (first_card_color + 1)
    discard_then_deck = 0 * 12 + 1 * 6

    mask = np.asarray(legal_action_mask(state), dtype=bool)

    assert not mask[discard_then_same_pile]
    assert mask[discard_then_deck]


def test_illegal_action_is_noop_with_zero_reward():
    state = reset_from_order(jnp.arange(N_CARDS, dtype=jnp.int8))
    hand = np.asarray(current_hand_sorted(state))
    color = int(hand[0]) // CARDS_PER_COLOR
    illegal_action = 0 * 12 + 1 * 6 + (color + 1)

    next_state, reward, done = step(state, jnp.int32(illegal_action))

    assert_states_equal(next_state, state)
    np.testing.assert_array_equal(np.asarray(reward), np.zeros(2, dtype=np.float32))
    assert bool(done) is False


def test_done_state_step_is_noop_and_has_no_repeated_terminal_reward():
    state = reset_from_order(jnp.arange(N_CARDS, dtype=jnp.int8))
    while not bool(state.done):
        action = first_deck_draw_action(legal_action_mask(state))
        state, _, _ = step(state, jnp.int32(action))

    after_done, reward, done = step(state, jnp.int32(0))

    assert_states_equal(after_done, state)
    np.testing.assert_array_equal(np.asarray(reward), np.zeros(2, dtype=np.float32))
    assert bool(done) is True


def test_deck_only_draws_finish_in_exactly_44_plies():
    state = reset_from_order(jnp.arange(N_CARDS, dtype=jnp.int8))
    plies = 0
    while not bool(state.done):
        action = first_deck_draw_action(legal_action_mask(state))
        state, reward, done = step(state, jnp.int32(action))
        plies += 1

    assert plies == 44
    assert int(state.draw_ptr) == N_CARDS
    assert bool(done) is True
    np.testing.assert_array_equal(np.asarray(reward), np.asarray(score(state)))
