"""Tests for environment wrappers."""

import gymnasium as gym
import numpy as np
import pytest

import oncosim  # noqa: F401
from oncosim.envs.wrappers import FlattenObsWrapper, NormalizeRewardWrapper


class TestFlattenObsWrapper:
    def test_flatten_obs_type(self):
        env = gym.make("oncosim/BeamSelection-v0")
        wrapped = FlattenObsWrapper(env)
        obs, _ = wrapped.reset(seed=42)
        assert isinstance(obs, np.ndarray)
        wrapped.close()

    def test_flatten_obs_shape(self):
        env = gym.make("oncosim/BeamSelection-v0")
        wrapped = FlattenObsWrapper(env)
        obs, _ = wrapped.reset(seed=42)
        assert obs.ndim == 1
        assert len(obs) > 0
        wrapped.close()

    def test_flatten_contains(self):
        env = gym.make("oncosim/BeamSelection-v0")
        wrapped = FlattenObsWrapper(env)
        obs, _ = wrapped.reset(seed=42)
        assert wrapped.observation_space.contains(obs)
        wrapped.close()

    def test_flatten_step(self):
        env = gym.make("oncosim/BeamSelection-v0")
        wrapped = FlattenObsWrapper(env)
        wrapped.reset(seed=42)
        obs, reward, terminated, truncated, info = wrapped.step(0)
        assert isinstance(obs, np.ndarray)
        assert obs.ndim == 1
        wrapped.close()

    def test_flatten_dose_fractionation(self):
        env = gym.make("oncosim/DoseFractionation-v0")
        wrapped = FlattenObsWrapper(env)
        obs, _ = wrapped.reset(seed=42)
        assert isinstance(obs, np.ndarray)
        wrapped.close()

    def test_flatten_adaptive_rt(self):
        env = gym.make("oncosim/AdaptiveRT-v0")
        wrapped = FlattenObsWrapper(env)
        obs, _ = wrapped.reset(seed=42)
        assert isinstance(obs, np.ndarray)
        wrapped.close()


class TestNormalizeRewardWrapper:
    def test_normalize_returns_float(self):
        env = gym.make("oncosim/BeamSelection-v0")
        wrapped = NormalizeRewardWrapper(env)
        wrapped.reset(seed=42)
        _, reward, _, _, _ = wrapped.step(0)
        assert isinstance(reward, float)
        wrapped.close()

    def test_normalize_finite(self):
        env = gym.make("oncosim/BeamSelection-v0")
        wrapped = NormalizeRewardWrapper(env)
        wrapped.reset(seed=42)
        for _ in range(5):
            _, reward, terminated, _, _ = wrapped.step(env.action_space.sample())
            assert np.isfinite(reward)
            if terminated:
                break
        wrapped.close()
