import numpy as np
from typing import Any, Iterable, Tuple
from .ext import EnvSpec
from .parallel import ParallelEnv
from ..prelude import Action, Array, State
from ..utils.rms import RunningMeanStd


class ParallelEnvWrapper(ParallelEnv[Action, State]):
    def __init__(self, penv: ParallelEnv) -> None:
        self.penv = penv

    def close(self) -> None:
        self.penv.close()

    def reset(self) -> Array[State]:
        return self.penv.reset()

    def step(
            self,
            actions: Iterable[Action]
    ) -> Tuple[Array[State], Array[float], Array[bool], Array[Any]]:
        return self.penv.step(actions)

    def seed(self, seeds: Iterable[int]) -> None:
        self.penv.seed(seeds)

    @property
    def num_envs(self) -> int:
        return self.penv.num_envs

    @property
    def spec(self) -> EnvSpec:
        return self.penv.spec

    def extract(self, states: Iterable[State]) -> Array:
        return self.penv.extract(states)


class FrameStackParallel(ParallelEnvWrapper):
    """Parallel version of atari_wrappers.FrameStack
    """
    def __init__(self, penv: ParallelEnv, nstack: int = 4, dtype: type = np.float32) -> None:
        super().__init__(penv)
        idx = 0
        shape = self.penv.state_dim
        for dim in shape:
            if dim == 1:
                idx += 1
            else:
                break
        self.shape = (nstack, *self.penv.state_dim[idx:])
        self.obs = np.zeros((self.num_envs, *self.shape), dtype=dtype)

    def step(
            self,
            actions: Iterable[Action]
    ) -> Tuple[Array, Array[float], Array[bool], Array[Any]]:
        state, reward, done, info = self.penv.step(actions)
        self.obs = np.roll(self.obs, shift=-1, axis=1)
        for i, _ in filter(lambda t: t[1], enumerate(done)):
            self.obs[i] = 0.0
        self.obs[:, -1] = self.extract(state).squeeze()
        return (self.obs, reward, done, info)

    def reset(self) -> Array[State]:
        self.obs.fill(0)
        state = self.penv.reset()
        self.obs[:, -1] = self.extract(state).squeeze()
        return self.obs

    @property
    def state_dim(self) -> Tuple[int, ...]:
        return self.shape


class NormalizeObs(ParallelEnvWrapper[Action, Array[float]]):
    def __init__(self, penv: ParallelEnv, obs_clip: float = 10.) -> None:
        super().__init__(penv)
        self.obs_clip = obs_clip
        self._rms = RunningMeanStd(shape=self.state_dim)
        self.training_mode = True

    def step(
            self,
            actions: Iterable[Action]
    ) -> Tuple[Array[Array[float]], Array[float], Array[bool], Array[Any]]:
        state, reward, done, info = self.penv.step(actions)
        return self._filter_obs(state), reward, done, info

    def _filter_obs(self, obs: Array[Array[float]]) -> Array[Array[float]]:
        if self.training_mode:
            self._rms.update(obs)  # type: ignore
        obs = np.clip((obs - self._rms.mean) / self._rms.std(), -self.obs_clip, self.obs_clip)
        return obs

    def reset(self) -> Array[Array[float]]:
        obs = self.penv.reset()
        return self._filter_obs(obs)


class NormalizeReward(ParallelEnvWrapper[Action, State]):
    def __init__(self, penv: ParallelEnv, reward_clip: float = 10., gamma: float = 0.99) -> None:
        super().__init__(penv)
        self.reward_clip = reward_clip
        self.gamma = gamma
        self._rms = RunningMeanStd(shape=())
        self.ret = np.zeros(self.num_envs)

    def step(
            self,
            actions: Iterable[Action]
    ) -> Tuple[Array[State], Array[float], Array[bool], Array[Any]]:
        state, reward, done, info = self.penv.step(actions)
        self.ret = self.ret * self.gamma + reward
        self._rms.update(self.ret)
        reward = np.clip(reward / self._rms.std(), -self.reward_clip, self.reward_clip)
        self.ret[done] = 0.0
        return state, reward, done, info

    def reset(self) -> Array[State]:
        self.ret = np.zeros(self.num_envs)
        return self.penv.reset()
