from typing import Sequence

import gym.spaces as spaces

from centralized_verification.envs.fast_grid_world import FastGridWorld


class FastGridWorldPartialObs(FastGridWorld):
    def agent_obs_spaces(self) -> Sequence[spaces.Space]:
        return [spaces.Discrete(len(self.grid_posns))] * self.num_agents

    def project_obs(self, state) -> Sequence[int]:
        return state  # Using the fact that a state is really just the position of each agent in a list
