"""Tests for benchmark suite."""

import gymnasium as gym
import numpy as np
import pytest

import oncosim  # noqa: F401
from oncosim.benchmarks.environments import BENCHMARK_CONFIGS, ENV_ID_MAP


class TestBenchmarkConfigs:
    def test_configs_exist(self):
        assert len(BENCHMARK_CONFIGS) == 3
        assert "BeamSelection" in BENCHMARK_CONFIGS
        assert "DoseFractionation" in BENCHMARK_CONFIGS
        assert "AdaptiveRT" in BENCHMARK_CONFIGS

    def test_all_tiers(self):
        for env_name, configs in BENCHMARK_CONFIGS.items():
            for tier in ["trivial", "easy", "medium", "hard", "expert"]:
                assert tier in configs, f"Missing {tier} for {env_name}"

    def test_env_id_map(self):
        assert len(ENV_ID_MAP) == 3
        for name, env_id in ENV_ID_MAP.items():
            assert env_id.startswith("oncosim/")

    def test_configs_create_valid_envs(self):
        for env_name, configs in BENCHMARK_CONFIGS.items():
            config = configs["medium"]
            env_id = ENV_ID_MAP[env_name]
            env = gym.make(env_id, difficulty=config["difficulty"], **config.get("kwargs", {}))
            obs, _ = env.reset(seed=42)
            assert env.observation_space.contains(obs)
            env.close()

    def test_trivial_easier_than_expert(self):
        """Trivial configs should have less penalizing parameters than expert."""
        beam_trivial = BENCHMARK_CONFIGS["BeamSelection"]["trivial"]
        beam_expert = BENCHMARK_CONFIGS["BeamSelection"]["expert"]
        assert beam_trivial["kwargs"]["oar_weight"] < beam_expert["kwargs"]["oar_weight"]

    def test_benchmark_result_format(self):
        from oncosim.benchmarks.runner import run_benchmarks
        results = run_benchmarks(tiers=["trivial"], n_episodes=3)
        assert isinstance(results, dict)
        for env_name, tiers in results.items():
            for tier, res in tiers.items():
                assert "random_mean_reward" in res
                assert np.isfinite(res["random_mean_reward"])

    def test_runner_with_multiple_tiers(self):
        from oncosim.benchmarks.runner import run_benchmarks
        results = run_benchmarks(tiers=["trivial", "easy"], n_episodes=2)
        for env_name in BENCHMARK_CONFIGS:
            assert "trivial" in results[env_name]
            assert "easy" in results[env_name]

    def test_all_envs_creatable(self):
        for env_id in ENV_ID_MAP.values():
            env = gym.make(env_id)
            env.reset(seed=42)
            env.close()
