"""Match-level evaluation, and the probe that says whether carry changed anything.

Single-deal win rate cannot tell you if the three-round rework worked. Two things
can:

- **Duplicate match play.** The same three deals and the same coin flips, played
  from both seats. Deal luck cancels, so what is left is the policy difference.
- **The carry probe.** Drop the policy into round three holding a fixed deficit or
  lead and watch what it does. A carry-blind policy plays a 60-point deficit
  exactly like a 60-point lead. If the numbers below do not move with carry, the
  rework bought nothing, no matter what the win rate says.
"""

from __future__ import annotations

import math

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

from lost_cities_jax.engine import reset_from_order
from lost_cities_jax.match import (
    N_ROUNDS,
    MatchState,
    match_legal_action_mask,
    match_reset_from,
    match_score,
    match_step,
)
from lost_cities_jax.match_obs import match_critic_observation, match_observation
from lost_cities_jax.match_ppo import MatchActorCritic, match_lead
from lost_cities_jax.ppo import JaxPPOConfig, mask_logits
from lost_cities_jax.types import DRAW_DECK, MAX_STEPS, N_CARDS

MATCH_SCAN_STEPS = N_ROUNDS * MAX_STEPS


def match_bank(seed: int, matches: int) -> tuple[jnp.ndarray, jnp.ndarray]:
    """Deterministic deals and coin flips, so runs are comparable."""
    rng = np.random.default_rng(seed)
    decks = np.stack(
        [np.stack([rng.permutation(N_CARDS) for _ in range(N_ROUNDS)]) for _ in range(matches)]
    )
    coins = rng.integers(0, 2, size=(matches, N_ROUNDS))
    return jnp.asarray(decks, dtype=jnp.int8), jnp.asarray(coins, dtype=jnp.int8)


def _wilson(wins: float, games: float) -> tuple[float, float]:
    if games == 0:
        return 0.0, 0.0
    z = 1.96
    p = wins / games
    denom = 1 + z * z / games
    centre = p + z * z / (2 * games)
    margin = z * math.sqrt(p * (1 - p) / games + z * z / (4 * games * games))
    return (centre - margin) / denom, (centre + margin) / denom


def make_match_runner(cfg: JaxPPOConfig, opponent_policy=None):
    """Play whole matches to the end. ``opponent_policy`` is a single-round policy."""

    model = MatchActorCritic(cfg.network.hidden_size, cfg.network.num_layers)

    @jax.jit
    def run(params, env: MatchState, learner_seat: jax.Array, rng: jax.Array):
        def body(carry, _):
            env, key, deck_draws, plies = carry
            key, opp_key = jax.random.split(key)

            to_move = env.round.to_move.astype(jnp.int32)
            obs = jax.vmap(match_observation)(env, to_move)
            critic = jax.vmap(match_critic_observation)(env, to_move)
            mask = jax.vmap(match_legal_action_mask)(env)
            logits, _ = model.apply(params, obs, critic)
            learner_action = jnp.argmax(mask_logits(logits, mask), axis=-1).astype(jnp.int32)

            if opponent_policy is None:
                action = learner_action
            else:
                opp_keys = jax.random.split(opp_key, env.done.shape[0])
                opp_action = jax.vmap(opponent_policy, in_axes=(0, 0, 0))(
                    env.round, to_move, opp_keys
                )
                action = jnp.where(to_move == learner_seat, learner_action, opp_action)

            live = ~env.done
            is_deck = (action % 6) == DRAW_DECK
            deck_draws = deck_draws + (is_deck & live).astype(jnp.int32)
            plies = plies + live.astype(jnp.int32)

            env, _, _ = jax.vmap(match_step, in_axes=(0, 0))(env, action)
            return (env, key, deck_draws, plies), None

        n = env.done.shape[0]
        zeros = jnp.zeros((n,), dtype=jnp.int32)
        (env, _, deck_draws, plies), _ = jax.lax.scan(
            body, (env, rng, zeros, zeros), xs=None, length=MATCH_SCAN_STEPS
        )
        return env, deck_draws, plies

    return run


def match_evaluate(cfg: JaxPPOConfig, params, *, matches: int = 2000, seed: int = 20260715) -> dict:
    """Duplicate match play: every deal seen from both seats."""

    decks, coins = match_bank(seed, matches)
    runner = make_match_runner(cfg)

    results = []
    for seat in (0, 1):
        env = jax.vmap(match_reset_from)(decks, coins)
        seats = jnp.full((matches,), seat, dtype=jnp.int32)
        final, deck_draws, plies = runner(params, env, seats, jax.random.PRNGKey(0))
        totals = np.asarray(jax.vmap(match_score)(final))
        lead = totals[:, seat] - totals[:, 1 - seat]
        results.append(
            {
                "lead": lead,
                "plies": np.asarray(plies),
                "deck_share": np.asarray(deck_draws) / np.maximum(np.asarray(plies), 1),
                "done": np.asarray(final.done),
            }
        )

    lead = np.concatenate([r["lead"] for r in results])
    games = float(lead.size)
    wins = float((lead > 0).sum())
    ties = float((lead == 0).sum())
    low, high = _wilson(wins, games)
    return {
        "matches": games,
        "wins": wins,
        "ties": ties,
        "losses": games - wins - ties,
        "match_win_rate": wins / games,
        "wilson_low": low,
        "wilson_high": high,
        "mean_total_lead": float(lead.mean()),
        "mean_plies": float(np.concatenate([r["plies"] for r in results]).mean()),
        "unfinished_rate": float(
            1.0 - np.concatenate([r["done"] for r in results]).astype(np.float32).mean()
        ),
    }


def head_to_head(
    cfg: JaxPPOConfig, params_a, params_b, *, matches: int = 4096, seed: int = 20260717
) -> dict:
    """Duplicate matches between two policies: same deals, same coins, both seats.

    Every deal is played twice with the seats swapped, so deal luck cancels and
    what is left is the difference between the two policies.
    """

    decks, coins = match_bank(seed, matches)
    model = MatchActorCritic(cfg.network.hidden_size, cfg.network.num_layers)

    @jax.jit
    def run(env: MatchState, a_seat: jax.Array):
        def body(carry, _):
            env, _unused = carry
            to_move = env.round.to_move.astype(jnp.int32)
            obs = jax.vmap(match_observation)(env, to_move)
            critic = jax.vmap(match_critic_observation)(env, to_move)
            mask = jax.vmap(match_legal_action_mask)(env)

            logits_a, _ = model.apply(params_a, obs, critic)
            logits_b, _ = model.apply(params_b, obs, critic)
            act_a = jnp.argmax(mask_logits(logits_a, mask), axis=-1).astype(jnp.int32)
            act_b = jnp.argmax(mask_logits(logits_b, mask), axis=-1).astype(jnp.int32)

            action = jnp.where(to_move == a_seat, act_a, act_b)
            env, _, _ = jax.vmap(match_step, in_axes=(0, 0))(env, action)
            return (env, _unused), None

        (env, _), _ = jax.lax.scan(body, (env, jnp.int32(0)), xs=None, length=MATCH_SCAN_STEPS)
        return env

    leads = []
    for seat in (0, 1):
        env = jax.vmap(match_reset_from)(decks, coins)
        final = run(env, jnp.full((matches,), seat, dtype=jnp.int32))
        totals = np.asarray(jax.vmap(match_score)(final))
        leads.append(totals[:, seat] - totals[:, 1 - seat])

    lead = np.concatenate(leads)
    games = float(lead.size)
    wins = float((lead > 0).sum())
    ties = float((lead == 0).sum())
    low, high = _wilson(wins, games)
    return {
        "matches": games,
        "a_win_rate": wins / games,
        "wilson_low": low,
        "wilson_high": high,
        "ties": ties,
        "a_mean_lead": float(lead.mean()),
    }


def carry_probe(
    cfg: JaxPPOConfig,
    params,
    *,
    carry_levels=(-60, -25, -1, 1, 25, 60),
    matches: int = 512,
    seed: int = 20260716,
) -> list[dict]:
    """Start the policy in round three holding ``carry`` and see what changes.

    A policy that reads carry should open more expeditions and lean on wagers when
    it is behind (it has to swing), and shut down when it is ahead. Flat rows mean
    the policy is ignoring the one thing the three-round rework added.
    """

    decks, coins = match_bank(seed, matches)
    runner = make_match_runner(cfg)
    rows = []

    for level in carry_levels:
        base = jax.vmap(match_reset_from)(decks, coins)
        carry = jnp.tile(jnp.asarray([level, 0], dtype=jnp.int32), (matches, 1))
        round_idx = jnp.full((matches,), N_ROUNDS - 1, dtype=jnp.int32)
        # Whoever is ahead leads the round, exactly as the match layer would.
        first = jnp.int8(0 if level > 0 else 1)
        rounds = jax.vmap(reset_from_order, in_axes=(0, None))(decks[:, N_ROUNDS - 1], first)

        env = base._replace(round=rounds, carry=carry, round_idx=round_idx)
        final, deck_draws, plies = runner(
            params, env, jnp.zeros((matches,), jnp.int32), jax.random.PRNGKey(1)
        )

        lead = np.asarray(jax.vmap(match_lead)(final))
        opened = np.asarray(jnp.sum(final.round.col_len[:, 0, :] > 0, axis=-1))
        wagers = np.asarray(jnp.sum(final.round.col_hs[:, 0, :], axis=-1))
        deck_share = np.asarray(deck_draws) / np.maximum(np.asarray(plies), 1)

        rows.append(
            {
                "carry": level,
                "win_rate": float((lead > 0).mean()),
                "opened_colors": float(opened.mean()),
                "wagers_played": float(wagers.mean()),
                "deck_race_rate": float(deck_share.mean()),
                "mean_plies": float(np.asarray(plies).mean()),
            }
        )
    return rows
