'''
Classes that follows a gym-like interface and implements stage two of the Messenger
environment.
'''

import json
import random
from collections import namedtuple
from pathlib import Path
from os import environ
import re

# hack to stop PyGame from printing to stdout
environ["PYGAME_HIDE_SUPPORT_PROMPT"] = "hide"

import numpy as np
from vgdl.interfaces.gym import VGDLEnv

from messenger.envs.base import MessengerEnv
from messenger.envs.config import Entity
import messenger.envs.config as config
from messenger.envs.manual import TextManual
from messenger.envs.utils import games_from_json

# Positions of the entities
Position = namedtuple('Position', ["x", "y"])

# specifies the game variant (e.g. chasing enemy, fleeing message, stationary goal)
# path is path to the vgdl domain file describing the variant.
GameVariant = namedtuple(
    "GameVariant", ["path", "enemy_type", "message_type", "goal_type"]
)


class Grid:
    '''
    Class which makes it easier to build a grid observation from the dict state
    return by VGDLEnv.
    '''
    def __init__(self, layers, shuffle=True):
        '''
        layers:
            Each add() operation will place a separate entity in a new layer.
            Thus, this is the upper-limit to the number of items to be added.
        shuffle:
            Place each items in a random order.
        '''
        self.grid = np.zeros((config.STATE_HEIGHT, config.STATE_WIDTH, layers))
        self.order = list(range(layers)) # insertion order
        if shuffle:
            random.shuffle(self.order)
        self.layers = layers
        self.entity_count = 0

    def add(self, entity:Entity, position:Position):
        '''
        Add entity entity and position position.
        '''
        assert self.entity_count < self.layers, \
            f"Tried to add entity no. {self.entity_count} with {self.layers} layers."

        self.grid[position.y, position.x, self.order[self.entity_count]] = entity.id
        self.entity_count += 1


class StageTwo(MessengerEnv):
    '''
    Full messenger environment with mobile sprites. Uses Py-VGDL as game engine.
    To avoid the need to instantiate a large number of games, (since there are
    P(12,3) = 1320 possible entity to role assignments) We apply a wrapper on top
    of the text and game state which masks the role archetypes (enemy, message goal)
    into entities (e.g. alien, knight, mage).
    '''

    def __init__(self, split:str, shuffle_obs=True):
        super().__init__()
        self.shuffle_obs = shuffle_obs # shuffle the entity layers

        this_folder = Path(__file__).parent
        # Get the games and manual
        games_json_path = this_folder.joinpath("games.json")
        if "train" in split and "mc" in split: # multi-combination games
            game_split = "train_multi_comb"
            text_json_path = this_folder.joinpath("texts", "text_train.json")
        elif "train" in split and "sc" in split: # single-combination games
            game_split = "train_single_comb"
            text_json_path = this_folder.joinpath("texts", "text_train.json")
        elif "val" in split:
            game_split = "val"
            text_json_path = this_folder.joinpath("texts", "text_val.json")
        elif "test" in split:
            game_split = "test"
            text_json_path = this_folder.joinpath("texts", "text_test.json")
        else:
            raise Exception(f"Split: {split} not understood.")

        # list of Game namedtuples
        self.all_games = games_from_json(json_path=games_json_path, split=game_split)
        self.text_manual = TextManual(json_path=text_json_path)

        # get the folder that has the game variants and init_states
        if "test" in split and "se" not in split: # new dynamics (se for state estimation)
            vgdl_files = this_folder.joinpath("vgdl_files", "stage_2_nd")
        else: # training dynamics
            vgdl_files = this_folder.joinpath("vgdl_files", "stage_2")
        
        # get the file paths to possible starting states
        self.init_states = [
            str(path) for path in vgdl_files.joinpath("init_states").glob("*.txt")
        ]
        # get all the game variants
        self.game_variants = [
            self._get_variant(path) for path in vgdl_files.joinpath("variants").glob("*.txt")
        ]

        # entities tracked by VGDLEnv
        self.notable_sprites = ["enemy", "message", "goal", "no_message", "with_message"]
        self.env = None # the VGDLEnv

    def _get_variant(self, variant_file:Path) -> GameVariant:
        '''
        Return the GameVariant for the variant specified by variant_file. 
        Searches through the vgdl code to find the correct type:
        {chaser, fleeing, immovable}
        '''

        code = variant_file.read_text()
        return GameVariant(
            path = str(variant_file),
            enemy_type = re.search(r'enemy > (\S+)', code)[1].lower(),
            message_type = re.search(r'message > (\S+)', code)[1].lower(),
            goal_type = re.search(r'goal > (\S+)', code)[1].lower()
        )

    def _convert_obs(self, vgdl_obs):
        '''
        Return a grid built from the vgdl observation which is a
        KeyValueObservation object (see vgdl code for details).
        '''
        entity_locs = Grid(layers=3, shuffle=self.shuffle_obs)
        avatar_locs = Grid(layers=1)

        # try to add each entity one by one, if it's not there move on.
        if 'enemy.1' in vgdl_obs:
            entity_locs.add(self.game.enemy, Position(*vgdl_obs['enemy.1']['position']))
        if 'message.1' in vgdl_obs:
            entity_locs.add(self.game.message, Position(*vgdl_obs['message.1']['position']))
        else:
            # advance the entity counter, Oracle model requires special order.
            # TODO: maybe used named layers to make this more understandable.
            entity_locs.entity_count += 1
        if 'goal.1' in vgdl_obs:
            entity_locs.add(self.game.goal, Position(*vgdl_obs['goal.1']['position']))

        if 'no_message.1' in vgdl_obs:
            '''
            Due to a quirk in VGDL, the avatar is no_message if it starts as no_message
            even if the avatar may have acquired the message at a later point.
            To check, if it has a message, check that the class vector corresponding to
            with_message is == 1.
            '''
            avatar_pos = Position(*vgdl_obs['no_message.1']['position'])
            # with_key is last position, see self.notable_sprites
            if vgdl_obs['no_message.1']['class'][-1] == 1:
                avatar = config.WITH_MESSAGE
            else:
                avatar = config.NO_MESSAGE

        elif "with_message.1" in vgdl_obs:
            # this case only occurs if avatar begins as with_message at start of episode
            avatar_pos = Position(*vgdl_obs['with_message.1']['position'])
            avatar = config.WITH_MESSAGE

        else: # the avatar is not in observation, so is probably dead
            return {"entities": entity_locs.grid, "avatar": avatar_locs.grid}

        avatar_locs.add(avatar, avatar_pos) # if not dead, add it.

        return {"entities": entity_locs.grid, "avatar": avatar_locs.grid}

    def reset(self, **kwargs):
        '''
        Resets the current environment. NOTE: We remake the environment each time.
        This is a workaround to a bug in py-vgdl, where env.reset() does not
        properly reset the environment. kwargs go to get_document().
        '''

        self.game = random.choice(self.all_games) # (e.g. enemy-alien, message-knight, goal - bear)

        # choose the game variant (e.g. enmey-chasing, message-fleeing, goal-static)
        # and initial starting location of the entities.
        variant = random.choice(self.game_variants)
        init_state = random.choice(self.init_states) # inital state file

        # args that will go into VGDL Env.
        self._envargs = {
            'game_file': variant.path,
            'level_file': init_state,
            'notable_sprites': self.notable_sprites.copy(),
            'obs_type': 'objects', # track the objects
            'block_size': 34  # rendering block size
        }
        self.env = VGDLEnv(**self._envargs)
        vgdl_obs = self.env.reset()

        manual = self.text_manual.get_document(
            enemy=self.game.enemy.name,
            message=self.game.message.name,
            goal=self.game.goal.name,
            enemy_type=variant.enemy_type,
            message_type=variant.message_type,
            goal_type=variant.goal_type,
            **kwargs
        )
            
        return self._convert_obs(vgdl_obs), manual

    def step(self, action):
        vgdl_obs, reward, done, info = self.env.step(action)
        return self._convert_obs(vgdl_obs), reward, done, info
    