from __future__ import annotations

import random

import jax
import jax.numpy as jnp
import numpy as np

from lost_cities_jax import reset_from_order
from lost_cities_jax.types import (
    CARDS_PER_COLOR,
    LOC_DISCARD,
    LOC_P0_BOARD,
    LOC_P0_HAND,
    MAX_PILE_SIZE,
    N_CARDS,
    N_COLORS,
    NO_CARD,
    State,
)


def shuffled_order(seed: int) -> list[int]:
    order = list(range(N_CARDS))
    random.Random(seed).shuffle(order)
    return order


def hs(color: int, index: int = 0) -> int:
    return color * CARDS_PER_COLOR + index


def num(color: int, rank: int) -> int:
    return color * CARDS_PER_COLOR + rank + 1


def state_with_columns(
    p0_columns: list[list[int]],
    p1_columns: list[list[int]] | None = None,
) -> State:
    if p1_columns is None:
        p1_columns = [[] for _ in range(N_COLORS)]

    state = reset_from_order(jnp.arange(N_CARDS, dtype=jnp.int8))
    card_loc = jnp.zeros((N_CARDS,), dtype=jnp.int8)
    col_top = jnp.zeros((2, N_COLORS), dtype=jnp.int8)
    col_hs = jnp.zeros((2, N_COLORS), dtype=jnp.int8)
    col_len = jnp.zeros((2, N_COLORS), dtype=jnp.int8)

    for player, columns in enumerate([p0_columns, p1_columns]):
        for color, cards in enumerate(columns):
            if not cards:
                continue
            idx = jnp.asarray(cards, dtype=jnp.int32)
            card_loc = card_loc.at[idx].set(jnp.int8(LOC_P0_BOARD + player))
            slots = [card % CARDS_PER_COLOR for card in cards]
            ranks = [slot - 1 for slot in slots if slot >= 3]
            col_top = col_top.at[player, color].set(max(ranks, default=0))
            col_hs = col_hs.at[player, color].set(sum(slot < 3 for slot in slots))
            col_len = col_len.at[player, color].set(len(cards))

    return state._replace(
        card_loc=card_loc,
        col_top=col_top,
        col_hs=col_hs,
        col_len=col_len,
    )


def first_deck_draw_action(mask) -> int:
    values = np.asarray(mask, dtype=bool)
    for action, legal in enumerate(values):
        if legal and action % 6 == 0:
            return action
    raise AssertionError("no legal deck-draw action")


def assert_state_consistent(state: State) -> None:
    loc = np.asarray(state.card_loc)
    public = np.asarray(state.hand_public)
    pile = np.asarray(state.pile)
    pile_len = np.asarray(state.pile_len)
    col_top = np.asarray(state.col_top)
    col_hs = np.asarray(state.col_hs)
    col_len = np.asarray(state.col_len)

    assert loc.shape == (N_CARDS,)
    assert np.all((0 <= loc) & (loc <= LOC_DISCARD))
    assert sum(np.bincount(loc, minlength=LOC_DISCARD + 1)) == N_CARDS
    assert int(np.asarray(state.just_discarded)) == NO_CARD

    for player in range(2):
        assert int(np.sum(loc == LOC_P0_HAND + player)) == 8
        for color in range(N_COLORS):
            cards = [
                card
                for card in range(N_CARDS)
                if loc[card] == LOC_P0_BOARD + player and card // CARDS_PER_COLOR == color
            ]
            slots = [card % CARDS_PER_COLOR for card in cards]
            ranks = [slot - 1 for slot in slots if slot >= 3]
            assert int(col_len[player, color]) == len(cards)
            assert int(col_hs[player, color]) == sum(slot < 3 for slot in slots)
            assert int(col_top[player, color]) == max(ranks, default=0)
            assert int(col_hs[player, color]) + len(ranks) == int(col_len[player, color])

    for color in range(N_COLORS):
        length = int(pile_len[color])
        assert 0 <= length < MAX_PILE_SIZE
        used = list(map(int, pile[color, :length]))
        assert all(card != NO_CARD for card in used)
        assert len(set(used)) == len(used)
        assert np.all(pile[color, length:] == NO_CARD)
        discard_cards = [
            card
            for card in range(N_CARDS)
            if loc[card] == LOC_DISCARD and card // CARDS_PER_COLOR == color
        ]
        assert sorted(used) == sorted(discard_cards)

    in_hand = (loc == LOC_P0_HAND) | (loc == LOC_P0_HAND + 1)
    assert not np.any(public & ~in_hand)


def assert_states_equal(left: State, right: State) -> None:
    for left_leaf, right_leaf in zip(
        jax.tree_util.tree_leaves(left), jax.tree_util.tree_leaves(right), strict=False
    ):
        np.testing.assert_array_equal(np.asarray(left_leaf), np.asarray(right_leaf))
