"""Tests for the UltrafiltrationControl-v0 Gymnasium environment."""

import gymnasium as gym
import numpy as np
import pytest

import nephrosim  # noqa: F401
from nephrosim.envs.ultrafiltration_control import UltrafiltrationControlEnv


@pytest.fixture
def env():
    e = gym.make("nephrosim/UltrafiltrationControl-v0")
    yield e
    e.close()


class TestUltrafiltrationControlEnv:
    def test_observation_space_shape(self, env):
        obs, _ = env.reset(seed=42)
        assert obs.shape == (8,)

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

    def test_action_space_shape(self, env):
        assert env.action_space.shape == (1,)

    def test_step_returns_five_tuple(self, env):
        env.reset(seed=42)
        result = env.step(np.array([0.0], dtype=np.float32))
        assert len(result) == 5

    def test_episode_terminates(self, env):
        env.reset(seed=42)
        done = False
        steps = 0
        while not done:
            _, _, terminated, truncated, _ = env.step(
                np.array([0.0], dtype=np.float32)
            )
            done = terminated or truncated
            steps += 1
        assert steps == 240

    def test_bp_in_info(self, env):
        env.reset(seed=42)
        _, _, _, _, info = env.step(np.array([0.0], dtype=np.float32))
        assert "systolic_bp" in info

    def test_hypotension_events_tracked(self, env):
        env.reset(seed=42)
        _, _, _, _, info = env.step(np.array([0.0], dtype=np.float32))
        assert "hypotension_events" in info
        assert info["hypotension_events"] >= 0

    def test_fluid_removal_increases(self, env):
        env.reset(seed=42)
        _, _, _, _, info1 = env.step(np.array([0.5], dtype=np.float32))
        _, _, _, _, info2 = env.step(np.array([0.5], dtype=np.float32))
        assert info2["total_removed_L"] >= info1["total_removed_L"]

    def test_completion_fraction_bounded(self, env):
        env.reset(seed=42)
        for _ in range(60):
            _, _, _, _, info = env.step(np.array([0.0], dtype=np.float32))
        assert info["completion_fraction"] >= 0

    def test_high_uf_triggers_bp_drop(self):
        env = UltrafiltrationControlEnv(difficulty="easy")
        env.reset(seed=42)
        for _ in range(120):
            obs, _, _, _, info = env.step(np.array([1.0], dtype=np.float32))
        # High UF should cause some BP drop from baseline
        assert info["systolic_bp"] < 140

    def test_obs_stays_in_bounds(self, env):
        env.reset(seed=42)
        for _ in range(100):
            obs, _, _, _, _ = env.step(env.action_space.sample())
            assert env.observation_space.contains(obs)

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