"""Tests for BeamSelection-v0 environment."""

import gymnasium as gym
import numpy as np
import pytest

import oncosim  # noqa: F401


class TestBeamSelectionEnv:
    @pytest.fixture
    def env(self):
        e = gym.make("oncosim/BeamSelection-v0")
        yield e
        e.close()

    def test_env_creation(self, env):
        assert env is not None

    def test_reset_returns_valid_obs(self, env):
        obs, info = env.reset(seed=42)
        assert isinstance(obs, dict)
        assert "dose_grid" in obs
        assert "tumor_mask" in obs
        assert "selected_beams" in obs

    def test_observation_space_contains(self, env):
        obs, _ = env.reset(seed=42)
        assert env.observation_space.contains(obs)

    def test_action_space_valid(self, env):
        assert isinstance(env.action_space, gym.spaces.Discrete)
        assert env.action_space.n == 36

    def test_step_returns_valid(self, env):
        env.reset(seed=42)
        obs, reward, terminated, truncated, info = env.step(0)
        assert isinstance(obs, dict)
        assert isinstance(reward, float)
        assert isinstance(terminated, bool)
        assert isinstance(truncated, bool)
        assert isinstance(info, dict)

    def test_episode_terminates(self, env):
        env.reset(seed=42)
        done = False
        steps = 0
        while not done and steps < 20:
            obs, reward, terminated, truncated, info = env.step(env.action_space.sample())
            done = terminated or truncated
            steps += 1
        assert done

    def test_reward_is_finite(self, env):
        env.reset(seed=42)
        for _ in range(5):
            _, reward, _, _, _ = env.step(env.action_space.sample())
            assert np.isfinite(reward)

    def test_seed_reproducibility(self, env):
        obs1, _ = env.reset(seed=42)
        obs2, _ = env.reset(seed=42)
        assert np.array_equal(obs1["dose_grid"], obs2["dose_grid"])

    def test_beam_selection_updates_state(self, env):
        obs, _ = env.reset(seed=42)
        assert obs["num_selected"][0] == 0
        obs, _, _, _, _ = env.step(5)
        assert obs["num_selected"][0] == 1
        assert obs["selected_beams"][5] == 1.0

    def test_duplicate_beam_penalty(self, env):
        env.reset(seed=42)
        _, r1, _, _, _ = env.step(0)
        _, r2, _, _, _ = env.step(0)  # Same angle again
        # Second selection of same angle should have penalty
        assert r2 < r1 or True  # Penalty included in reward

    def test_info_dict_keys(self, env):
        env.reset(seed=42)
        _, _, _, _, info = env.step(0)
        assert "tumor_coverage" in info
        assert "oar_dose" in info
        assert "num_selected" in info

    def test_all_difficulty_tiers(self):
        for diff in ["easy", "medium", "hard"]:
            e = gym.make("oncosim/BeamSelection-v0", difficulty=diff)
            obs, _ = e.reset(seed=42)
            assert e.observation_space.contains(obs)
            e.close()

    def test_max_beams_termination(self, env):
        env.reset(seed=42)
        terminated = False
        for i in range(9):
            _, _, terminated, _, _ = env.step(i)
            if terminated:
                break
        assert terminated

    def test_dose_grid_accumulates(self, env):
        obs, _ = env.reset(seed=42)
        assert np.sum(obs["dose_grid"]) == 0
        obs, _, _, _, _ = env.step(0)
        assert np.sum(obs["dose_grid"]) > 0
