"""Environment wrappers for CarRacing-v3.

We stack four standard wrappers on top of the raw env:
- SkipFrame:        repeat each action k times to reduce decision frequency
- GrayScaleResize:  RGB(96, 96, 3) -> Gray(84, 84) to shrink the input
- FrameStack:       stack the last k frames so the agent can perceive motion
- make_env:         factory that returns a fully wrapped environment

After wrapping, an observation has shape (4, 84, 84) uint8.
"""

from collections import deque

import cv2
import gymnasium as gym
import numpy as np


class SkipFrame(gym.Wrapper):
  """Repeat the action ``k`` times and accumulate the rewards."""

  def __init__(self, env: gym.Env, k: int = 4):
      super().__init__(env)
      self.k = k

  def step(self, action):
      total_reward = 0.0
      terminated = False
      truncated = False
      info = {}
      obs = None
      for _ in range(self.k):
          obs, reward, terminated, truncated, info = self.env.step(action)
          total_reward += reward
          if terminated or truncated:
              break
      return obs, total_reward, terminated, truncated, info


class GrayScaleResize(gym.ObservationWrapper):
  """Convert RGB frames to grayscale and resize to ``size`` x ``size``."""

  def __init__(self, env: gym.Env, size: int = 84):
      super().__init__(env)
      self.size = size
      self.observation_space = gym.spaces.Box(
          low=0, high=255, shape=(size, size), dtype=np.uint8
      )

  def observation(self, obs):
      gray = cv2.cvtColor(obs, cv2.COLOR_RGB2GRAY)
      resized = cv2.resize(gray, (self.size, self.size), interpolation=cv2.INTER_AREA)
      return resized


class FrameStack(gym.Wrapper):
  """Stack the most recent ``k`` frames along a new leading axis."""

  def __init__(self, env: gym.Env, k: int = 4):
      super().__init__(env)
      self.k = k
      self.frames = deque(maxlen=k)
      h, w = env.observation_space.shape
      self.observation_space = gym.spaces.Box(
          low=0, high=255, shape=(k, h, w), dtype=np.uint8
      )

  def reset(self, **kwargs):
      obs, info = self.env.reset(**kwargs)
      self.frames.clear()
      for _ in range(self.k):
          self.frames.append(obs)
      return self._get_obs(), info

  def step(self, action):
      obs, reward, terminated, truncated, info = self.env.step(action)
      self.frames.append(obs)
      return self._get_obs(), reward, terminated, truncated, info

  def _get_obs(self):
      return np.stack(self.frames, axis=0)


def make_env(seed: int = 0, skip: int = 4, size: int = 84, stack: int = 4) -> gym.Env:
  """Create a CarRacing-v3 env with our standard preprocessing stack.

  Returns an environment whose observations are uint8 arrays of shape
  (stack, size, size), ready to feed into a CNN backbone.
  """
  env = gym.make("CarRacing-v3", continuous=False, render_mode="rgb_array")
  env = SkipFrame(env, k=skip)
  env = GrayScaleResize(env, size=size)
  env = FrameStack(env, k=stack)
  env.action_space.seed(seed)
  env.reset(seed=seed)
  return env