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

import gymnasium as gym
import numpy as np
import pytest

import nephrosim  # noqa: F401
from nephrosim.envs.phosphate_mgmt import PhosphateMgmtEnv


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


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

    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_episode_length(self, env):
        env.reset(seed=42)
        steps = 0
        done = False
        while not done:
            _, _, terminated, truncated, _ = env.step(
                np.array([0.5], dtype=np.float32)
            )
            done = terminated or truncated
            steps += 1
        assert steps == 21  # 7 days * 3 periods

    def test_dialysis_sessions_occur(self, env):
        env.reset(seed=42)
        dialysis_count = 0
        for step in range(21):
            _, _, _, _, info = env.step(np.array([0.5], dtype=np.float32))
            if info["is_dialysis_step"]:
                dialysis_count += 1
        assert dialysis_count == 3  # MWF

    def test_session_count_reaches_three(self, env):
        env.reset(seed=42)
        for _ in range(21):
            _, _, _, _, info = env.step(np.array([0.5], dtype=np.float32))
        assert info["session_count"] == 3

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

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

    def test_binder_dose_in_info(self, env):
        env.reset(seed=42)
        _, _, _, _, info = env.step(np.array([0.75], dtype=np.float32))
        assert info["binder_dose"] == pytest.approx(1500.0)

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

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