"""Player-view observation for a three-round match.

Everything the single-round observation had, plus the four things a match policy
cannot play without:

- **carry**, as a scalar *and* a one-hot over bins. Round three is a threshold
  problem -- "win by 41 or lose" plays nothing like "win by 39" -- and a lone
  scalar makes the network learn a cliff from a smooth input. (The old
  ``score_diff`` divided by ``MAX_ABS_SCORE`` = 780, which squashed a decisive
  50-point lead to 0.06. That scaling is fixed here too.)
- **which round it is**, and how many are left.
- **whose turn it is.** The single-round obs never had this, yet the critic is
  trained on opponent-turn states as well, leaving it to recover turn parity
  from ``step_count / 400``.
- **the deck clock.** A round ends when the last deck card is drawn, so turn
  parity decides who gets the final placement -- and players bend that parity by
  drawing from discard piles instead.

Live points per colour are split by where the card is (hand / discard pile /
unseen) because a discard pile is public and recoverable: lumping it in with the
unseen cards hides the high cards sitting in plain view.
"""

from __future__ import annotations

import jax
import jax.numpy as jnp

from lost_cities_jax.engine import board_score
from lost_cities_jax.match import N_ROUNDS, MatchState
from lost_cities_jax.obs import observation
from lost_cities_jax.types import (
    CARDS_PER_COLOR,
    LOC_DECK,
    LOC_DISCARD,
    LOC_P0_HAND,
    N_CARDS,
    N_COLORS,
    OBS_DIM,
)

# Bin edges for the carry one-hot. Packed tightly around zero: that is where the
# round-three decision actually flips.
CARRY_BIN_EDGES = jnp.asarray([-60.0, -30.0, -12.0, -1.0, 1.0, 12.0, 30.0, 60.0])
N_CARRY_BINS = CARRY_BIN_EDGES.shape[0] + 1

CARRY_SCALE = 75.0
"""Typical round margin, not the theoretical 780 maximum."""

MAX_COLOR_POINTS = 54.0
"""2+3+...+10, the most a single expedition can be worth before multipliers."""

N_MATCH_SCALARS = (
    1  # carry, scaled
    + N_CARRY_BINS  # carry, binned
    + N_ROUNDS  # which round
    + 1  # rounds left
    + 1  # is it my turn
    + 1  # did I lead this round
    + 1  # do I take the last deck card at the current parity
    + 2 * N_COLORS * 3  # live points: both players x colour x {hand, discard, unseen}
)
MATCH_OBS_DIM = OBS_DIM + N_MATCH_SCALARS


def _live_points(state, viewer: jax.Array, subject: jax.Array) -> jax.Array:
    """Points still reachable for ``subject``, split by where the card sits.

    Returns ``float32[N_COLORS, 3]`` over {in hand, in a discard pile, unseen}.
    Seen through ``viewer``'s eyes: a card in the opponent's hand only counts as
    "in hand" if it is public, otherwise it is unseen.
    """

    ids = jnp.arange(N_CARDS, dtype=jnp.int32)
    colors = ids // CARDS_PER_COLOR
    slots = ids % CARDS_PER_COLOR
    ranks = jnp.where(slots >= 3, slots - 1, 0).astype(jnp.float32)

    loc = state.card_loc.astype(jnp.int32)
    subject_hand = loc == (LOC_P0_HAND + subject)
    is_mine = subject == viewer

    # An ascending column can only take cards above its current top.
    above_top = ranks > state.col_top[subject][colors].astype(jnp.float32)

    in_hand = subject_hand & (is_mine | state.hand_public)
    in_discard = loc == LOC_DISCARD
    unseen = (loc == LOC_DECK) | (subject_hand & ~is_mine & ~state.hand_public)
    # Cards in the *other* player's hidden hand are unseen to the viewer too, but
    # they are not reachable by the subject, so they are deliberately excluded.

    def by_color(mask: jax.Array) -> jax.Array:
        weighted = jnp.where(mask & above_top, ranks, 0.0)
        return jax.ops.segment_sum(weighted, colors, num_segments=N_COLORS)

    stacked = jnp.stack([by_color(in_hand), by_color(in_discard), by_color(unseen)], axis=1)
    return stacked / MAX_COLOR_POINTS


def match_observation(state: MatchState, player: jax.Array) -> jax.Array:
    """Return ``float32[MATCH_OBS_DIM]`` from ``player``'s seat."""

    player = jnp.asarray(player, dtype=jnp.int32)
    opponent = 1 - player
    round_state = state.round

    base = observation(round_state, player)

    totals = state.carry + board_score(round_state).astype(jnp.int32)
    lead = (totals[player] - totals[opponent]).astype(jnp.float32)
    carry_scaled = jnp.clip(lead / CARRY_SCALE, -2.0, 2.0).reshape((1,))
    carry_bins = jax.nn.one_hot(jnp.digitize(lead, CARRY_BIN_EDGES), N_CARRY_BINS)

    round_onehot = jax.nn.one_hot(state.round_idx, N_ROUNDS)
    rounds_left = ((N_ROUNDS - 1 - state.round_idx).astype(jnp.float32) / (N_ROUNDS - 1)).reshape(
        (1,)
    )

    to_move = round_state.to_move.astype(jnp.int32)
    my_turn = (to_move == player).astype(jnp.float32).reshape((1,))
    # to_move flips every ply, so the round's opener is recoverable from parity.
    parity = round_state.step_count.astype(jnp.int32) & 1
    round_opener = to_move ^ parity
    i_opened = (round_opener == player).astype(jnp.float32).reshape((1,))

    # If both players drew from the deck from here, the last deck card falls to
    # whoever is on move after `remaining - 1` more plies.
    remaining = jnp.maximum(N_CARDS - round_state.draw_ptr, 1).astype(jnp.int32)
    last_drawer = to_move ^ ((remaining - 1) & 1)
    i_take_last = (last_drawer == player).astype(jnp.float32).reshape((1,))

    live_me = _live_points(round_state, player, player).reshape((-1,))
    live_opp = _live_points(round_state, player, opponent).reshape((-1,))

    extra = jnp.concatenate(
        [
            carry_scaled,
            carry_bins,
            round_onehot,
            rounds_left,
            my_turn,
            i_opened,
            i_take_last,
            live_me,
            live_opp,
        ],
        axis=0,
    ).astype(jnp.float32)

    return jnp.concatenate([base, extra], axis=0).reshape((MATCH_OBS_DIM,))


CRITIC_EXTRA_CHANNELS = 3
MATCH_CRITIC_OBS_DIM = MATCH_OBS_DIM + N_CARDS * CRITIC_EXTRA_CHANNELS


def match_critic_observation(state: MatchState, player: jax.Array) -> jax.Array:
    """The same view, plus the hidden state. Training only -- never played from.

    A value baseline may condition on anything the action does not depend on
    without biasing the policy gradient, and the deal is exactly the noise that
    swamps a match-terminal reward. So the critic gets the opponent's hand and
    the deck *in order* -- it knows what is coming -- while the actor keeps the
    masked view. The two run on separate trunks so none of this can leak into
    the logits.
    """

    player = jnp.asarray(player, dtype=jnp.int32)
    opponent = 1 - player
    round_state = state.round
    loc = round_state.card_loc.astype(jnp.int32)

    opp_hand = (loc == (LOC_P0_HAND + opponent)).astype(jnp.float32)
    in_deck = (loc == LOC_DECK).astype(jnp.float32)

    # Position of each card in the undrawn deck, so the critic can see the order
    # rather than just the multiset.
    deck_order = round_state.deck_order.astype(jnp.int32)
    position = (
        jnp.zeros((N_CARDS,), dtype=jnp.int32)
        .at[deck_order]
        .set(jnp.arange(N_CARDS, dtype=jnp.int32))
    )
    depth = (position - round_state.draw_ptr).astype(jnp.float32) / N_CARDS
    deck_depth = jnp.where(in_deck > 0, depth, 0.0)

    hidden = jnp.stack([opp_hand, in_deck, deck_depth], axis=1).reshape((-1,))
    base = match_observation(state, player)
    return jnp.concatenate([base, hidden], axis=0).reshape((MATCH_CRITIC_OBS_DIM,))


__all__ = [
    "MATCH_CRITIC_OBS_DIM",
    "MATCH_OBS_DIM",
    "N_CARRY_BINS",
    "match_critic_observation",
    "match_observation",
]
