"""RL 扩展验收：特征/动作布局、训练环境闭环与模型接口守卫。

RL 依赖属于可选组 [dependency-groups] rl；未安装 numpy/gymnasium 时
全部跳过（运行闭环不依赖它们）。
"""

import unittest
from pathlib import Path

pytest = None
try:
    import pytest
    numpy = pytest.importorskip("numpy")
    pytest.importorskip("gymnasium")
except Exception:  # pragma: no cover - 无 RL 依赖环境
    numpy = None

if numpy is not None:
    from mcm26b.belief import CandidateRegion
    from mcm26b.config import ProblemConfig
    from mcm26b.rl._constants import GLOBAL_DIM, NUM_ACTIONS, OBS_DIM, PER_ACTION_DIM
    from mcm26b.rl.env import LocalizationEnv
    from mcm26b.rl.features import build_observation, candidate_actions
    from mcm26b.rl.policy import default_model_path

CONFIG = ProblemConfig()
STATION = (0.0, 0.0)
BEARING = 30.0


def _region():
    region = CandidateRegion.initial(STATION, CONFIG)
    region.clip_wedge(STATION, BEARING, 1.005)
    return region


@unittest.skipIf(numpy is None, "RL 依赖组未安装")
class FeatureLayoutTests(unittest.TestCase):
    def test_action_candidates_have_stable_count_and_order(self):
        region = _region()
        first = candidate_actions(region, STATION, CONFIG, [])
        self.assertEqual(len(first), NUM_ACTIONS)
        second = candidate_actions(region, STATION, CONFIG, [])
        self.assertEqual(first, second)
        self.assertEqual(first[0], region.mec()[0])   # 0 号候选 = MEC 圆心

    def test_observation_dimension_matches_model_contract(self):
        region = _region()
        obs = build_observation(region, STATION, 3, CONFIG, problem=3,
                                blocked=[], virtual_time_s=0.0)
        self.assertEqual(obs.shape, (OBS_DIM,))
        self.assertEqual(OBS_DIM, GLOBAL_DIM + NUM_ACTIONS * PER_ACTION_DIM)

    def test_blocked_changes_observation_but_not_layout(self):
        region = _region()
        clean = build_observation(region, STATION, 3, CONFIG, problem=4,
                                  blocked=[], virtual_time_s=0.0)
        blocked = build_observation(region, STATION, 3, CONFIG, problem=4,
                                    blocked=[(100.0, 0.0)], virtual_time_s=0.0)
        self.assertEqual(clean.shape, blocked.shape)
        self.assertFalse(numpy.allclose(clean, blocked))


@unittest.skipIf(numpy is None, "RL 依赖组未安装")
class TrainingEnvironmentTests(unittest.TestCase):
    def test_q3_episode_reset_and_step_shapes(self):
        env = LocalizationEnv(problem=3)
        obs, info = env.reset(seed=1)
        self.assertEqual(obs.shape, (OBS_DIM,))
        self.assertIn("source", info)
        obs, reward, terminated, truncated, info = env.step(0)
        self.assertEqual(obs.shape, (OBS_DIM,))
        self.assertIsInstance(reward, float)
        self.assertFalse(terminated and truncated)

    def test_q4_directional_reports_blocked_or_fallback(self):
        env = LocalizationEnv(problem=4, directional_probability=1.0)
        _, info = env.reset(seed=7)
        self.assertTrue(info["source"]["directional"])
        terminated = truncated = False
        for _ in range(env.max_episode_steps + 2):
            _, _, terminated, truncated, info = env.step(
                env.action_space.sample())
            if terminated or truncated:
                break
        self.assertTrue(terminated or truncated)
        # 定向分支：要么清除，要么走过后备清除/截断语义。
        self.assertIn("fallback_attempts", info)

    def test_missing_model_raises_clear_error(self):
        from mcm26b.rl.policy import PointSelectionPolicy
        with self.assertRaises(RuntimeError):
            PointSelectionPolicy("output/rl/does-not-exist.zip")

    def test_default_model_paths(self):
        self.assertEqual(default_model_path(3), Path("output/rl/q3_ppo.zip"))
        self.assertEqual(default_model_path(4), Path("output/rl/q4_ppo.zip"))


if __name__ == "__main__":
    unittest.main()