#   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import cv2
from atari import AtariPlayer
from atari_wrapper import FrameStack, MapState, FireResetEnv


def get_player(rom,
               image_size,
               viz=False,
               train=False,
               frame_skip=1,
               context_len=1):
    env = AtariPlayer(
        rom,
        frame_skip=frame_skip,
        viz=viz,
        live_lost_as_eoe=train,
        max_num_frames=60000)
    env = FireResetEnv(env)
    env = MapState(env, lambda im: cv2.resize(im, image_size))
    if not train:
        # in training, context is taken care of in expreplay buffer
        env = FrameStack(env, context_len)
    return env