import jax.numpy as jnp
from gymnax.environments import EnvParams, spaces

from memorax.utils.wrappers import GymnaxWrapper, MaskObservationWrapper

mask_dims = {
    "ant": {
        "F": list(range(27)),
        "P": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
        "V": [13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26],
    },
    "halfcheetah": {
        "F": list(range(17)),
        "P": [0, 1, 2, 3, 8, 9, 10, 11, 12],
        "V": [4, 5, 6, 7, 13, 14, 15, 16],
    },
    "hopper": {
        "F": list(range(11)),
        "P": [0, 1, 2, 3, 4],
        "V": [5, 6, 7, 8, 9, 10],
    },
    "walker2d": {
        "F": list(range(17)),
        "P": [0, 1, 2, 3, 4, 5, 6, 7],
        "V": [8, 9, 10, 11, 12, 13, 14, 15, 16],
    },
}


class BraxGymnaxWrapper(GymnaxWrapper):
    def __init__(self, env):
        super().__init__(env)
        self.action_size = env.action_size
        self.observation_size = (env.observation_size,)

    @property
    def default_params(self) -> EnvParams:
        return EnvParams(max_steps_in_episode=1000)

    def reset(self, key, params):
        state = self._env.reset(key)
        return state.obs, state

    def step(self, key, state, action, params):
        next_state = self._env.step(state, action)
        return (
            next_state.obs,
            next_state,
            next_state.reward,
            next_state.done.astype(jnp.bool),
            {},
        )

    def observation_space(self, params):
        return spaces.Box(
            low=-jnp.inf,
            high=jnp.inf,
            shape=(self._env.observation_size,),
        )

    def action_space(self, params):
        return spaces.Box(
            low=-1.0,
            high=1.0,
            shape=(self._env.action_size,),
        )


def make(env_id, mode, backend="mjx", **kwargs):
    from brax import envs
    from brax.envs.wrappers.training import AutoResetWrapper, EpisodeWrapper

    env = envs.get_environment(env_name=env_id, backend=backend, **kwargs)
    env = EpisodeWrapper(env, episode_length=1000, action_repeat=1)
    env = AutoResetWrapper(env)
    env = BraxGymnaxWrapper(
        env,
    )
    env = MaskObservationWrapper(env, mask_dims=mask_dims[env_id][mode])

    env_params = env.default_params
    return env, env_params
