"""Wrappers for PeptideGym environments — SB3 compatibility and reward shaping."""
from __future__ import annotations

from collections import OrderedDict

import gymnasium
import numpy as np
from gymnasium import spaces


class FlattenObservation(gymnasium.ObservationWrapper):
    """Flatten a Dict observation space into a single Box.

    Useful for algorithms (e.g. Stable-Baselines3 MLP policies) that
    require a flat observation vector.
    """

    def __init__(self, env: gymnasium.Env) -> None:
        super().__init__(env)
        assert isinstance(env.observation_space, spaces.Dict), (
            "FlattenObservation requires a Dict observation space."
        )

        # Compute total flat dimension
        self._keys: list[str] = []
        self._shapes: list[tuple[int, ...]] = []
        self._dims: list[int] = []
        low_parts: list[np.ndarray] = []
        high_parts: list[np.ndarray] = []

        for key, subspace in env.observation_space.spaces.items():
            assert isinstance(subspace, spaces.Box), (
                f"FlattenObservation only supports Box sub-spaces, got {type(subspace)} for '{key}'."
            )
            self._keys.append(key)
            self._shapes.append(subspace.shape)
            flat_dim = int(np.prod(subspace.shape))
            self._dims.append(flat_dim)
            low_parts.append(subspace.low.flatten())
            high_parts.append(subspace.high.flatten())

        total_dim = sum(self._dims)
        self.observation_space = spaces.Box(
            low=np.concatenate(low_parts),
            high=np.concatenate(high_parts),
            shape=(total_dim,),
            dtype=np.float32,
        )

    def observation(self, observation: dict) -> np.ndarray:
        parts = []
        for key in self._keys:
            arr = np.asarray(observation[key], dtype=np.float32).flatten()
            parts.append(arr)
        return np.concatenate(parts)


class RewardShapingWrapper(gymnasium.RewardWrapper):
    """Scale rewards to the [-1, 1] range using tanh squashing.

    This prevents reward magnitude from dominating the gradient signal
    for policy-gradient methods operating on PeptideGym environments.
    """

    def __init__(self, env: gymnasium.Env, scale: float = 1.0) -> None:
        super().__init__(env)
        self._scale = scale

    def reward(self, reward: float) -> float:
        return float(np.tanh(reward * self._scale))
