"""Gymnasium wrappers for NephroSim environments."""

import gymnasium as gym
import numpy as np


class NormalizeObservation(gym.ObservationWrapper):
    """Normalize observations to [0, 1] range using observation space bounds."""

    def __init__(self, env: gym.Env):
        super().__init__(env)
        self.obs_low = self.observation_space.low
        self.obs_high = self.observation_space.high
        self.obs_range = np.where(
            self.obs_high - self.obs_low > 0,
            self.obs_high - self.obs_low,
            1.0,
        )
        self.observation_space = gym.spaces.Box(
            low=np.zeros_like(self.obs_low),
            high=np.ones_like(self.obs_high),
            dtype=np.float32,
        )

    def observation(self, obs):
        return ((obs - self.obs_low) / self.obs_range).astype(np.float32)


class ClipAction(gym.ActionWrapper):
    """Clip actions to the valid action space bounds."""

    def __init__(self, env: gym.Env):
        super().__init__(env)

    def action(self, action):
        return np.clip(action, self.action_space.low, self.action_space.high)
