from vizdoom import * 
from .base.abenv import Env
from ..utils.error import WadNotFoundError

class VizdoomEnv(Env):

    RES_640X480 = ScreenResolution.RES_640X480 
    RES_320X240 = ScreenResolution.RES_320X240
    RES_160X120 = ScreenResolution.RES_160X120 

    def __init__(self, wad, doommap="map01", _id="vizdoom", render=True, reset_done=True, num_steps=1000, res=RES_160X120, auto_map=True):
        if wad != None and wad != "":
            super().__init__(_id, render, reset_done)
            self.num_steps = num_steps
            self.wad = wad
            self.doommap = doommap
            self.game = DoomGame()
            self.auto_map = auto_map

            self.res = res
            if self.res == VizdoomEnv.RES_160X120:
                self.res_w = 160
                self.res_h = 120
            elif self.res == VizdoomEnv.RES_320X240:
                self.res_w = 320 
                self.res_h = 240 
            elif self.res == VizdoomEnv.RES_640X480:
                self.res_w = 640 
                self.res_h = 480
            else:
                raise UnsupportedVizDoomRes(self.res + " is an unsupported vizdoom resolution, use only VizdoomEnv.RES_640X480, VizdoomEnv.RES_320X240 or VizdoomEnv.RES_160X120.")

        else:
            raise WadNotFoundError("A wad file is needed.")

    def start(self):
        '''
            Start vizdoom 
        '''

        #Init object, set wad and map
        self.game.set_doom_scenario_path(self.wad)
        if (self.doommap != None):
            self.game.set_doom_map(self.doommap)

        #Set available controls
        self.game.add_available_button(Button.ATTACK)
        self.game.add_available_button(Button.USE)
        self.game.add_available_button(Button.MOVE_FORWARD)
        self.game.add_available_button(Button.MOVE_BACKWARD)
        self.game.add_available_button(Button.MOVE_RIGHT)
        self.game.add_available_button(Button.MOVE_LEFT)
        self.game.add_available_button(Button.TURN_RIGHT)
        self.game.add_available_button(Button.TURN_LEFT)

        #Set which variables will be set in observation
        #self.game.get_state().game_variables will be a nparray
        #where every position will be one of the variables
        #below. All variables are floats.
        #The list is ordered on the same order the variable
        #are added here
        self.game.add_available_game_variable(GameVariable.KILLCOUNT)
        self.game.add_available_game_variable(GameVariable.ITEMCOUNT)
        self.game.add_available_game_variable(GameVariable.SECRETCOUNT)
        self.game.add_available_game_variable(GameVariable.DEATHCOUNT)
        self.game.add_available_game_variable(GameVariable.HITCOUNT)
        self.game.add_available_game_variable(GameVariable.HITS_TAKEN)
        self.game.add_available_game_variable(GameVariable.DAMAGECOUNT)
        self.game.add_available_game_variable(GameVariable.DAMAGE_TAKEN)
        self.game.add_available_game_variable(GameVariable.HEALTH)
        self.game.add_available_game_variable(GameVariable.ARMOR)
        self.game.add_available_game_variable(GameVariable.DEAD)
        self.game.add_available_game_variable(GameVariable.ATTACK_READY)
        self.game.add_available_game_variable(GameVariable.SELECTED_WEAPON)
        self.game.add_available_game_variable(GameVariable.SELECTED_WEAPON_AMMO)
        self.game.add_available_game_variable(GameVariable.POSITION_X)
        self.game.add_available_game_variable(GameVariable.POSITION_Y)
        self.game.add_available_game_variable(GameVariable.POSITION_Z)

        #Set mini_map available in game obs
        #To get automap: observation.automap_buffer
        #Example: https://github.com/mwydmuch/ViZDoom/blob/master/examples/python/automap.py
        if self.auto_map:
            self.game.set_automap_buffer_enabled(True)
            self.game.set_automap_mode(AutomapMode.OBJECTS_WITH_SIZE)



        #To get gamescreen on state, you should:
        #game.get_state().screen_buffer
        #Some graphic properties
        #screen format may be:
        # game.set_screen_format(ScreenFormat.RGB24)
        # game.set_screen_format(ScreenFormat.RGBA32)
        # game.set_screen_format(ScreenFormat.ARGB32)
        # game.set_screen_format(ScreenFormat.BGRA32)
        # game.set_screen_format(ScreenFormat.ABGR32)
        # game.set_screen_format(ScreenFormat.GRAY8)
        # self.game.set_screen_format(ScreenFormat.RGB24)
        self.game.set_screen_format(ScreenFormat.GRAY8)
        self.game.set_screen_resolution(self.res)
        self.game.set_render_hud(True)
        self.game.set_render_crosshair(False)
        self.game.set_render_weapon(True)
        self.game.set_render_decals(False)
        self.game.set_render_particles(False)
        self.game.set_window_visible(self.render)

        #Set reward -1 for every movement the agent makes
        self.game.set_living_reward(-1)


        #Set episode limit and a little delay to spawn monsters
        self.game.set_episode_timeout(self.num_steps)
        self.game.set_episode_start_time(0)


        self.game.init()

    def step(self, action):
        '''
            Executes action and return observation, reward and done
        '''

        reward = self.game.make_action(action)
        done = self.game.is_episode_finished()
        state = self.game.get_state()
        if state != None:
            self.observation = state
        
        return self.observation, reward, done

    def reset(self):
        '''
            Resets the game to a new episode
        '''
        self.close()
        self.start()
        if self.game.get_state() != None:
            self.observation = self.game.get_state()
        return self.observation

    def close(self):
        ''' 
            Stops vizdoom
        '''
        if self.game != None:
            self.game.close()

    def restart(self):
        self.reset()

    def get_screen_width(self):
        return self.res_w 

    def get_screen_height(self):
        return self.res_h 

