import numpy as np
from gym import Env
from gym.spaces import Discrete


class FixedSequenceEnv(Env):
    def __init__(
            self,
            n_actions=10,
            episode_len=100
    ):
        self.action_space = Discrete(n_actions)
        self.observation_space = Discrete(1)
        self.np_random = np.random.RandomState(0)
        self.episode_len = episode_len
        self.sequence = [self.np_random.randint(0, self.action_space.n)
            for _ in range(self.episode_len)]
        self.time = 0


    def reset(self):
        self.time = 0
        return 0

    def step(self, actions):
        rew = self._get_reward(actions)
        self._choose_next_state()
        done = False
        if self.episode_len and self.time >= self.episode_len:
            done = True

        return 0, rew, done, {}

    def seed(self, seed=None):
        self.np_random.seed(seed)

    def _choose_next_state(self):
        self.time += 1

    def _get_reward(self, actions):
        return 1 if actions == self.sequence[self.time] else 0