"""Tests for AnestheSim Gymnasium environments."""

import gymnasium as gym
import numpy as np
import pytest

import anesthesim  # noqa: F401


# ─── PropofolControl ───

class TestPropofolControl:
    ENV_ID = "anesthesim/PropofolControl-v0"

    def test_make(self):
        env = gym.make(self.ENV_ID)
        assert env is not None
        env.close()

    def test_reset_returns_obs(self):
        env = gym.make(self.ENV_ID)
        obs, info = env.reset(seed=42)
        assert obs.shape == (6,)
        assert isinstance(info, dict)
        env.close()

    def test_step_returns_correct_types(self):
        env = gym.make(self.ENV_ID)
        env.reset(seed=42)
        action = env.action_space.sample()
        obs, reward, terminated, truncated, info = env.step(action)
        assert obs.shape == (6,)
        assert isinstance(reward, float)
        assert isinstance(terminated, bool)
        assert isinstance(truncated, bool)
        env.close()

    def test_observation_in_space(self):
        env = gym.make(self.ENV_ID)
        obs, _ = env.reset(seed=42)
        assert env.observation_space.contains(obs)
        env.close()

    def test_episode_truncates(self):
        env = gym.make(self.ENV_ID)
        env.reset(seed=42)
        truncated = False
        for _ in range(250):
            action = np.array([100.0], dtype=np.float32)
            _, _, terminated, truncated, _ = env.step(action)
            if terminated or truncated:
                break
        assert truncated or terminated
        env.close()

    def test_difficulty_easy(self):
        env = gym.make(self.ENV_ID, difficulty="easy")
        obs, _ = env.reset(seed=42)
        assert obs.shape == (6,)
        env.close()

    def test_difficulty_hard(self):
        env = gym.make(self.ENV_ID, difficulty="hard")
        obs, _ = env.reset(seed=42)
        assert obs.shape == (6,)
        env.close()

    def test_zero_action_safe(self):
        env = gym.make(self.ENV_ID)
        env.reset(seed=42)
        for _ in range(10):
            obs, reward, terminated, truncated, _ = env.step(np.array([0.0]))
            assert not terminated
        env.close()

    def test_info_contains_bis(self):
        env = gym.make(self.ENV_ID)
        env.reset(seed=42)
        _, _, _, _, info = env.step(np.array([100.0]))
        assert "BIS" in info
        assert "in_target" in info
        env.close()

    def test_bis_responds_to_drug(self):
        env = gym.make(self.ENV_ID, difficulty="easy")
        env.reset(seed=42)
        for _ in range(50):
            _, _, _, _, info = env.step(np.array([150.0]))
        assert info["BIS"] < 97.0
        env.close()


# ─── RemifentanilDosing ───

class TestRemifentanilDosing:
    ENV_ID = "anesthesim/RemifentanilDosing-v0"

    def test_make(self):
        env = gym.make(self.ENV_ID)
        assert env is not None
        env.close()

    def test_reset_returns_obs(self):
        env = gym.make(self.ENV_ID)
        obs, info = env.reset(seed=42)
        assert obs.shape == (6,)
        env.close()

    def test_step_returns_correct_types(self):
        env = gym.make(self.ENV_ID)
        env.reset(seed=42)
        action = env.action_space.sample()
        obs, reward, terminated, truncated, info = env.step(action)
        assert obs.shape == (6,)
        assert isinstance(reward, float)
        env.close()

    def test_observation_in_space(self):
        env = gym.make(self.ENV_ID)
        obs, _ = env.reset(seed=42)
        assert env.observation_space.contains(obs)
        env.close()

    def test_episode_truncates(self):
        env = gym.make(self.ENV_ID)
        env.reset(seed=42)
        for _ in range(250):
            _, _, terminated, truncated, _ = env.step(np.array([0.1]))
            if terminated or truncated:
                break
        assert truncated or terminated
        env.close()

    def test_difficulty_easy(self):
        env = gym.make(self.ENV_ID, difficulty="easy")
        obs, _ = env.reset(seed=42)
        assert obs.shape == (6,)
        env.close()

    def test_difficulty_hard(self):
        env = gym.make(self.ENV_ID, difficulty="hard")
        obs, _ = env.reset(seed=42)
        assert obs.shape == (6,)
        env.close()

    def test_info_contains_error(self):
        env = gym.make(self.ENV_ID)
        env.reset(seed=42)
        _, _, _, _, info = env.step(np.array([0.2]))
        assert "error" in info
        assert "on_target" in info
        env.close()

    def test_concentration_responds(self):
        env = gym.make(self.ENV_ID, difficulty="easy")
        env.reset(seed=42)
        for _ in range(30):
            _, _, _, _, info = env.step(np.array([0.3]))
        assert info["effect_site_concentration"] > 0.0
        env.close()

    def test_zero_infusion_no_concentration(self):
        env = gym.make(self.ENV_ID, difficulty="easy")
        env.reset(seed=42)
        _, _, _, _, info = env.step(np.array([0.0]))
        assert info["effect_site_concentration"] == pytest.approx(0.0, abs=0.01)
        env.close()


# ─── CombinedAnesthesia ───

class TestCombinedAnesthesia:
    ENV_ID = "anesthesim/CombinedAnesthesia-v0"

    def test_make(self):
        env = gym.make(self.ENV_ID)
        assert env is not None
        env.close()

    def test_reset_returns_obs(self):
        env = gym.make(self.ENV_ID)
        obs, info = env.reset(seed=42)
        assert obs.shape == (8,)
        env.close()

    def test_step_returns_correct_types(self):
        env = gym.make(self.ENV_ID)
        env.reset(seed=42)
        action = env.action_space.sample()
        obs, reward, terminated, truncated, info = env.step(action)
        assert obs.shape == (8,)
        assert isinstance(reward, float)
        env.close()

    def test_action_space_shape(self):
        env = gym.make(self.ENV_ID)
        assert env.action_space.shape == (2,)
        env.close()

    def test_observation_in_space(self):
        env = gym.make(self.ENV_ID)
        obs, _ = env.reset(seed=42)
        assert env.observation_space.contains(obs)
        env.close()

    def test_episode_truncates(self):
        env = gym.make(self.ENV_ID)
        env.reset(seed=42)
        for _ in range(250):
            _, _, terminated, truncated, _ = env.step(np.array([100.0, 0.15]))
            if terminated or truncated:
                break
        assert truncated or terminated
        env.close()

    def test_difficulty_easy(self):
        env = gym.make(self.ENV_ID, difficulty="easy")
        obs, _ = env.reset(seed=42)
        assert obs.shape == (8,)
        env.close()

    def test_difficulty_hard(self):
        env = gym.make(self.ENV_ID, difficulty="hard")
        obs, _ = env.reset(seed=42)
        assert obs.shape == (8,)
        env.close()

    def test_info_contains_both_drugs(self):
        env = gym.make(self.ENV_ID)
        env.reset(seed=42)
        _, _, _, _, info = env.step(np.array([100.0, 0.2]))
        assert "propofol_Ce" in info
        assert "remifentanil_Ce" in info
        assert "BIS" in info
        env.close()

    def test_synergy_effect(self):
        env = gym.make(self.ENV_ID, difficulty="easy")
        env.reset(seed=42)
        for _ in range(50):
            _, _, _, _, info = env.step(np.array([120.0, 0.2]))
        # Both drugs should produce lower BIS than propofol alone
        assert info["BIS"] < 60.0
        env.close()
