"""Tests for NephroSim Gymnasium wrappers."""

import gymnasium as gym
import numpy as np
import pytest

import nephrosim  # noqa: F401
from nephrosim.envs.wrappers import NormalizeObservation, ClipAction


class TestNormalizeObservation:
    def test_output_range_zero_one(self):
        env = NormalizeObservation(gym.make("nephrosim/UreaClearing-v0"))
        obs, _ = env.reset(seed=42)
        assert np.all(obs >= -0.01)
        assert np.all(obs <= 1.01)
        env.close()

    def test_shape_preserved(self):
        base = gym.make("nephrosim/UreaClearing-v0")
        wrapped = NormalizeObservation(base)
        obs_base, _ = base.reset(seed=42)
        obs_wrap, _ = wrapped.reset(seed=42)
        assert obs_base.shape == obs_wrap.shape
        wrapped.close()

    def test_step_returns_normalized(self):
        env = NormalizeObservation(gym.make("nephrosim/UreaClearing-v0"))
        env.reset(seed=42)
        obs, _, _, _, _ = env.step(env.action_space.sample())
        assert np.all(obs >= -0.01)
        assert np.all(obs <= 1.01)
        env.close()

    def test_observation_space_updated(self):
        env = NormalizeObservation(gym.make("nephrosim/UreaClearing-v0"))
        assert np.allclose(env.observation_space.low, 0.0)
        assert np.allclose(env.observation_space.high, 1.0)
        env.close()


class TestClipAction:
    def test_action_clipped(self):
        env = ClipAction(gym.make("nephrosim/UreaClearing-v0"))
        env.reset(seed=42)
        # Action outside bounds
        obs, _, _, _, _ = env.step(np.array([5.0, -5.0], dtype=np.float32))
        assert obs is not None  # didn't crash
        env.close()

    def test_valid_action_unchanged(self):
        env = ClipAction(gym.make("nephrosim/UreaClearing-v0"))
        env.reset(seed=42)
        action = np.array([0.5, -0.3], dtype=np.float32)
        clipped = env.action(action)
        assert np.allclose(action, clipped)
        env.close()

    def test_out_of_bounds_action_clipped(self):
        env = ClipAction(gym.make("nephrosim/UreaClearing-v0"))
        action = np.array([2.0, -3.0], dtype=np.float32)
        clipped = env.action(action)
        assert clipped[0] == pytest.approx(1.0)
        assert clipped[1] == pytest.approx(-1.0)
        env.close()
