"""Phase 0a: GAE bootstrap, in-scan auto-reset, and episode-boundary metrics."""

import jax
import jax.numpy as jnp
import pytest

from lost_cities_jax.engine import reset
from lost_cities_jax.opponents import policy_by_name
from lost_cities_jax.ppo import (
    EpisodeEnd,
    JaxPPOConfig,
    Transition,
    compute_gae,
    create_train_state,
    episode_metrics,
    make_rollout_fn,
    shaping_coefficient,
)


def _cfg(**overrides) -> JaxPPOConfig:
    cfg = JaxPPOConfig()
    cfg.ppo.batch_games = 8
    cfg.ppo.rollout_steps = 120
    cfg.network.hidden_size = 16
    cfg.network.num_layers = 1
    for key, value in overrides.items():
        setattr(cfg.run, key, value)
    return cfg


# --- GAE bootstrap ---------------------------------------------------------


def test_gae_bootstraps_truncated_episode_from_last_value():
    """A rollout that ends mid-episode must carry V(s_T), not zero."""
    rewards = jnp.zeros((3, 1))
    values = jnp.zeros((3, 1))
    dones = jnp.zeros((3, 1), dtype=bool)  # never terminates inside the window

    _, returns_zero = compute_gae(rewards, values, dones, gamma=1.0, gae_lambda=1.0)
    _, returns_boot = compute_gae(
        rewards, values, dones, gamma=1.0, gae_lambda=1.0, last_value=jnp.array([5.0])
    )

    # Without a bootstrap the truncated episode looks worthless.
    assert jnp.allclose(returns_zero, 0.0)
    # With it, every step inherits the tail value.
    assert jnp.allclose(returns_boot, 5.0)


def test_gae_ignores_last_value_when_episode_terminates():
    """A terminal step must not bootstrap past the episode boundary."""
    rewards = jnp.array([[0.0], [1.0]])
    values = jnp.zeros((2, 1))
    dones = jnp.array([[False], [True]])

    _, returns = compute_gae(
        rewards, values, dones, gamma=1.0, gae_lambda=1.0, last_value=jnp.array([99.0])
    )
    # Terminal step sees only its own reward; the step before it sees that too.
    assert jnp.allclose(returns, jnp.array([[1.0], [1.0]]))


def test_gae_done_flag_cuts_credit_between_episodes():
    """With in-scan resets two episodes share a slot; credit must not leak."""
    rewards = jnp.array([[1.0], [7.0]])
    values = jnp.zeros((2, 1))
    dones = jnp.array([[True], [False]])  # first step ends episode A

    _, returns = compute_gae(rewards, values, dones, gamma=1.0, gae_lambda=1.0)
    # Episode A keeps its own reward; episode B's +7 must not flow backwards.
    assert jnp.allclose(returns[0], 1.0)
    assert jnp.allclose(returns[1], 7.0)


# --- episode-boundary metrics ---------------------------------------------


def _episodes(done, episode_return, length) -> EpisodeEnd:
    shape = jnp.asarray(done).shape
    return EpisodeEnd(
        done=jnp.asarray(done),
        episode_return=jnp.asarray(episode_return, dtype=jnp.float32),
        length=jnp.asarray(length, dtype=jnp.int32),
        opened_colors=jnp.zeros(shape, dtype=jnp.int32),
        positive_expeditions=jnp.zeros(shape, dtype=jnp.int32),
        hit_max_steps=jnp.zeros(shape, dtype=bool),
    )


def _transitions(steps: int, games: int) -> Transition:
    zeros = jnp.zeros((steps, games), dtype=jnp.float32)
    false = jnp.zeros((steps, games), dtype=bool)
    return Transition(
        obs=jnp.zeros((steps, games, 1)),
        legal_mask=jnp.zeros((steps, games, 1), dtype=bool),
        action=jnp.zeros((steps, games), dtype=jnp.int32),
        log_prob=zeros,
        value=zeros,
        reward=zeros,
        done=false,
        active=~false,
        actor_mask=~false,
        entropy=zeros,
        play_action=false,
    )


def test_episode_metrics_only_counts_finished_episodes():
    """Unfinished episodes carry a partial return; they must not be averaged in."""
    # Slot 0 finishes twice (returns 10 and 20); slot 1 never finishes.
    done = jnp.array([[True, False], [False, False], [True, False]])
    episode_return = jnp.array([[10.0, 3.0], [0.0, 6.0], [20.0, 9.0]])
    length = jnp.array([[50, 1], [0, 2], [60, 3]])

    metrics = episode_metrics(_transitions(3, 2), _episodes(done, episode_return, length))

    assert float(metrics["episodes_completed"]) == 2.0
    assert float(metrics["return_mean"]) == pytest.approx(15.0)
    assert float(metrics["game_length_mean"]) == pytest.approx(55.0)
    assert int(metrics["game_length_max"]) == 60


def test_episode_metrics_survives_a_rollout_with_no_completions():
    done = jnp.zeros((2, 2), dtype=bool)
    metrics = episode_metrics(
        _transitions(2, 2), _episodes(done, jnp.zeros((2, 2)), jnp.zeros((2, 2)))
    )
    assert float(metrics["episodes_completed"]) == 0.0
    assert float(metrics["return_mean"]) == 0.0  # guarded denominator, not a NaN


# --- in-scan auto-reset ----------------------------------------------------


def test_rollout_never_returns_a_done_env_and_completes_many_episodes():
    """The whole point of in-scan reset: no dead steps, several games per slot."""
    cfg = _cfg()
    rng = jax.random.PRNGKey(0)
    rng, init_key, reset_key, roll_key = jax.random.split(rng, 4)
    state = create_train_state(cfg, init_key)
    env_state = jax.jit(jax.vmap(reset))(jax.random.split(reset_key, cfg.ppo.batch_games))

    rollout_fn = make_rollout_fn(cfg, policy_by_name("discard_only"))
    next_env, transitions, last_value, metrics = rollout_fn(
        state, env_state, roll_key, jnp.asarray(0.0, dtype=jnp.float32)
    )

    # Every slot is mid-episode, never parked on a finished game.
    assert not bool(jnp.any(next_env.done))
    # 120 steps at ~50 plies a game means each of the 8 slots finished at least one.
    assert float(metrics["episodes_completed"]) >= cfg.ppo.batch_games
    # No step is wasted stepping an already-done env.
    assert bool(jnp.all(transitions.active))
    assert int(metrics["active_steps"]) == cfg.ppo.rollout_steps * cfg.ppo.batch_games
    assert last_value.shape == (cfg.ppo.batch_games,)


def test_rollout_game_length_matches_the_rules_clock():
    """A round is 44 deck draws plus one turn per discard-pile draw."""
    cfg = _cfg()
    rng = jax.random.PRNGKey(1)
    rng, init_key, reset_key, roll_key = jax.random.split(rng, 4)
    state = create_train_state(cfg, init_key)
    env_state = jax.jit(jax.vmap(reset))(jax.random.split(reset_key, cfg.ppo.batch_games))

    rollout_fn = make_rollout_fn(cfg, policy_by_name("discard_only"))
    _, _, _, metrics = rollout_fn(state, env_state, roll_key, jnp.asarray(0.0, dtype=jnp.float32))

    # 44 deck draws is the floor; discard-pile draws only ever extend a round.
    assert float(metrics["game_length_mean"]) >= 44.0
    assert float(metrics["max_steps_rate"]) == 0.0


# --- shaping anneal accounting --------------------------------------------


def test_shaping_anneal_tracks_learner_actions_not_padded_steps():
    cfg = _cfg()
    cfg.reward.potential_shaping_initial = 1.0
    cfg.reward.potential_shaping_final = 0.0
    cfg.reward.potential_shaping_anneal_steps = 1_000

    assert shaping_coefficient(cfg, 0) == pytest.approx(1.0)
    assert shaping_coefficient(cfg, 500) == pytest.approx(0.5)
    assert shaping_coefficient(cfg, 1_000) == pytest.approx(0.0)
    assert shaping_coefficient(cfg, 10_000) == pytest.approx(0.0)
