"""Shared constants and state types for the JAX Lost Cities engine."""

from __future__ import annotations

from typing import NamedTuple

import jax

N_PLAYERS = 2
N_COLORS = 5
CARDS_PER_COLOR = 12
N_CARDS = N_COLORS * CARDS_PER_COLOR
HAND_SIZE = 8
INITIAL_DEAL = N_PLAYERS * HAND_SIZE
DECK_DRAWS = N_CARDS - INITIAL_DEAL
MAX_PILE_SIZE = CARDS_PER_COLOR + 1

N_PLACE_TYPES = 2
N_DRAW_SOURCES = 6
N_ACTIONS = HAND_SIZE * N_PLACE_TYPES * N_DRAW_SOURCES

MAX_STEPS = 400

CARD_CHANNELS = 7
CARD_OBS_DIM = N_CARDS * CARD_CHANNELS
SCALAR_OBS_DIM = 34
OBS_DIM = CARD_OBS_DIM + SCALAR_OBS_DIM

LOC_DECK = 0
LOC_P0_HAND = 1
LOC_P1_HAND = 2
LOC_P0_BOARD = 3
LOC_P1_BOARD = 4
LOC_DISCARD = 5

PLAY = 0
DISCARD = 1
DRAW_DECK = 0

UNKNOWN_CARD = N_CARDS
NO_CARD = -1

MAX_ABS_SCORE = 780.0


class State(NamedTuple):
    deck_order: jax.Array
    draw_ptr: jax.Array
    card_loc: jax.Array
    hand_public: jax.Array
    col_top: jax.Array
    col_hs: jax.Array
    col_len: jax.Array
    pile: jax.Array
    pile_len: jax.Array
    to_move: jax.Array
    just_discarded: jax.Array
    step_count: jax.Array
    done: jax.Array
