import os
import pathlib
import yaml

from enum import IntEnum
from gym_minigrid.roomgrid import MiniGridEnv, spaces, WorldObj, fill_coords, point_in_circle, point_in_rect, Key, Wall, Door, Grid, Goal
from gym_minigrid.register import register
from gym_minigrid.minigrid import COLORS, OBJECT_TO_IDX, COLOR_TO_IDX


class Gradient(WorldObj):

    def __init__(self, max_value, color):
        super().__init__('wall', color)
        self.current_value = 0
        self.max_value = max_value
        self.timer_color = COLORS[color]  # np.array([255, 255, 255])

    def see_behind(self):
        return False

    def update(self, value):
        self.current_value = value

    def encode(self):
        return OBJECT_TO_IDX[self.type], COLOR_TO_IDX[self.color], self.current_value

    def render(self, img):
        rate = self.current_value / self.max_value
        fill_coords(img, point_in_rect(0, 1, 0, 1), self.timer_color * rate)


class Apple(WorldObj):

    def __init__(self):
        super().__init__('ball', 'green')
        self.eaten = False
        self.reward_to_collect = 0

    def render(self, img):
        fill_coords(img, point_in_circle(0.5, 0.5, 0.31), COLORS['green'])

    def collect(self):
        reward_to_collect = self.reward_to_collect
        self.reward_to_collect = 0
        return reward_to_collect

    def toggle(self, env, pos):
        # Replace with wall
        self.eaten = True
        env.grid.set(*pos, None)
        self.reward_to_collect = 1
        return True


class DistractorEnv(MiniGridEnv):

    def __init__(self, seed=None, max_steps=None):
        filename = 'distractor_maps.yaml'
        with open(os.path.join(pathlib.Path(__file__).parent.resolve(), filename)) as f:
            self.stages = [stage.strip('\n').split('\n') for stage in yaml.load(f, yaml.FullLoader)['stages']]

        self.stage_steps = [15, 50, 10]
        self.max_reward = 5.
        self.items = None
        self.wait_step = None
        self.current_stage = None
        self.steps_in_stage = None
        self.max_steps = sum(self.stage_steps)
        self.apple_positions = [((2, 1), (1, 1)), ((1, 1), (2, 1)), ((2, 1), (2, 1))]
        width = len(self.stages[0])
        height = len(self.stages[0][0])

        class Actions(IntEnum):
            left = 0
            right = 1
            forward = 2
            pickup = 3
            toggle = 4

        super().__init__(
            width=width,
            height=height,
            max_steps=self.max_steps,
            seed=seed,
        )

        # Action enumeration for this environment
        self.actions = Actions

        # Actions are discrete integer values
        self.action_space = spaces.Discrete(len(self.actions))

    def _gen_grid(self, width, height):
        self.grid = Grid(width, height)
        self.current_stage = 0
        self.wait_step = False
        self.episode_events = list()
        self.items = dict()
        self.switch_to_stage(self.current_stage)
        self.mission = 'do nothing'

    def switch_to_stage(self, stage):
        self.agent_pos = None

        for x, row in enumerate(self.stages[stage]):
            for y, cell in enumerate(row):
                try:
                    self.grid.set(x, y, None)
                    if cell == '.':
                        self.grid.set(x, y, None)
                    elif cell == 'w':
                        self.place_obj(Wall(), top=(x, y), size=(1, 1), max_tries=100)
                    elif cell == 'k':
                        self.items['key'] = Key(color='yellow')
                        self.place_obj(self.items['key'], top=(x, y), size=(1, 1), max_tries=100)
                    elif cell == 'r':
                        self.items['reward'] = Gradient(self.max_reward, color='green')
                        self.place_obj(self.items['reward'], top=(x, y), size=(1, 1), max_tries=100)
                    elif cell == 't':
                        self.items['timer'] = Gradient(self.stage_steps[stage], color='grey')
                        self.place_obj(self.items['timer'], top=(x, y), size=(1, 1), max_tries=100)
                    elif cell == 'x':
                        self.place_agent(top=(x, y), size=(1, 1), max_tries=100)
                    elif cell == 'a':
                        if 'apples' not in self.items:
                            self.items['apples'] = list()

                        self.items['apples'].append(Apple())
                        self.place_obj(self.items['apples'][-1], top=(x, y), size=(1, 1), max_tries=100)
                    elif cell == 'd':
                        # self.items['door'] = Door(is_locked=True, color='yellow')
                        self.items['door'] = Goal()
                        self.place_obj(self.items['door'], top=(x, y), size=(1, 1))
                    else:
                        raise NotImplementedError()
                except RecursionError as e:
                    print('error!!!', x, y, cell, stage)
                    raise e

        self.steps_in_stage = 0

    def step(self, action):
        prev_agent_state = (self.agent_dir, self.agent_pos[0], self.agent_pos[1])
        obs, reward, done, info = super().step(action)
        self.steps_in_stage += 1
        info['events'] = list()

        if reward > 0:
            reward = self.max_reward if self.carrying else 0

        reward += sum(apple.collect() for apple in self.items.get('apples', list()))
        if 'reward' in self.items:
            self.items['reward'].update(reward)

        if 'timer' in self.items:
            self.items['timer'].update(self.steps_in_stage)

        # Switch stage
        if not done and self.steps_in_stage >= self.stage_steps[self.current_stage]:
            self.current_stage += 1
            self.switch_to_stage(self.current_stage)

        # if 'door' in self.items and not self.items['door'].is_locked:
        #     reward += self.max_reward
        #     done = True

        # Allow the agent to see effects of actions
        # if self.wait_step:
        #     reward += self.max_reward
        #     done = True

        if self.agent_dir != prev_agent_state[0]:
            info['events'].append(('agent', 'move'))

        if self.agent_pos[0] != prev_agent_state[1] or self.agent_pos[1] != prev_agent_state[2]:
            info['events'].append(('agent', 'move'))

        # if 'door' in self.items and not self.items['door'].is_locked:
        #     self.wait_step = True

        self.episode_events.append(info['events'])

        to_index = [('nothing', 'nothing'), ('agent', 'move'), ('key', 'picked'), ('apple', 'eaten'), ('door', 'opened')]
        info['event_index'] = to_index.index(info['events'][0]) if info['events'] else 0
        return obs, reward, done, info


register(id='MiniGrid-DistractorEnv-v0', entry_point='gym_minigrid.envs:DistractorEnv')
