"""modified from gymnax, to update to gymnasium 1.0.0a2."""

from typing import Optional, Tuple, Union, List

import chex
import gymnasium
import jax.random
from gymnasium.core import ActType, ObsType, RenderFrame
from gymnasium.vector.utils import batch_space
from gymnax.environments.spaces import Space, Discrete, Box
from gymnax.environments.environment import Environment, EnvParams
from copy import deepcopy

from typing import Tuple, Sequence, Any, Dict
import chex
import numpy as np
from gymnasium import spaces as gspc


def gymnax_space_to_gym_space(space: Space) -> gspc.Space:
    """Convert Gymnax space to equivalent Gym space"""
    if isinstance(space, Discrete):
        return gspc.Discrete(space.n)
    elif isinstance(space, Box):
        low = (
            float(space.low)
            if (np.isscalar(space.low) or space.low.size == 1)
            else np.array(space.low)
        )
        high = (
            float(space.high)
            if (np.isscalar(space.high) or space.low.size == 1)
            else np.array(space.high)
        )
        return gspc.Box(low, high, space.shape, space.dtype)
    elif isinstance(space, Dict):
        return gspc.Dict({k: gymnax_space_to_gym_space(v) for k, v in space.spaces})
    elif isinstance(space, Tuple):
        return gspc.Tuple(space.spaces)
    else:
        raise NotImplementedError(
            f"Conversion of {space.__class__.__name__} not supported"
        )


class GymnaxToGymWrapper(gymnasium.Env):
    def __init__(
        self,
        env: Environment,
        params: Optional[EnvParams] = None,
        seed: Optional[int] = None,
    ):
        """Wrap Gymnax environment as OOP Gym environment

        Args:
            env: Gymnax Environment instance
            params: If provided, gymnax EnvParams for environment (otherwise uses default)
            seed: If provided, seed for JAX PRNG (otherwise picks 0)
        """
        super().__init__()
        self._env = deepcopy(env)
        self.env_params = params if params is not None else env.default_params
        self.metadata.update(
            {
                "name": env.name,
                "render_modes": ["human", "rgb_array"]
                if hasattr(env, "render")
                else [],
            }
        )
        self.rng: chex.PRNGKey = jax.random.PRNGKey(0)  # Placeholder
        self._seed(seed)
        _, self.env_state = self._env.reset(self.rng, self.env_params)

    @property
    def action_space(self):
        """Dynamically adjust action space depending on params"""
        return gymnax_space_to_gym_space(self._env.action_space(self.env_params))

    @property
    def observation_space(self):
        """Dynamically adjust state space depending on params"""
        return gymnax_space_to_gym_space(self._env.observation_space(self.env_params))

    def _seed(self, seed: Optional[int] = None):
        """Set RNG seed (or use 0)"""
        self.rng = jax.random.PRNGKey(seed or 0)

    def step(
        self, action: ActType
    ) -> Union[
        Tuple[ObsType, float, bool, bool, dict],
        Tuple[ObsType, float, bool, dict],
    ]:
        """Step environment, follow new step API"""
        self.rng, step_key = jax.random.split(self.rng)
        o, self.env_state, r, d, info = self._env.step(
            step_key, self.env_state, action, self.env_params
        )
        return o, r, d, d, info

    def reset(
        self,
        *,
        seed: Optional[int] = None,
        return_info: bool = False,
        options: Optional[dict] = None,
    ) -> Tuple[ObsType, dict]:
        """Reset environment, update parameters and seed if provided"""
        if seed is not None:
            self._seed(seed)
        if options is not None:
            self.env_params = options.get(
                "env_params", self.env_params
            )  # Allow changing environment parameters on reset
        self.rng, reset_key = jax.random.split(self.rng)
        o, self.env_state = self._env.reset(reset_key, self.env_params)
        return o, {}

    def render(self, mode="human") -> Optional[Union[RenderFrame, List[RenderFrame]]]:
        """use underlying environment rendering if it exists, otherwise return None"""
        return getattr(self._env, "render", lambda x, y: None)(
            self.env_state, self.env_params
        )


class GymnaxToVectorGymWrapper(gymnasium.vector.VectorEnv):
    def __init__(
        self,
        env: Environment,
        num_envs: int = 1,
        params: Optional[EnvParams] = None,
        seed: Optional[int] = None,
    ):
        """Wrap Gymnax environment as OOP Gym Vector Environment

        Args:
            env: Gymnax Environment instance
            num_envs: Desired number of environments to run in parallel
            params: If provided, gymnax EnvParams for environment (otherwise uses default)
            seed: If provided, seed for JAX PRNG (otherwise picks 0)
        """
        self._env = deepcopy(env)
        self.num_envs = num_envs
        self.is_vector_env = True
        self.new_step_api = True
        self.closed = False
        self.viewer = None
        self.rng: chex.PRNGKey = jax.random.PRNGKey(0)  # Placeholder
        self._seed(seed)
        # Jit-of-vmap is faster than vmap-of-jit. Map over leading axis of all but env params
        self._env.reset = jax.jit(jax.vmap(self._env.reset, in_axes=(0, None)))
        self._env.step = jax.jit(jax.vmap(self._env.step, in_axes=(0, 0, 0, None)))
        self.env_params = params if params is not None else env.default_params
        _, self.env_state = self._env.reset(self.rng, self.env_params)  # Placeholder
        self._batched_rng_split = jax.jit(
            jax.vmap(jax.random.split, in_axes=0, out_axes=1)
        )  # Split all rng keys

    @property
    def single_action_space(self):
        """Dynamically adjust action space depending on params"""
        return gymnax_space_to_gym_space(self._env.action_space(self.env_params))

    @property
    def single_observation_space(self):
        """Dynamically adjust state space depending on params"""
        return gymnax_space_to_gym_space(self._env.observation_space(self.env_params))

    @property
    def action_space(self):
        """Dynamically adjust action space depending on params"""
        return batch_space(self.single_action_space, self.num_envs)

    @property
    def observation_space(self):
        """Dynamically adjust state space depending on params"""
        return batch_space(self.single_observation_space, self.num_envs)

    def _seed(self, seed: Optional[int] = None):
        """Set RNG seed (or use 0)"""
        self.rng = jax.random.split(
            jax.random.PRNGKey(seed or 0), self.num_envs
        )  # 1 RNG per env

    def reset(
        self,
        *,
        seed: Optional[Union[int, List[int]]] = None,
        return_info: bool = False,
        options: Optional[dict] = None,
    ) -> Tuple[ObsType, dict]:
        """Reset environment, update parameters and seed if provided"""
        if seed is not None:
            self._seed(seed)
        if options is not None:
            self.env_params = options.get(
                "env_params", self.env_params
            )  # Allow changing environment parameters on reset
        self.rng, reset_key = self._batched_rng_split(self.rng)  # Split all keys
        o, self.env_state = self._env.reset(reset_key, self.env_params)
        return o, {}

    def step(self, action: ActType) -> Tuple[ObsType, float, bool, bool, dict]:
        """Step environment, follow new step API"""
        self.rng, step_key = self._batched_rng_split(self.rng)
        o, self.env_state, r, d, info = self._env.step(
            step_key, self.env_state, action, self.env_params
        )
        return o, r, d, d, info

    def render(self, mode="human") -> Optional[Union[RenderFrame, List[RenderFrame]]]:
        """use underlying environment rendering if it exists (for first environment), otherwise return None"""
        return getattr(self._env, "render", lambda x, y: None)(
            jax.tree_map(lambda x: x[0], self.env_state), self.env_params
        )
