"""Observation construction for the JAX Lost Cities engine."""

from __future__ import annotations

import jax
import jax.numpy as jnp

from lost_cities_jax.engine import board_score
from lost_cities_jax.types import (
    CARDS_PER_COLOR,
    DECK_DRAWS,
    HAND_SIZE,
    LOC_DECK,
    LOC_DISCARD,
    LOC_P0_BOARD,
    LOC_P0_HAND,
    MAX_ABS_SCORE,
    MAX_STEPS,
    N_CARDS,
    N_COLORS,
    NO_CARD,
    OBS_DIM,
    State,
)


def observation(state: State, player: jax.Array) -> jax.Array:
    """Return a masked player-view observation as ``float32[OBS_DIM]``."""

    player = jnp.asarray(player, dtype=jnp.int32)
    opponent = 1 - player
    loc = state.card_loc.astype(jnp.int32)
    ids = jnp.arange(N_CARDS, dtype=jnp.int32)
    colors = ids // CARDS_PER_COLOR

    my_hand_loc = LOC_P0_HAND + player
    opp_hand_loc = LOC_P0_HAND + opponent
    my_board_loc = LOC_P0_BOARD + player
    opp_board_loc = LOC_P0_BOARD + opponent

    pile_len = state.pile_len.astype(jnp.int32)
    top_idx = jnp.maximum(pile_len - 1, 0)
    pile_tops = jnp.where(
        pile_len > 0,
        state.pile[jnp.arange(N_COLORS), top_idx].astype(jnp.int32),
        NO_CARD,
    )
    is_discard = loc == LOC_DISCARD
    is_discard_top = is_discard & (pile_tops[colors] == ids)
    is_discard_not_top = is_discard & ~is_discard_top
    is_opp_hand = loc == opp_hand_loc
    is_opp_public = is_opp_hand & state.hand_public

    channels = jnp.stack(
        [
            loc == my_hand_loc,
            loc == my_board_loc,
            loc == opp_board_loc,
            is_discard_top,
            is_discard_not_top,
            is_opp_public,
            (loc == LOC_DECK) | (is_opp_hand & ~state.hand_public),
        ],
        axis=1,
    ).astype(jnp.float32)

    scores = board_score(state)
    score_diff = ((scores[player] - scores[opponent]) / MAX_ABS_SCORE).reshape((1,))
    scalars = jnp.concatenate(
        [
            ((N_CARDS - state.draw_ptr).astype(jnp.float32) / DECK_DRAWS).reshape((1,)),
            (jnp.sum(is_opp_hand & ~state.hand_public).astype(jnp.float32) / HAND_SIZE).reshape(
                (1,)
            ),
            (state.step_count.astype(jnp.float32) / MAX_STEPS).reshape((1,)),
            state.col_top[player].astype(jnp.float32) / 10.0,
            state.col_top[opponent].astype(jnp.float32) / 10.0,
            state.col_hs[player].astype(jnp.float32) / 3.0,
            state.col_hs[opponent].astype(jnp.float32) / 3.0,
            state.col_len[player].astype(jnp.float32) / CARDS_PER_COLOR,
            state.col_len[opponent].astype(jnp.float32) / CARDS_PER_COLOR,
            score_diff.astype(jnp.float32),
        ],
        axis=0,
    )
    obs = jnp.concatenate([channels.reshape((-1,)), scalars.astype(jnp.float32)], axis=0)
    return obs.reshape((OBS_DIM,))


__all__ = ["observation"]
