"""Wrappers for SB3 compatibility and observation normalization."""

from __future__ import annotations

import gymnasium as gym
import numpy as np
from gymnasium import spaces


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

    def __init__(self, env: gym.Env) -> None:
        super().__init__(env)
        obs_space = env.observation_space
        assert isinstance(obs_space, spaces.Box)
        self._low = obs_space.low.copy()
        self._range = (obs_space.high - obs_space.low).copy()
        self._range[self._range == 0] = 1.0

        self.observation_space = spaces.Box(
            low=-np.ones_like(obs_space.low),
            high=np.ones_like(obs_space.high),
            dtype=np.float32,
        )

    def observation(self, obs: np.ndarray) -> np.ndarray:
        normalized = 2.0 * (obs - self._low) / self._range - 1.0
        return np.clip(normalized, -1.0, 1.0).astype(np.float32)


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

    def __init__(self, env: gym.Env) -> None:
        super().__init__(env)
        assert isinstance(env.action_space, spaces.Box)

    def action(self, action: np.ndarray) -> np.ndarray:
        return np.clip(
            action,
            self.env.action_space.low,
            self.env.action_space.high,
        )
