"""Tests for DoseFractionation-v0 environment."""

import gymnasium as gym
import numpy as np
import pytest

import oncosim  # noqa: F401


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

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

    def test_reset_valid(self, env):
        obs, info = env.reset(seed=42)
        assert isinstance(obs, dict)
        assert "fraction_number" in obs
        assert "tumor_volume" in obs
        assert "tcp" in obs

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

    def test_action_space_bounds(self, env):
        assert isinstance(env.action_space, gym.spaces.Box)
        assert env.action_space.low[0] == pytest.approx(0.5)
        assert env.action_space.high[0] == pytest.approx(4.0)

    def test_step_valid_action(self, env):
        env.reset(seed=42)
        action = np.array([2.0], dtype=np.float32)
        obs, reward, terminated, truncated, info = env.step(action)
        assert isinstance(reward, float)
        assert np.isfinite(reward)

    def test_episode_length(self, env):
        env.reset(seed=42)
        steps = 0
        done = False
        while not done and steps < 50:
            action = np.array([2.0], dtype=np.float32)
            _, _, terminated, truncated, _ = env.step(action)
            done = terminated or truncated
            steps += 1
        assert steps <= 35  # max_fractions=30 + some buffer

    def test_tumor_volume_decreases(self, env):
        obs, _ = env.reset(seed=42)
        initial_vol = obs["tumor_volume"][0]
        # Apply high dose
        obs, _, _, _, _ = env.step(np.array([3.5], dtype=np.float32))
        # Volume should decrease (cell kill > regrowth for high dose)
        assert obs["tumor_volume"][0] < initial_vol

    def test_tcp_increases(self, env):
        env.reset(seed=42)
        tcps = []
        for _ in range(15):
            obs, _, terminated, _, info = env.step(np.array([2.0], dtype=np.float32))
            tcps.append(info["tcp"])
            if terminated:
                break
        # TCP should generally increase with sufficient treatment
        assert tcps[-1] >= tcps[0]

    def test_cumulative_dose_increases(self, env):
        env.reset(seed=42)
        prev_dose = 0.0
        for _ in range(3):
            obs, _, _, _, _ = env.step(np.array([2.0], dtype=np.float32))
            assert obs["cumulative_dose"][0] > prev_dose
            prev_dose = obs["cumulative_dose"][0]

    def test_high_dose_toxicity(self, env):
        env.reset(seed=42)
        # Give very high doses repeatedly
        for _ in range(30):
            _, _, terminated, _, info = env.step(np.array([4.0], dtype=np.float32))
            if terminated:
                break
        # NTCP should be significant after high doses
        assert info["ntcp"] > 0.1 or info["tcp"] > 0.9

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

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

    def test_fraction_number_increments(self, env):
        env.reset(seed=42)
        for i in range(5):
            obs, _, _, _, _ = env.step(np.array([2.0], dtype=np.float32))
            assert obs["fraction_number"][0] == i + 1

    def test_info_dict_keys(self, env):
        env.reset(seed=42)
        _, _, _, _, info = env.step(np.array([2.0], dtype=np.float32))
        assert "fraction" in info
        assert "dose_given" in info
        assert "tcp" in info
        assert "ntcp" in info

    def test_action_clipping(self, env):
        env.reset(seed=42)
        # Action outside bounds should be clipped
        obs, _, _, _, info = env.step(np.array([10.0], dtype=np.float32))
        assert info["dose_given"] <= 4.0
