"""Utility functions for environment, device detection, and TensorBoard."""
import gymnasium as gym
import numpy as np
import torch
from collections import deque


class GrayScaleWrapper(gym.ObservationWrapper):
    """Convert RGB observation to grayscale."""

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

    def observation(self, obs):
        # RGB to grayscale: weighted average
        gray = 0.299 * obs[:, :, 0] + 0.587 * obs[:, :, 1] + 0.114 * obs[:, :, 2]
        return gray.astype(np.uint8)


class ResizeWrapper(gym.ObservationWrapper):
    """Resize observation to target size."""

    def __init__(self, env, size=(84, 84)):
        super().__init__(env)
        self.size = size

    def observation(self, obs):
        import cv2
        return cv2.resize(obs, self.size, interpolation=cv2.INTER_AREA)


class FrameStackWrapper(gym.ObservationWrapper):
    """Stack last N frames."""

    def __init__(self, env, num_stack=4):
        super().__init__(env)
        self.num_stack = num_stack
        self.frames = deque(maxlen=num_stack)
        obs_shape = env.observation_space.shape
        self.observation_space = gym.spaces.Box(
            low=0, high=255,
            shape=(num_stack, *obs_shape[-2:]),
            dtype=np.uint8
        )

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

    def observation(self, obs):
        self.frames.append(obs)
        return self._get_observation()

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


def make_env(env_id="CarRacing-v3", gray_scale=True, resize=True, frame_stack=4):
    """Create preprocessed CarRacing environment."""
    env = gym.make(env_id, render_mode="rgb_array")
    if resize:
        env = ResizeWrapper(env, size=(84, 84))
    if gray_scale:
        env = GrayScaleWrapper(env)
    if frame_stack > 1:
        env = FrameStackWrapper(env, num_stack=frame_stack)
    return env


def get_device():
    """Detect and return available device."""
    if torch.cuda.is_available():
        device = torch.device("cuda")
        print(f"Using GPU: {torch.cuda.get_device_name(0)}")
    else:
        device = torch.device("cpu")
        print("Using CPU")
    return device


def preprocess_obs(obs):
    """Ensure observation is in correct format for network."""
    if len(obs.shape) == 2:  # single channel
        obs = np.expand_dims(obs, axis=0)
    return obs