# coding=utf-8
# Copyright 2021 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""Multi-agent goal-seeking task with many static obstacles.
"""

import gym_minigrid.minigrid as minigrid
from . import multigrid
from . import register


class ClutteredMultiGrid(multigrid.MultiGridEnv):
    """Goal seeking environment with obstacles."""

    def __init__(
        self,
        size=15,
        n_agents=3,
        n_clutter=25,
        randomize_goal=True,
        agent_view_size=5,
        max_steps=250,
        walls_are_lava=False,
        **kwargs
    ):
        self.n_clutter = n_clutter
        self.randomize_goal = randomize_goal
        self.walls_are_lava = walls_are_lava
        super().__init__(
            grid_size=size,
            max_steps=max_steps,
            n_agents=n_agents,
            agent_view_size=agent_view_size,
            **kwargs
        )

    def _gen_grid(self, width, height):
        self.grid = multigrid.Grid(width, height)
        self.grid.wall_rect(0, 0, width, height)
        if self.randomize_goal:
            self.place_obj(minigrid.Goal(), max_tries=100)
        else:
            self.put_obj(minigrid.Goal(), width - 2, height - 2)
        for _ in range(self.n_clutter):
            if self.walls_are_lava:
                self.place_obj(minigrid.Lava(), max_tries=100)
            else:
                self.place_obj(minigrid.Wall(), max_tries=100)

        self.place_agent()

        self.mission = "get to the green square"

    def step(self, action):
        obs, reward, done, info = multigrid.MultiGridEnv.step(self, action)
        return obs, reward, done, info


class ClutteredMultiGridSingle6x6(ClutteredMultiGrid):

    def __init__(self):
        super().__init__(
            n_agents=1,
            size=6,
            n_clutter=5,
            randomize_goal=True,
            agent_view_size=5,
            max_steps=50,
        )


class ClutteredMultiGridSingle(ClutteredMultiGrid):

    def __init__(self):
        super().__init__(
            n_agents=1,
            size=15,
            n_clutter=25,
            randomize_goal=True,
            agent_view_size=5,
            max_steps=250,
        )


class Cluttered40Minigrid(ClutteredMultiGrid):

    def __init__(self):
        super().__init__(n_agents=1, n_clutter=40, minigrid_mode=True)


class Cluttered10Minigrid(ClutteredMultiGrid):

    def __init__(self):
        super().__init__(n_agents=1, n_clutter=10, minigrid_mode=True)


class Cluttered50Minigrid(ClutteredMultiGrid):

    def __init__(self):
        super().__init__(n_agents=1, n_clutter=50, minigrid_mode=True)


class Cluttered5Minigrid(ClutteredMultiGrid):

    def __init__(self):
        super().__init__(n_agents=1, n_clutter=5, minigrid_mode=True)


class Cluttered1MinigridMini(ClutteredMultiGrid):

    def __init__(self):
        super().__init__(n_agents=1, n_clutter=1, minigrid_mode=True, size=6)


class Cluttered6MinigridMini(ClutteredMultiGrid):

    def __init__(self):
        super().__init__(n_agents=1, n_clutter=6, minigrid_mode=True, size=6)


class Cluttered7MinigridMini(ClutteredMultiGrid):

    def __init__(self):
        super().__init__(n_agents=1, n_clutter=7, minigrid_mode=True, size=6)


class ClutteredMinigridLava(ClutteredMultiGrid):

    def __init__(self):
        super().__init__(n_agents=1, walls_are_lava=True, minigrid_mode=True)


class ClutteredMinigridLavaMini(ClutteredMultiGrid):

    def __init__(self):
        super().__init__(
            n_agents=1, n_clutter=4, walls_are_lava=True, size=6, minigrid_mode=True
        )


class ClutteredMinigridLavaMedium(ClutteredMultiGrid):

    def __init__(self):
        super().__init__(
            n_agents=1, n_clutter=15, walls_are_lava=True, size=10, minigrid_mode=True
        )


class Cluttered15MinigridMedium(ClutteredMultiGrid):

    def __init__(self):
        super().__init__(n_agents=1, n_clutter=15, minigrid_mode=True, size=10)


if hasattr(__loader__, "name"):
    module_path = __loader__.name
elif hasattr(__loader__, "fullname"):
    module_path = __loader__.fullname

register.register(
    env_id="MultiGrid-Cluttered-v0", entry_point=module_path + ":ClutteredMultiGrid"
)

register.register(
    env_id="MultiGrid-Cluttered-Single-v0",
    entry_point=module_path + ":ClutteredMultiGridSingle",
)

register.register(
    env_id="MultiGrid-Cluttered-Single-6x6-v0",
    entry_point=module_path + ":ClutteredMultiGridSingle6x6",
)

register.register(
    env_id="MultiGrid-Cluttered40-Minigrid-v0",
    entry_point=module_path + ":Cluttered40Minigrid",
)

register.register(
    env_id="MultiGrid-Cluttered10-Minigrid-v0",
    entry_point=module_path + ":Cluttered10Minigrid",
)

register.register(
    env_id="MultiGrid-Cluttered50-Minigrid-v0",
    entry_point=module_path + ":Cluttered50Minigrid",
)

register.register(
    env_id="MultiGrid-Cluttered5-Minigrid-v0",
    entry_point=module_path + ":Cluttered5Minigrid",
)

register.register(
    env_id="MultiGrid-MiniCluttered1-Minigrid-v0",
    entry_point=module_path + ":Cluttered1MinigridMini",
)

register.register(
    env_id="MultiGrid-MiniCluttered6-Minigrid-v0",
    entry_point=module_path + ":Cluttered6MinigridMini",
)

register.register(
    env_id="MultiGrid-MiniCluttered7-Minigrid-v0",
    entry_point=module_path + ":Cluttered7MinigridMini",
)

register.register(
    env_id="MultiGrid-Cluttered-Lava-Minigrid-v0",
    entry_point=module_path + ":ClutteredMinigridLava",
)

register.register(
    env_id="MultiGrid-MiniCluttered-Lava-Minigrid-v0",
    entry_point=module_path + ":ClutteredMinigridLavaMini",
)

register.register(
    env_id="MultiGrid-MediumCluttered-Lava-Minigrid-v0",
    entry_point=module_path + ":ClutteredMinigridLavaMedium",
)

register.register(
    env_id="MultiGrid-MediumCluttered15-Minigrid-v0",
    entry_point=module_path + ":Cluttered15MinigridMedium",
)
