import numpy as np
import os
import gym
from gym import error, spaces
from gym import utils
from gym.utils import seeding

try:
    import atari_py
except ImportError as e:
    raise error.DependencyNotInstalled("{}. (HINT: you can install Atari dependencies by running 'pip install gym[atari]'.)".format(e))

import logging
logger = logging.getLogger(__name__)

def to_ram(ale):
    ram_size = ale.getRAMSize()
    ram = np.zeros((ram_size),dtype=np.uint8)
    ale.getRAM(ram)
    return ram

class AtariEnv(gym.Env, utils.EzPickle):
    metadata = {'render.modes': ['human', 'rgb_array']}

    def __init__(self, game='pong', obs_type='ram', frameskip=(2, 5), repeat_action_probability=0.):
        """Frameskip should be either a tuple (indicating a random range to
        choose from, with the top value exclude), or an int."""

        utils.EzPickle.__init__(self, game, obs_type)
        assert obs_type in ('ram', 'image')

        self.game_path = atari_py.get_game_path(game)
        if not os.path.exists(self.game_path):
            raise IOError('You asked for game %s but path %s does not exist'%(game, self.game_path))
        self._obs_type = obs_type
        self.frameskip = frameskip
        self.ale = atari_py.ALEInterface()
        self.viewer = None

        # Tune (or disable) ALE's action repeat:
        # https://github.com/openai/gym/issues/349
        assert isinstance(repeat_action_probability, (float, int)), "Invalid repeat_action_probability: {!r}".format(repeat_action_probability)
        self.ale.setFloat('repeat_action_probability'.encode('utf-8'), repeat_action_probability)

        self._seed()

        (screen_width, screen_height) = self.ale.getScreenDims()

        self._action_set = self.ale.getMinimalActionSet()
        self.action_space = spaces.Discrete(len(self._action_set))

        (screen_width,screen_height) = self.ale.getScreenDims()
        if self._obs_type == 'ram':
            self.observation_space = spaces.Box(low=np.zeros(128), high=np.zeros(128)+255)
        elif self._obs_type == 'image':
            self.observation_space = spaces.Box(low=0, high=255, shape=(screen_height, screen_width, 3))
        else:
            raise error.Error('Unrecognized observation type: {}'.format(self._obs_type))

        ### BEGIN CURIOSITY SEARCH (ADDED BY CMS 10/01/17)
        self._initCS = True # used to enable/disable CS without having to comment out the below
        self.curiositySearch = False
        self.ATARI_BASE_RAM_OFFSET = 128
        self.Y_DOWN = False
        self.CS_ROOM_SWITCH = False
        self.CS_ROOM_DIM_LENGTH = 0
        self.ADDR_ROOM = None
        self.ADDR_ROOM_2 = None
        self.ADDR_POSITION_OFFSET = None
        self.ADDR_ROOM_REFRESH_VALUE = None
        self.SLICE_CS_GRID = False # currently only y-slicing (for room-panning) is supported
        self.POPULATION_MEMORY = False

        if (self._initCS and "seaquest" in game.lower()):
            self.curiositySearch = True
            self.Y_DOWN = True
            self.CS_GRID_WIDTH = 114
            self.CS_GRID_HEIGHT = 96
            self.ADDR_X = int("C6", 16) - self.ATARI_BASE_RAM_OFFSET
            self.ADDR_X_OFFSET = 21
            self.ADDR_Y = int("E1", 16) - self.ATARI_BASE_RAM_OFFSET
            self.ADDR_Y_OFFSET = 13
        elif (self._initCS and "montezuma" in game.lower()):
            self.curiositySearch = True
            self.CS_GRID_WIDTH = 153
            self.CS_GRID_HEIGHT = 122
            self.ADDR_X = int("AA", 16) - self.ATARI_BASE_RAM_OFFSET
            self.ADDR_X_OFFSET = 0
            self.ADDR_Y = int("AB", 16) - self.ATARI_BASE_RAM_OFFSET
            self.ADDR_Y_OFFSET = 134
            self.ADDR_ROOM = int("83", 16) - self.ATARI_BASE_RAM_OFFSET
        elif (self._initCS and "amidar" in game.lower()):
            self.curiositySearch = True
            self.CS_ROOM_SWITCH = True
            self.Y_DOWN = True
            self.CS_GRID_WIDTH = 125
            self.CS_GRID_HEIGHT = 147
            self.ADDR_X = int("C2", 16) - self.ATARI_BASE_RAM_OFFSET
            self.ADDR_X_OFFSET = 6
            self.ADDR_Y = int("BB", 16) - self.ATARI_BASE_RAM_OFFSET
            self.ADDR_Y_OFFSET = 7
            self.ADDR_ROOM = int("B6", 16) - self.ATARI_BASE_RAM_OFFSET # only first two rooms
            self.ADDR_POSITION_OFFSET = int("D7", 16) - self.ATARI_BASE_RAM_OFFSET
        elif (self._initCS and "gravitar" in game.lower()):
            self.curiositySearch = True
            self.Y_DOWN = True
            self.CS_GRID_WIDTH = 154
            self.CS_GRID_HEIGHT = 157
            self.ADDR_X = int("9F", 16) - self.ATARI_BASE_RAM_OFFSET
            self.ADDR_X_OFFSET = 0
            self.ADDR_Y = int("A0", 16) - self.ATARI_BASE_RAM_OFFSET
            self.ADDR_Y_OFFSET = 0
            self.ADDR_ROOM = int("81", 16) - self.ATARI_BASE_RAM_OFFSET
            self.ADDR_ROOM_2 = int("82", 16) - self.ATARI_BASE_RAM_OFFSET
            self.ROOM_DIM_LENGTH = 10
        elif (self._initCS and "bank_heist" in game.lower()):
            self.curiositySearch = True
            self.Y_DOWN = True
            self.CS_GRID_WIDTH = 129
            self.CS_GRID_HEIGHT = 121
            self.ADDR_X = int("9C", 16) - self.ATARI_BASE_RAM_OFFSET
            self.ADDR_X_OFFSET = 12
            self.ADDR_Y = int("88", 16) - self.ATARI_BASE_RAM_OFFSET
            self.ADDR_Y_OFFSET = 9
            self.ADDR_ROOM = int("80", 16) - self.ATARI_BASE_RAM_OFFSET
        elif (self._initCS and "tutankham" in game.lower()):
            self.curiositySearch = True
            self.Y_DOWN = True
            self.CS_GRID_WIDTH = 145
            self.CS_GRID_HEIGHT = 256
            self.ADDR_X = int("87", 16) - self.ATARI_BASE_RAM_OFFSET
            self.ADDR_X_OFFSET = 4
            self.ADDR_Y = int("C1", 16) - self.ATARI_BASE_RAM_OFFSET
            self.ADDR_Y_OFFSET = 0
            self.ADDR_ROOM_REFRESH_VALUE = int("FB", 16) - self.ATARI_BASE_RAM_OFFSET
            self.ROOM_REFRESH_VALUE = int("35", 16)
            self.SLICE_CS_GRID = True
            self.CS_SLICE_WIDTH = 145
            self.CS_SLICE_HEIGHT = 27

        if self.curiositySearch:
            self.csGrid = [None] * 256 # Atari encodes rooms with bytes (limit 255); most unused
            self.csGridPop = [None] * 256 # will never be reset during training
            self.CS_EXPLORED = 255
            self.CS_UNEXPLORED = 0
            if self.SLICE_CS_GRID:
                self.slicedCSGrid = self.CS_EXPLORED * np.ones(shape=(self.CS_SLICE_WIDTH, self.CS_SLICE_HEIGHT, 1), dtype=np.uint8)
                self.slicedCSGridPop = self.CS_EXPLORED * np.ones(shape=(self.CS_SLICE_WIDTH, self.CS_SLICE_HEIGHT, 1), dtype=np.uint8)
            self.INTRINSIC_REWARD = 1
            self.currIntrinsic = 0
            self.currIntrinsicPop = 0
            self.CS_GRID_SIZE = 16
            self.currGrid = None
            self.currGridPop = None
            self.lastCurrRoom = -1
            self.resetRoomForLevel = False
            self.totalSwitchedRooms = 0
            self.roomOverride = -1
        ### END CURIOSITY SEARCH

    def setCSGridSize(self, size):
        '''Sets the CS grid size, overriding the game default. Tiles are always square. Added by CMS 11/1/18.'''
        if self.curiositySearch:
            self.CS_GRID_SIZE = size
            return
        raise RuntimeError("Cannot set CS grid size unless curiosity search was enabled during init.")

    def enablePopulationMemory(self):
        '''Enables population-memory for CS, creating a separate grid on which population exploration is recorded. Added by CMS 22/05/18.'''
        if self.curiositySearch:
            self.POPULATION_MEMORY = True
            return
        raise RuntimeError("Cannot enable population memory unless curiosity search was enabled during init.")

    def restorePopulationMemory(self, csGridPop):
        '''Restores the population-level memory from a (previously serialized) list. Added by CMS 22/05/18.'''
        if self.curiositySearch and self.POPULATION_MEMORY:
            self.csGridPop = csGridPop
            return
        raise RuntimeError("Cannot restore population memory unless population-level memory was enabled during init.")

    def getPopulationMemory(self):
        '''Obtains the CS population memory for serialization. Added by CMS 22/05/18.'''
        if self.curiositySearch and self.POPULATION_MEMORY:
            return self.csGridPop
        raise RuntimeError("Cannot obtain population memory unless population-level memory was enabled during init.")

    def _clearCSGrid(self):
        '''Clears the current curiosity grid. Added by CMS 10/01/17.'''
        if (self.curiositySearch):
            for i in range(0, len(self.csGrid)):
                self.csGrid[i] = None
            if self.SLICE_CS_GRID:
                self.slicedCSGrid.fill(self.CS_EXPLORED)
            return
        raise RuntimeError("Cannot clear CS grid unless curiosity search was enabled during init.")

    def _getCSRAM(self):
        '''Obtains a tuple (room, x, y) for the current agent. Added by CSM 10/01/17.'''
        if (self.curiositySearch):
            currRAM = self._get_ram()
            addrX = self.ADDR_X
            addrY = self.ADDR_Y
            if self.ADDR_POSITION_OFFSET is not None:
                addrX += currRAM[self.ADDR_POSITION_OFFSET]
                addrY += currRAM[self.ADDR_POSITION_OFFSET]
            currX = currRAM[addrX] - self.ADDR_X_OFFSET
            currY = currRAM[addrY] - self.ADDR_Y_OFFSET
            if (self.ADDR_ROOM is None):
                currRoom = 0
            else:
                currRoom = currRAM[self.ADDR_ROOM]
            if (self.ADDR_ROOM_2 is not None):
                currRoom = currRoom * self.ROOM_DIM_LENGTH + currRAM[self.ADDR_ROOM_2]
            if self.Y_DOWN:
                currY = self.CS_GRID_HEIGHT - currY - 1
            roomReset = None
            if self.ADDR_ROOM_REFRESH_VALUE is not None:
                roomReset = currRAM[self.ADDR_ROOM_REFRESH_VALUE]
            return currRoom, currX, currY, roomReset
        raise RuntimeError("Cannot obtain CS RAM values unless curiosity search was enabled during init.")

    def _getIntrinsic(self):
        '''Obtains intrinsic reward for the agent. Always call before getCurrCSGrid()! Added by CMS 10/01/17.'''
        if (self.curiositySearch):
            currRoom, currX, currY, _ = self._getCSRAM()
            if self.roomOverride != -1:
                currRoom = self.roomOverride

            # If agent is off-screen, no rewards of any kind.
            self.currIntrinsic = 0
            self.currIntrinsicPop = 0
            if (currX >= self.CS_GRID_WIDTH or currY >= self.CS_GRID_HEIGHT) and self.ale.lives() == 0:
                return # some games (e.g. Amidar) place the player off-screen on end-of-game

       	    # Get individual and population-level intrinsic reward. The latter may or may not be used.
       	    self.currIntrinsicPop = 0
       	    if ((self.csGridPop[currRoom] is None) or (self.csGridPop[currRoom][currX][currY] != self.CS_EXPLORED)):
       	       	self.currIntrinsicPop =	self.INTRINSIC_REWARD
            if ((self.csGrid[currRoom] is None) or (self.csGrid[currRoom][currX][currY] != self.CS_EXPLORED)):
                self.currIntrinsic = self.INTRINSIC_REWARD
            return
        raise RuntimeError("Cannot obtain intrinsic reward unless curiosity search was enabled during init.")

    def getIntrinsic(self):
        '''Obtains the current intrinsic reward for the agent. Added by CMS 10/01/17.'''
        if (self.curiositySearch):
            return self.currIntrinsic
        raise RuntimeError("Cannot obtain intrinsic reward unless curiosity search was enabled during init.")

    def getIntrinsicPop(self):
        '''Obtains the current population intrinsic reward for the agent. Added by CMS 22/05/18.'''
        if (self.curiositySearch and self.POPULATION_MEMORY):
            return self.currIntrinsicPop
        raise RuntimeError("Cannot obtain population-level intrinsic reward unless curiosity search / population memory was enabled.")

    def _getCurrCSGrid(self):
        '''Gets the current curiosity grid. Added by CMS 10/01/17.'''
        if (self.curiositySearch):
            # Get copy of current RAM state.
            currRoom, currX, currY, _ = self._getCSRAM()
            if self.roomOverride != -1:
                currRoom = self.roomOverride

            # Initialize CS grid for this room if necessary.
            if (self.csGrid[currRoom] is None):
                self.csGrid[currRoom] = np.zeros(shape=(self.CS_GRID_WIDTH, self.CS_GRID_HEIGHT, 1), dtype=np.uint8)
            if (self.csGridPop[currRoom] is None):
                self.csGridPop[currRoom] = np.zeros(shape=(self.CS_GRID_WIDTH, self.CS_GRID_HEIGHT, 1), dtype=np.uint8)

            # Some games (e.g. Amidar) place the player off-screen on end-of-game; return grid without changes.
            if (currX >= self.CS_GRID_WIDTH or currY >= self.CS_GRID_HEIGHT) and self.ale.lives() == 0:
                self.currGrid = self.csGrid[currRoom]
                self.currGridPop = self.csGridPop[currRoom]
                return

            # Update CS grid with current position.
            startX = (currX // self.CS_GRID_SIZE) * self.CS_GRID_SIZE # the left edge of the current tile
            startY = (currY // self.CS_GRID_SIZE) * self.CS_GRID_SIZE # the bottom edge of the current tile
            endX = min(self.CS_GRID_WIDTH, startX + self.CS_GRID_SIZE) # the right edge of the tile or screen
            endY = min(self.CS_GRID_HEIGHT, startY + self.CS_GRID_SIZE) # the top edge of the tile or screen
            for x in range(startX, endX):
                for y in range(startY, endY):
                    self.csGrid[currRoom][x][y] = self.CS_EXPLORED # mark each "pixel" in tile as explored
                    self.csGridPop[currRoom][x][y] = self.CS_EXPLORED
            self.currGrid = self.csGrid[currRoom]
            self.currGridPop = self.csGridPop[currRoom]

            # Slice the CS grid for partial input.
            if self.SLICE_CS_GRID:
                self.slicedCSGrid.fill(self.CS_EXPLORED)
                self.slicedCSGridPop.fill(self.CS_EXPLORED)
                startY = currY - (self.CS_SLICE_HEIGHT // 2)
                endY = currY + (self.CS_SLICE_HEIGHT // 2)
                tempOffsetY = 0
                if startY < 0:
                    tempOffsetY = abs(startY)
                tempMaxY = self.CS_SLICE_HEIGHT
                if endY >= self.CS_GRID_HEIGHT:
                    tempMaxY -= (endY - self.CS_GRID_HEIGHT + 1)
                self.slicedCSGrid[:, tempOffsetY:tempMaxY] = self.currGrid[:, max(0, startY):min(self.CS_GRID_HEIGHT, endY + 1)]
                self.slicedCSGridPop[:, tempOffsetY:tempMaxY] = self.currGridPop[:, max(0, startY):min(self.CS_GRID_HEIGHT, endY + 1)]
                #for x in range(0, self.CS_GRID_WIDTH):
                #    currCounter = startY
                #    for y in range(0, self.CS_SLICE_HEIGHT):
                #        if currCounter >= 0 and currCounter < self.CS_GRID_HEIGHT: 
                #        	self.slicedCSGrid[x][y] = self.currGrid[x][currCounter]
                #        currCounter += 1
                self.currGrid = self.slicedCSGrid
                self.currGridPop = self.slicedCSGridPop
            return
        raise RuntimeError("Cannot access CS grid unless curiosity search was enabled during init.")

    def getCurrCSGrid(self):
        '''Gets the current curiosity grid. Added by CMS 10/01/17.'''
        if (self.curiositySearch):
            return self.currGrid
        raise RuntimeError("Cannot obtain CS grid unless curiosity search was enabled during init.")

    def getCurrCSGridPop(self):
        '''Gets the current population-level curiosity grid. Added by CMS 22/05/18.'''
        if (self.curiositySearch and self.POPULATION_MEMORY):
            return self.currGridPop
        raise RuntimeError("Cannot obtain population-level CS grid unless population memory was enabled during init.")

    def _seed(self, seed=None):
        self.np_random, seed1 = seeding.np_random(seed)
        # Derive a random seed. This gets passed as a uint, but gets
        # checked as an int elsewhere, so we need to keep it below
        # 2**31.
        seed2 = seeding.hash_seed(seed1 + 1) % 2**31
        # Empirically, we need to seed before loading the ROM.
        self.ale.setInt(b'random_seed', seed2)
        self.ale.loadROM(self.game_path)
        return [seed1, seed2]

    def _step(self, a):
        reward = 0.0
        action = self._action_set[a]

        if isinstance(self.frameskip, int):
            num_steps = self.frameskip
        else:
            num_steps = self.np_random.randint(self.frameskip[0], self.frameskip[1])
        for _ in range(num_steps):
            reward += self.ale.act(action)
        ob = self._get_obs()

        # BEGIN CURIOSITY SEARCH (ADDED BY CMS 10/01/17)
        if (self.curiositySearch):
            currRoom, _, _, roomReset = self._getCSRAM()
            if self.CS_ROOM_SWITCH and self.lastCurrRoom != -1 and self.lastCurrRoom != currRoom:
                #self._clearCSGrid()
                self.roomOverride = max(self.roomOverride, currRoom) + 1
                self.totalSwitchedRooms += 1
                self.lastCurrRoom = -1
            elif self.CS_ROOM_SWITCH:
                self.lastCurrRoom = currRoom
            if self.ADDR_ROOM_REFRESH_VALUE is not None and roomReset == self.ROOM_REFRESH_VALUE and not self.resetRoomForLevel:
                #self._clearCSGrid()
                self.roomOverride = max(self.roomOverride, currRoom) + 1
                self.totalSwitchedRooms += 1
                self.resetRoomForLevel = True
            if self.ADDR_ROOM_REFRESH_VALUE is not None and roomReset != self.ROOM_REFRESH_VALUE and self.resetRoomForLevel:
                self.resetRoomForLevel = False # prevents bugs with frame repeating on level start
            self._getIntrinsic() # updates current intrinsic reward value
            self._getCurrCSGrid() # updates current grid observation
            return ob, reward, self.ale.game_over(), {"ale.lives": self.ale.lives(), "intrinsic": self.currIntrinsic, "totalRooms": self._totalRooms(), "intrinsicPop": self.currIntrinsicPop}
        # END CURIOSITY SEARCH

        return ob, reward, self.ale.game_over(), {"ale.lives": self.ale.lives()}

    def _totalRooms(self):
        '''Returns the total number of explored rooms, if applicable to this game. Added by CMS 15/12/17.'''
        if self.CS_ROOM_SWITCH or self.ADDR_ROOM_REFRESH_VALUE is not None:
            return self.totalSwitchedRooms + 1
        toReturn = 0
        for i in range(0, len(self.csGrid)):
            if (self.csGrid[i] is not None):
                toReturn +=1
        return toReturn

    def _get_image(self):
        return self.ale.getScreenRGB2()

    def _get_ram(self):
        return to_ram(self.ale)

    @property
    def _n_actions(self):
        return len(self._action_set)

    def _get_obs(self):
        if self._obs_type == 'ram':
            return self._get_ram()
        elif self._obs_type == 'image':
            img = self._get_image()
        return img

    # return: (states, observations)
    def _reset(self):
        self.ale.reset_game()
        ### BEGIN CURIOSITY SEARCH (ADDED BY CMS 10/01/17)
        if (self.curiositySearch):
            self._clearCSGrid()
            self.currIntrinsic = 0
            self.currGrid = None
            self.lastCurrRoom = -1
            self.totalSwitchedRooms = 0
            self.resetRoomForLevel = False
            self.roomOverride = -1
        ### END CURIOSITY SEARCH
        return self._get_obs()

    def _render(self, mode='human', close=False):
        if close:
            if self.viewer is not None:
                self.viewer.close()
                self.viewer = None
            return
        img = self._get_image()
        if mode == 'rgb_array':
            return img
        elif mode == 'human':
            from gym.envs.classic_control import rendering
            if self.viewer is None:
                self.viewer = rendering.SimpleImageViewer()
            self.viewer.imshow(img)

    def get_action_meanings(self):
        return [ACTION_MEANING[i] for i in self._action_set]

    def get_keys_to_action(self):
        KEYWORD_TO_KEY = {
            'UP':      ord('w'),
            'DOWN':    ord('s'),
            'LEFT':    ord('a'),
            'RIGHT':   ord('d'),
            'FIRE':    ord(' '),
        }

        keys_to_action = {}

        for action_id, action_meaning in enumerate(self.get_action_meanings()):
            keys = []
            for keyword, key in KEYWORD_TO_KEY.items():
                if keyword in action_meaning:
                    keys.append(key)
            keys = tuple(sorted(keys))

            assert keys not in keys_to_action
            keys_to_action[keys] = action_id

        return keys_to_action

    def clone_state(self):
        """Clone emulator state w/o system state. Restoring this state will
        *not* give an identical environment. For complete cloning and restoring
        of the full state, see `{clone,restore}_full_state()`."""
        state_ref = self.ale.cloneState()
        state = self.ale.encodeState(state_ref)
        self.ale.deleteState(state_ref)
        return state

    def restore_state(self, state):
        """Restore emulator state w/o system state."""
        state_ref = self.ale.decodeState(state)
        self.ale.restoreState(state_ref)
        self.ale.deleteState(state_ref)

    def clone_full_state(self):
        """Clone emulator state w/ system state including pseudorandomness.
        Restoring this state will give an identical environment."""
        state_ref = self.ale.cloneSystemState()
        state = self.ale.encodeState(state_ref)
        self.ale.deleteState(state_ref)
        return state

    def restore_full_state(self, state):
        """Restore emulator state w/ system state including pseudorandomness."""
        state_ref = self.ale.decodeState(state)
        self.ale.restoreSystemState(state_ref)
        self.ale.deleteState(state_ref)

ACTION_MEANING = {
    0 : "NOOP",
    1 : "FIRE",
    2 : "UP",
    3 : "RIGHT",
    4 : "LEFT",
    5 : "DOWN",
    6 : "UPRIGHT",
    7 : "UPLEFT",
    8 : "DOWNRIGHT",
    9 : "DOWNLEFT",
    10 : "UPFIRE",
    11 : "RIGHTFIRE",
    12 : "LEFTFIRE",
    13 : "DOWNFIRE",
    14 : "UPRIGHTFIRE",
    15 : "UPLEFTFIRE",
    16 : "DOWNRIGHTFIRE",
    17 : "DOWNLEFTFIRE",
}
