"""SB3-compatible wrappers for OncoSim environments."""

from __future__ import annotations

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


class FlattenObsWrapper(gym.ObservationWrapper):
    """Flatten Dict observations into a single Box for SB3 compatibility."""

    def __init__(self, env: gym.Env):
        super().__init__(env)
        # Compute total flat size
        total_size = 0
        self._keys: list[str] = sorted(env.observation_space.spaces.keys())
        self._slices: dict[str, tuple[int, int]] = {}
        for key in self._keys:
            space = env.observation_space[key]
            flat_size = int(np.prod(space.shape))
            self._slices[key] = (total_size, total_size + flat_size)
            total_size += flat_size

        self.observation_space = spaces.Box(
            low=-np.inf, high=np.inf, shape=(total_size,), dtype=np.float64
        )

    def observation(self, obs: dict) -> np.ndarray:
        parts = []
        for key in self._keys:
            parts.append(obs[key].flatten())
        return np.concatenate(parts).astype(np.float64)


class NormalizeRewardWrapper(gym.RewardWrapper):
    """Normalize rewards to approximately [-1, 1] range using running statistics."""

    def __init__(self, env: gym.Env, clip: float = 10.0):
        super().__init__(env)
        self._clip = clip
        self._reward_sum = 0.0
        self._reward_sq_sum = 0.0
        self._count = 0

    def reward(self, reward: float) -> float:
        self._count += 1
        self._reward_sum += reward
        self._reward_sq_sum += reward ** 2

        mean = self._reward_sum / self._count
        var = max(self._reward_sq_sum / self._count - mean ** 2, 1e-8)
        std = np.sqrt(var)

        normalized = (reward - mean) / std
        return float(np.clip(normalized, -self._clip, self._clip))
