from __future__ import annotations

import jax
import jax.numpy as jnp

from lost_cities_jax import reset
from lost_cities_jax.league import (
    LeagueConfig,
    PoolMember,
    active_training_pool,
    load_league_config,
    pool_sampling_probabilities,
)
from lost_cities_jax.opponents import policy_by_name
from lost_cities_jax.ppo import (
    JaxPPOConfig,
    NetworkConfig,
    OpponentConfig,
    PPOHyperConfig,
    RunConfig,
    create_train_state,
    make_league_train_iteration,
    sample_league_assignments,
)


def tiny_league_ppo_config(tmp_path) -> JaxPPOConfig:
    return JaxPPOConfig(
        run=RunConfig(
            experiment_name="pytest-jax-ppo-league",
            seed=123,
            total_updates=1,
            checkpoint_every=1,
            artifact_root=str(tmp_path),
        ),
        opponent=OpponentConfig(name="discard_only"),
        network=NetworkConfig(hidden_size=32, num_layers=1),
        ppo=PPOHyperConfig(batch_games=8, rollout_steps=16, epochs=1, minibatches=2),
    )


def test_league_config_loads():
    cfg = load_league_config("configs/jax_ppo/league-v1.yaml")
    assert cfg.experiment_name == "jax-ppo-league-v1"
    assert cfg.cycles == 5
    assert len(cfg.anchors) == 7
    assert any(member.name == "heuristic_expert" for member in cfg.anchors)


def test_pfsp_probabilities_cap_stalling_anchors():
    cfg = LeagueConfig(
        base_config="configs/jax_ppo/ladder-v2-expert.yaml",
        warm_start_checkpoint="unused",
        stalling_anchor_floor=0.02,
        stalling_anchor_cap=0.05,
        uniform_mix=0.1,
    )
    pool = [
        PoolMember("discard_only", "static", stalling=True, recent_win_rate=0.01),
        PoolMember("heuristic_expert", "static", recent_win_rate=0.5),
        PoolMember("snapshot", "checkpoint", anchor=False, recent_win_rate=0.2),
    ]
    probs = pool_sampling_probabilities(pool, cfg)
    assert abs(sum(probs) - 1.0) < 1.0e-6
    assert 0.02 <= probs[0] <= 0.05


def test_active_training_pool_keeps_anchors_and_hard_nonanchors():
    cfg = LeagueConfig(
        base_config="configs/jax_ppo/ladder-v2-expert.yaml",
        warm_start_checkpoint="unused",
        max_active_pool_members=4,
    )
    pool = [
        PoolMember("anchor_a", "static", anchor=True),
        PoolMember("anchor_b", "static", anchor=True),
        PoolMember("easy", "checkpoint", anchor=False, recent_win_rate=0.9),
        PoolMember("hard", "checkpoint", anchor=False, recent_win_rate=0.1),
        PoolMember("medium", "checkpoint", anchor=False, recent_win_rate=0.5),
    ]
    active = active_training_pool(pool, cfg)
    assert [member.name for member in active] == ["anchor_a", "anchor_b", "hard", "medium"]


def test_one_jitted_league_train_iteration_shapes(tmp_path):
    cfg = tiny_league_ppo_config(tmp_path)
    train_state = create_train_state(cfg, jax.random.PRNGKey(0))
    env_state = jax.jit(jax.vmap(reset))(jax.random.split(jax.random.PRNGKey(1), 8))
    probs = jnp.asarray([0.5, 0.5], dtype=jnp.float32)
    assignments = sample_league_assignments(jax.random.PRNGKey(2), 8, probs, 0.5)
    train_iteration = make_league_train_iteration(
        cfg,
        [policy_by_name("discard_only"), policy_by_name("heuristic_expert")],
        probs,
        0.5,
    )

    train_state, env_state, assignments, rng, metrics = train_iteration(
        train_state,
        env_state,
        assignments,
        jax.random.PRNGKey(3),
        jnp.asarray(0.0, dtype=jnp.float32),
    )

    assert env_state.to_move.shape == (8,)
    assert assignments.learner_seat.shape == (8,)
    assert rng.shape == (2,)
    assert "opened_colors_mean" in metrics