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

class CustomEnv(gym.Env):
    """Custom Environment that follows gym interface."""

    metadata = {"render_modes": ["human"], "render_fps": 1}

    def __init__(self, s_dim=0):
        super().__init__()
        self.S_DIM = s_dim
        self.action_space = spaces.Box(low=-1, high=1, shape=(3,), dtype=np.float32)
        self.observation_space = spaces.Box(
            low=-1, high=1, shape=(s_dim,), dtype=np.float32
        )

    def step(self, action):
        return np.zeros(self.S_DIM, dtype=np.float32), float(0), False, False, {}

    def change_s_dim(self, s_dim=0):
        self.S_DIM = s_dim
        self.observation_space = spaces.Box(
            low=-1, high=1, shape=(s_dim,), dtype=np.float32
        )

    def reset(self, seed=None):
        return np.zeros(self.S_DIM, dtype=np.float32), {}

    def close(self):
        self.__del__()