"""Tests for environment wrappers."""

import gymnasium as gym
import numpy as np
import pytest

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


class TestNormalizeObservation:
    def test_wraps_env(self):
        env = NormalizeObservation(gym.make("anesthesim/PropofolControl-v0"))
        obs, _ = env.reset(seed=42)
        assert obs.shape == (6,)
        assert np.all(obs >= -1.0)
        assert np.all(obs <= 1.0)
        env.close()

    def test_normalized_after_step(self):
        env = NormalizeObservation(gym.make("anesthesim/PropofolControl-v0"))
        env.reset(seed=42)
        obs, _, _, _, _ = env.step(np.array([100.0]))
        assert np.all(obs >= -1.0)
        assert np.all(obs <= 1.0)
        env.close()

    def test_works_with_remifentanil(self):
        env = NormalizeObservation(gym.make("anesthesim/RemifentanilDosing-v0"))
        obs, _ = env.reset(seed=42)
        assert np.all(obs >= -1.0) and np.all(obs <= 1.0)
        env.close()

    def test_works_with_combined(self):
        env = NormalizeObservation(gym.make("anesthesim/CombinedAnesthesia-v0"))
        obs, _ = env.reset(seed=42)
        assert obs.shape == (8,)
        env.close()


class TestClipAction:
    def test_clips_high(self):
        env = ClipAction(gym.make("anesthesim/PropofolControl-v0"))
        env.reset(seed=42)
        obs, _, _, _, _ = env.step(np.array([999.0]))
        assert obs.shape == (6,)
        env.close()

    def test_clips_low(self):
        env = ClipAction(gym.make("anesthesim/PropofolControl-v0"))
        env.reset(seed=42)
        obs, _, _, _, _ = env.step(np.array([-100.0]))
        assert obs.shape == (6,)
        env.close()

    def test_combined_wrappers(self):
        base = gym.make("anesthesim/PropofolControl-v0")
        env = ClipAction(NormalizeObservation(base))
        obs, _ = env.reset(seed=42)
        assert np.all(obs >= -1.0) and np.all(obs <= 1.0)
        obs2, _, _, _, _ = env.step(np.array([300.0]))
        assert np.all(obs2 >= -1.0) and np.all(obs2 <= 1.0)
        env.close()
