"""Pure-JAX static opponent policies for PPO diagnostics."""

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass

import jax
import jax.numpy as jnp

from lost_cities_jax.engine import board_score, current_hand_sorted, legal_action_mask
from lost_cities_jax.types import (
    CARDS_PER_COLOR,
    DECK_DRAWS,
    DISCARD,
    DRAW_DECK,
    LOC_DISCARD,
    LOC_P0_BOARD,
    LOC_P0_HAND,
    N_ACTIONS,
    N_CARDS,
    N_COLORS,
    PLAY,
    State,
)

ActionPolicy = Callable[[State, jax.Array, jax.Array], jax.Array]

ACTION_IDS = jnp.arange(N_ACTIONS, dtype=jnp.int32)
ACTION_HAND_SLOT = ACTION_IDS // 12
ACTION_PLACE_TYPE = (ACTION_IDS % 12) // 6
ACTION_DRAW_SOURCE = ACTION_IDS % 6


@dataclass(frozen=True)
class HeuristicExpertConfig:
    max_open_colors: int | None = None
    target_open_colors: int = 3
    fourth_color_margin: float = 38.0
    extra_color_margin: float = 58.0
    base_open_threshold: float = 12.0
    late_open_threshold: float = 26.0
    future_inflow_scale: float = 0.16
    handshake_risk_weight: float = 0.85
    open_value_scale: float = 1.15
    followup_base: float = 44.0
    followup_rank_weight: float = 2.4
    followup_gap_penalty: float = 8.0
    hs_followup_base: float = 22.0
    new_open_base: float = 18.0
    play_bias: float = 8.0
    defensive_sacrifice_bonus: float = 0.0
    discard_base: float = 20.0
    discard_rank_penalty: float = 1.05
    discard_own_open_penalty: float = 16.0
    discard_own_potential_penalty: float = 0.42
    denial_immediate_penalty: float = 24.0
    denial_near_penalty: float = 8.0
    denial_hs_multiplier: float = 0.8
    deck_draw_base: float = 4.5
    clock_draw_weight: float = 3.0
    useful_pile_draw_base: float = 13.0
    useful_pile_draw_rank_weight: float = 1.0
    useless_pile_draw_penalty: float = 36.0


def policy_by_name(name: str) -> ActionPolicy:
    normalized = name.replace("-", "_").lower()
    if normalized == "discard_only":
        return discard_only_action
    if normalized == "heuristic_balanced":
        return heuristic_balanced_action
    if normalized == "heuristic_cautious":
        return heuristic_cautious_action
    if normalized == "heuristic_expert":
        return heuristic_expert_action
    if normalized in {"expert_cap2", "heuristic_expert_cap2"}:
        return make_heuristic_expert_policy(HeuristicExpertConfig(max_open_colors=2))
    if normalized in {"expert_cap3", "heuristic_expert_cap3"}:
        return make_heuristic_expert_policy(HeuristicExpertConfig(max_open_colors=3))
    if normalized in {"expert_capn", "expert_cap_n", "heuristic_expert_capn"}:
        return make_heuristic_expert_policy(HeuristicExpertConfig(max_open_colors=None))
    msg = f"unknown opponent policy: {name}"
    raise ValueError(msg)


def discard_only_action(state: State, player: jax.Array, rng: jax.Array) -> jax.Array:
    del player, rng
    mask = legal_action_mask(state)
    scores = (
        (ACTION_PLACE_TYPE == DISCARD).astype(jnp.float32) * 100.0
        + (ACTION_DRAW_SOURCE == DRAW_DECK).astype(jnp.float32) * 10.0
        - ACTION_IDS.astype(jnp.float32) * 1.0e-4
    )
    return _argmax_legal(mask, scores)


def heuristic_balanced_action(state: State, player: jax.Array, rng: jax.Array) -> jax.Array:
    del rng
    return _heuristic_action(state, player, cautious=False)


def heuristic_cautious_action(state: State, player: jax.Array, rng: jax.Array) -> jax.Array:
    del rng
    return _heuristic_action(state, player, cautious=True)


def heuristic_expert_action(state: State, player: jax.Array, rng: jax.Array) -> jax.Array:
    del rng
    return _heuristic_expert_action(state, player, HeuristicExpertConfig())


def make_heuristic_expert_policy(
    config: HeuristicExpertConfig | None = None,
) -> ActionPolicy:
    cfg = HeuristicExpertConfig() if config is None else config

    def policy(state: State, player: jax.Array, rng: jax.Array) -> jax.Array:
        del rng
        return _heuristic_expert_action(state, player, cfg)

    return policy


def random_legal_action(state: State, player: jax.Array, rng: jax.Array) -> jax.Array:
    del player
    mask = legal_action_mask(state)
    logits = jnp.where(mask, 0.0, -1.0e9)
    return jax.random.categorical(rng, logits).astype(jnp.int32)


def _heuristic_action(state: State, player: jax.Array, *, cautious: bool) -> jax.Array:
    player = player.astype(jnp.int32)
    mask = legal_action_mask(state)
    cards = current_hand_sorted(state, player).astype(jnp.int32)
    action_cards = cards[ACTION_HAND_SLOT]
    valid_card = action_cards < N_CARDS
    color = jnp.clip(action_cards // CARDS_PER_COLOR, 0, N_COLORS - 1)
    slot = action_cards % CARDS_PER_COLOR
    is_handshake = slot < 3
    rank = jnp.where(is_handshake, 0, slot - 1).astype(jnp.float32)

    col_top = state.col_top[player, color].astype(jnp.float32)
    col_len = state.col_len[player, color].astype(jnp.float32)
    col_hs = state.col_hs[player, color].astype(jnp.float32)
    is_open = col_len > 0

    if cautious:
        play_number_score = jnp.where(
            is_open,
            16.0 + rank * 2.5 - (rank - col_top - 1.0) * 0.25,
            rank * 2.0 - 8.0,
        )
        play_hs_score = 4.0 - col_hs * 0.75
        discard_score = 13.0 - rank * 1.2 - is_open.astype(jnp.float32) * 5.0
    else:
        play_number_score = jnp.where(
            is_open,
            14.0 + rank * 2.3 - (rank - col_top - 1.0) * 0.2,
            rank * 2.1 - 5.0,
        )
        play_hs_score = 8.0 - col_hs * 0.5
        discard_score = 12.0 - rank - is_open.astype(jnp.float32) * 4.0

    place_score = jnp.where(
        ACTION_PLACE_TYPE == PLAY,
        jnp.where(is_handshake, play_hs_score, play_number_score),
        discard_score,
    )
    draw_score = _draw_source_score(state, player, ACTION_DRAW_SOURCE)
    scores = jnp.where(valid_card, place_score + draw_score, -1.0e9)
    scores = scores - ACTION_IDS.astype(jnp.float32) * 1.0e-4
    return _argmax_legal(mask, scores)


def _heuristic_expert_action(
    state: State, player: jax.Array, cfg: HeuristicExpertConfig
) -> jax.Array:
    player = player.astype(jnp.int32)
    opponent = 1 - player
    mask = legal_action_mask(state)
    hand = current_hand_sorted(state, player).astype(jnp.int32)
    action_cards = hand[ACTION_HAND_SLOT]
    valid_card = action_cards < N_CARDS
    safe_cards = jnp.clip(action_cards, 0, N_CARDS - 1)
    color = safe_cards // CARDS_PER_COLOR
    slot = safe_cards % CARDS_PER_COLOR
    is_handshake = slot < 3
    rank_i = jnp.where(is_handshake, 0, slot - 1)
    rank = rank_i.astype(jnp.float32)

    hand_numeric_sum, hand_hs, hand_high_count = _hand_color_features(hand)
    unseen_numeric = _unseen_numeric_sum(state, player, hand)
    deck_frac = (N_CARDS - state.draw_ptr).astype(jnp.float32) / DECK_DRAWS
    future_inflow = unseen_numeric * deck_frac * cfg.future_inflow_scale
    multiplier = 1.0 + hand_hs
    hand_total = hand_numeric_sum + future_inflow
    downside = jnp.maximum(20.0 - hand_numeric_sum, 0.0)
    open_value = (
        (hand_total - 20.0) * multiplier
        - downside * (1.0 + hand_hs * cfg.handshake_risk_weight)
        + hand_high_count * 2.5
    )

    progress = (state.draw_ptr - 16).astype(jnp.float32) / DECK_DRAWS
    active_colors = jnp.sum(state.col_len[player] > 0).astype(jnp.float32)
    hard_cap_allows_new_color = (
        jnp.asarray(True)
        if cfg.max_open_colors is None
        else active_colors < jnp.asarray(cfg.max_open_colors, dtype=jnp.float32)
    )
    open_threshold = (
        cfg.base_open_threshold
        + progress * cfg.late_open_threshold
        + jnp.maximum(active_colors - 2.0, 0.0) * 9.0
    )
    color_open_value = open_value[color]
    concentration_margin = jnp.where(
        active_colors < cfg.target_open_colors,
        0.0,
        jnp.where(active_colors < 4.0, cfg.fourth_color_margin, cfg.extra_color_margin),
    )
    ev_open_allowed = color_open_value > (open_threshold + concentration_margin)
    tempo_open_allowed = (
        (active_colors < cfg.target_open_colors)
        & ~is_handshake
        & (rank >= 3.0)
        & (rank <= 6.0)
        & (color_open_value > -4.0)
    )
    fourth_tempo_allowed = (
        (active_colors < 4.0) & ~is_handshake & (rank >= 8.0) & (color_open_value > 8.0)
    )
    open_allowed = ev_open_allowed | tempo_open_allowed | fourth_tempo_allowed
    hs_open_allowed = (
        open_allowed
        & (hand_numeric_sum[color] >= 18.0)
        & (hand_numeric_sum[color] + future_inflow[color] > 28.0)
        & (progress < 0.55)
    )
    hs_followup_allowed = (
        (hand_numeric_sum[color] >= 18.0)
        & (hand_numeric_sum[color] + future_inflow[color] > 28.0)
        & (progress < 0.55)
    )

    top = state.col_top[player, color].astype(jnp.float32)
    hs_count = state.col_hs[player, color].astype(jnp.float32)
    gap = jnp.maximum(rank - top - 1.0, 0.0)
    followup_number_score = (
        cfg.followup_base
        + rank * cfg.followup_rank_weight
        - gap * cfg.followup_gap_penalty
        + jnp.maximum(color_open_value, -12.0) * 0.12
    )
    followup_hs_score = (
        cfg.hs_followup_base
        + jnp.maximum(color_open_value, -10.0) * 0.10
        - hs_count * 4.0
        - progress * 10.0
    )
    new_open_number_score = (
        cfg.new_open_base + color_open_value * cfg.open_value_scale - rank * 0.35
    )
    new_open_hs_score = (
        cfg.new_open_base - 14.0 + color_open_value * 0.9 + hand_hs[color] * 1.0 - progress * 8.0
    )

    own_open = state.col_len[player, color] > 0
    discard_danger = _discard_danger(state, opponent, safe_cards, cfg)
    defensive_sacrifice = (
        (discard_danger > 22.0)
        & ~own_open
        & hard_cap_allows_new_color
        & (active_colors < cfg.target_open_colors)
        & (rank_i <= 5)
        & (color_open_value > -18.0)
    )
    new_color_allowed = (
        jnp.where(is_handshake, hs_open_allowed, open_allowed) & hard_cap_allows_new_color
    )
    play_score = jnp.where(
        own_open,
        jnp.where(
            is_handshake,
            jnp.where(hs_followup_allowed, followup_hs_score, -60.0),
            followup_number_score,
        ),
        jnp.where(
            new_color_allowed | defensive_sacrifice,
            jnp.where(is_handshake, new_open_hs_score, new_open_number_score),
            -60.0,
        ),
    )
    play_score = jnp.where(
        defensive_sacrifice, play_score + cfg.defensive_sacrifice_bonus, play_score
    )

    discard_score = (
        cfg.discard_base
        - rank * cfg.discard_rank_penalty
        - own_open.astype(jnp.float32) * cfg.discard_own_open_penalty
        - jnp.maximum(color_open_value, 0.0) * cfg.discard_own_potential_penalty
        - discard_danger
    )
    discard_score = jnp.where(is_handshake, discard_score - 5.0, discard_score)

    place_score = jnp.where(ACTION_PLACE_TYPE == PLAY, play_score + cfg.play_bias, discard_score)
    draw_score = _expert_draw_score(state, player, ACTION_DRAW_SOURCE, open_value, cfg)
    scores = jnp.where(valid_card, place_score + draw_score, -1.0e9)
    scores = scores - ACTION_IDS.astype(jnp.float32) * 1.0e-4
    return _argmax_legal(mask, scores)


def _hand_color_features(hand: jax.Array) -> tuple[jax.Array, jax.Array, jax.Array]:
    valid = hand < N_CARDS
    safe = jnp.clip(hand, 0, N_CARDS - 1)
    colors = safe // CARDS_PER_COLOR
    slots = safe % CARDS_PER_COLOR
    is_hs = (slots < 3) & valid
    ranks = jnp.where((slots >= 3) & valid, slots - 1, 0).astype(jnp.float32)
    numeric_sum = []
    hs_count = []
    high_count = []
    for color in range(N_COLORS):
        in_color = colors == color
        numeric_sum.append(jnp.sum(jnp.where(in_color, ranks, 0.0)))
        hs_count.append(jnp.sum(jnp.where(in_color & is_hs, 1.0, 0.0)))
        high_count.append(jnp.sum(jnp.where(in_color & (ranks >= 7.0), 1.0, 0.0)))
    return jnp.stack(numeric_sum), jnp.stack(hs_count), jnp.stack(high_count)


def _unseen_numeric_sum(state: State, player: jax.Array, hand: jax.Array) -> jax.Array:
    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)
    visible = (
        (loc == LOC_DISCARD)
        | (loc == LOC_P0_BOARD)
        | (loc == (LOC_P0_BOARD + 1))
        | (loc == (LOC_P0_HAND + player.astype(jnp.int32)))
    )
    visible = visible | state.hand_public
    unseen_sum = []
    for color in range(N_COLORS):
        unseen_sum.append(jnp.sum(jnp.where((colors == color) & ~visible, ranks, 0.0)))
    return jnp.stack(unseen_sum)


def _discard_danger(
    state: State, opponent: jax.Array, card: jax.Array, cfg: HeuristicExpertConfig
) -> jax.Array:
    color = card // CARDS_PER_COLOR
    slot = card % CARDS_PER_COLOR
    is_hs = slot < 3
    rank = jnp.where(is_hs, 0, slot - 1)
    opp_top = state.col_top[opponent, color].astype(jnp.int32)
    opp_len = state.col_len[opponent, color].astype(jnp.int32)
    opp_hs = state.col_hs[opponent, color].astype(jnp.float32)
    immediate = jnp.where(is_hs, opp_top == 0, (opp_len > 0) & (rank > opp_top))
    near = jnp.where(is_hs, opp_top == 0, rank > opp_top)
    multiplier = 1.0 + opp_hs * cfg.denial_hs_multiplier
    return (
        immediate.astype(jnp.float32) * cfg.denial_immediate_penalty * multiplier
        + (~immediate & near).astype(jnp.float32) * cfg.denial_near_penalty * multiplier
    )


def _expert_draw_score(
    state: State,
    player: jax.Array,
    draw_source: jax.Array,
    open_value: jax.Array,
    cfg: HeuristicExpertConfig,
) -> jax.Array:
    scores = board_score(state)
    opponent = 1 - player.astype(jnp.int32)
    score_diff = scores[player.astype(jnp.int32)] - scores[opponent]
    ahead = jnp.tanh(score_diff / 40.0)
    deck_score = jnp.where(
        draw_source == DRAW_DECK,
        cfg.deck_draw_base + ahead * cfg.clock_draw_weight,
        0.0,
    )

    pile_color = jnp.clip(draw_source - 1, 0, N_COLORS - 1)
    pile_len = state.pile_len[pile_color].astype(jnp.int32)
    top_idx = jnp.maximum(pile_len - 1, 0)
    top_card = jnp.clip(state.pile[pile_color, top_idx].astype(jnp.int32), 0, N_CARDS - 1)
    slot = top_card % CARDS_PER_COLOR
    is_hs = slot < 3
    rank = jnp.where(is_hs, 0, slot - 1)
    color = top_card // CARDS_PER_COLOR
    own_top = state.col_top[player.astype(jnp.int32), color].astype(jnp.int32)
    own_open = state.col_len[player.astype(jnp.int32), color] > 0
    can_play_open = jnp.where(is_hs, own_top == 0, rank > own_top)
    worth_opening = open_value[color] > 18.0
    high_rank_plan = (~is_hs) & (rank >= 7) & (open_value[color] > -4.0)
    useful = (pile_len > 0) & can_play_open & (own_open | worth_opening | high_rank_plan)
    pile_score = jnp.where(
        draw_source > 0,
        jnp.where(
            useful,
            cfg.useful_pile_draw_base
            + rank.astype(jnp.float32) * cfg.useful_pile_draw_rank_weight
            - ahead * cfg.clock_draw_weight,
            -cfg.useless_pile_draw_penalty,
        ),
        0.0,
    )
    return deck_score + pile_score


def _draw_source_score(state: State, player: jax.Array, draw_source: jax.Array) -> jax.Array:
    deck_score = jnp.where(draw_source == DRAW_DECK, 1.0, 0.0)
    pile_color = jnp.clip(draw_source - 1, 0, N_COLORS - 1)
    pile_len = state.pile_len[pile_color].astype(jnp.int32)
    top_idx = jnp.maximum(pile_len - 1, 0)
    top_card = jnp.clip(state.pile[pile_color, top_idx].astype(jnp.int32), 0, N_CARDS - 1)
    slot = top_card % CARDS_PER_COLOR
    is_handshake = slot < 3
    rank = jnp.where(is_handshake, 0, slot - 1)
    color = top_card // CARDS_PER_COLOR
    col_top = state.col_top[player.astype(jnp.int32), color]
    can_play = jnp.where(is_handshake, col_top == 0, rank > col_top)
    pile_score = jnp.where(
        draw_source > 0,
        jnp.where(can_play & (pile_len > 0), 9.0 + rank.astype(jnp.float32), 0.5),
        0.0,
    )
    return deck_score + pile_score


def _argmax_legal(mask: jax.Array, scores: jax.Array) -> jax.Array:
    legal_scores = jnp.where(mask, scores, -1.0e9)
    return jnp.argmax(legal_scores).astype(jnp.int32)


__all__ = [
    "ActionPolicy",
    "HeuristicExpertConfig",
    "discard_only_action",
    "heuristic_balanced_action",
    "heuristic_cautious_action",
    "heuristic_expert_action",
    "make_heuristic_expert_policy",
    "policy_by_name",
    "random_legal_action",
]
