from __future__ import print_function
from __future__ import division

import os
import gym
import numpy as np

from skimage.transform import resize
from skimage.color import rgb2gray


class Atari(object):
    s_dim = [84, 84, 1]
    a_dim = 3

    def __init__(self, args, record_video=False):
        self.env = gym.make('BreakoutNoFrameskip-v4')
        self.ale = self.env.env.ale  # ale interface
        if record_video:
            video_dir = os.path.join(args.save_path, 'videos')
            if not os.path.exists(video_dir):
                os.makedirs(video_dir)
            self.env = gym.wrappers.Monitor(
                self.env, video_dir, video_callable=lambda x: True, resume=True)
            self.ale = self.env.env.env.ale

        self.screen_size = Atari.s_dim[:2]  # 84x84
        self.noop_max = 30
        self.frame_skip = 4
        self.frame_feq = 4
        self.s_dim = Atari.s_dim
        self.a_dim = Atari.a_dim

        self.action_space = [1, 2, 3]  # Breakout specify
        self.done = True

    def new_round(self):
        if not self.done:  # dead but not done
            # no-op step to advance from terminal/lost life state
            obs, _, _, _ = self.env.step(0)
            obs = self.preprocess(obs)
        else:  # terminal
            self.env.reset()
            # No-op
            for _ in range(np.random.randint(1, self.noop_max + 1)):
                obs, _, done, _ = self.env.step(0)
            obs = self.preprocess(obs)
        return obs

    def preprocess(self, observ):
        return resize(rgb2gray(observ), self.screen_size)

    def step(self, action):
        observ, reward, dead = None, 0, False
        for _ in range(self.frame_skip):
            lives_before = self.ale.lives()
            o, r, self.done, _ = self.env.step(self.action_space[action])
            lives_after = self.ale.lives()
            reward += r
            if lives_before > lives_after:
                dead = True
                break
        observ = self.preprocess(o)
        observ = np.reshape(observ, newshape=self.screen_size + [1])
        self.state = np.append(self.state[:, :, 1:], observ, axis=2)

        return self.state, reward, dead, self.done

[(17, 'gym.make', 'gym.make', False, 'import gym\n'), (64, 'numpy.reshape', 'np.reshape', True, 'import numpy as np\n'), (65, 'numpy.append', 'np.append', True, 'import numpy as np\n'), (20, 'os.path.join', 'os.path.join', False, 'import os\n'), (23, 'gym.wrappers.Monitor', 'gym.wrappers.Monitor', False, 'import gym\n'), (51, 'skimage.color.rgb2gray', 'rgb2gray', False, 'from skimage.color import rgb2gray\n'), (21, 'os.path.exists', 'os.path.exists', False, 'import os\n'), (22, 'os.makedirs', 'os.makedirs', False, 'import os\n'), (45, 'numpy.random.randint', 'np.random.randint', True, 'import numpy as np\n')]