"""低层导航运动学与训练入口的行为回归；不验证策略收益。"""

import importlib.util
import math
from pathlib import Path

import pytest

np = pytest.importorskip("numpy")
pytest.importorskip("gymnasium")

from mcm26b.hrl.executor import LowLevelExecutor, navigation_transition
from mcm26b.hrl.low_level_env import LowLevelEnv


def _training_script():
    spec = importlib.util.spec_from_file_location(
        "train_hrl_low_revision", Path(__file__).parents[1]
        / "scripts" / "train_hrl_low.py")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def test_shortened_step_preserves_selected_direction_without_snapping():
    position, moved = navigation_transition((0.0, 0.0), (30.0, 40.0), 0)
    assert position == pytest.approx((50.0, 0.0))
    assert moved == pytest.approx(50.0)
    assert position != (30.0, 40.0)
    with pytest.raises(ValueError):
        navigation_transition((0, 0), (300, 0), 8)


def test_near_goal_reward_and_time_charge_actual_distance():
    env = LowLevelEnv()
    env.reset(seed=1)
    env._position, env._goal = (0.0, 0.0), (30.0, 0.0)
    _, reward, terminated, truncated, info = env.step(0)
    assert terminated and not truncated
    assert info["path_length_m"] == pytest.approx(30.0)
    assert info["step_virtual_time_s"] == pytest.approx(6.0)
    assert reward == pytest.approx(env.reach_bonus - 6.0 / env.step_cost_scale)


def test_expert_reaches_all_200_fixed_seed_targets():
    env = LowLevelEnv()
    for seed in range(200):
        env.reset(seed=seed)
        for _ in range(env.max_steps):
            obs, _, terminated, truncated, info = env.step(env.expert_action())
            assert env.observation_space.contains(obs)
            if terminated or truncated:
                break
        assert terminated, (seed, info)
        assert info["distance_m"] <= env.reach_m
        assert info["virtual_time_s"] == pytest.approx(
            info["path_length_m"] / env.config.cost.speed_mps)


def test_failed_executor_reports_failure_without_free_terminal_move():
    class WestPolicy:
        def select(self, observation):
            return 4

    executor = LowLevelExecutor(WestPolicy(), max_steps=3)
    path = executor.roll((0, 0), (300, 0))
    assert len(path) == 3
    assert path[-1] == pytest.approx((-300, 0))
    assert executor.last_roll_info == pytest.approx({
        "reached": False, "truncated": True,
        "final_distance_m": 600.0, "path_length_m": 300.0})


def test_env_and_executor_share_step_transitions():
    class EastPolicy:
        def select(self, observation):
            return 0

    env = LowLevelEnv(max_steps=3)
    env.reset(seed=3)
    start, goal = (0.0, 0.0), (350.0, 25.0)
    env._position, env._goal = start, goal
    actual = []
    for _ in range(3):
        _, _, _, _, info = env.step(0)
        actual.append(info["position"])
    executor = LowLevelExecutor(EastPolicy(), max_steps=3)
    assert executor.roll(start, goal) == actual
    assert executor.last_roll_info["path_length_m"] == info["path_length_m"]


def test_bc_holdout_keeps_each_episode_wholly_in_one_partition():
    script = _training_script()
    episode_ids = np.repeat(np.arange(20), 5)
    mask = script.bc_validation_mask(episode_ids, seed=7)
    assert mask.any() and not mask.all()
    for episode in np.unique(episode_ids):
        assert len(set(mask[episode_ids == episode])) == 1
    assert np.array_equal(mask, script.bc_validation_mask(episode_ids, seed=7))


def test_bc_copy_and_maskable_ppo_learn_with_two_environments():
    torch = pytest.importorskip("torch")
    pytest.importorskip("sb3_contrib")
    from stable_baselines3.common.vec_env import DummyVecEnv

    previous_threads = torch.get_num_threads()
    torch.set_num_threads(1)
    script = _training_script()
    env = DummyVecEnv([LowLevelEnv, LowLevelEnv])
    try:
        bc = script._bc_module(10, 8, script.HIDDEN)()
        model = script.copy_into_ppo(bc, "cpu", env, seed=9, n_steps=8)
        model.learn(total_timesteps=16)
        assert model.n_envs == 2
        assert model.num_timesteps == 16
        assert np.all(env.env_method("action_masks"))
    finally:
        env.close()
        torch.set_num_threads(previous_threads)