# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import gym
import cv2
import numpy as np
import matplotlib.pyplot as plt

env = gym.make('Pong-v0')
obs = env.reset()

obs.shape

plt.imshow(obs)

# +
obs_new = cv2.resize(obs, (84,110), interpolation=cv2.INTER_AREA)
obs_new = cv2.cvtColor(obs_new, cv2.COLOR_RGB2GRAY)
obs_new = obs_new[110-84-8:110-8, :]

obs_new.shape
# -

plt.imshow(obs_new, cmap='gray')


class AtariStateProcessing():
    def __init__(self, input_shape=(84,84), frames=4):
        self.state = [np.zeros(input_shape) for _ in range(frames)]
    

    def preprocess(self, obs, play_area_shift=8):
        obs_new = cv2.resize(obs, (84,110), interpolation=cv2.INTER_AREA)
        obs_new = cv2.cvtColor(obs_new, cv2.COLOR_RGB2GRAY)
        obs_processed = obs_new[110-84-play_area_shift:110-play_area_shift, :]

        return obs_processed
    

    def get_state(self, obs):
        obs_processed = self.preprocess(obs)
        self.state.pop(0)
        self.state.append(obs_processed)

        return self.state.copy()



StarcoderdataJupyterScriptsDedupFiltered