#!/usr/bin/env python
# python_example.py
# Author: Ben Goodrich
#
# This is a direct port to python of the shared library example from
# ALE provided in doc/examples/sharedLibraryInterfaceExample.cpp
import sys
import numpy as np
from random import randrange
from ale_python_interface import ALEInterface


def squareLoop(I,J):
  for i in (I, I+4):
    for j in (J, J+4):
      if(trimmed_data2[i,j] != 87):
        pooling_data[I/4,J/4] = 1
        return

if len(sys.argv) < 2:
  print 'Usage:', sys.argv[0], 'rom_file'
  sys.exit()

ale = ALEInterface()

# Get & Set the desired settings
ale.setInt('random_seed', 123)

# Set USE_SDL to true to display the screen. ALE must be compilied
# with SDL enabled for this to work. On OSX, pygame init is used to
# proxy-call SDL_main.
USE_SDL = True
if USE_SDL:
  if sys.platform == 'darwin':
    import pygame
    pygame.init()
    ale.setBool('sound', False) # Sound doesn't work on OSX
  elif sys.platform.startswith('linux'):
    ale.setBool('sound', True)
  ale.setBool('display_screen', True)

# Load the ROM file
ale.loadROM(sys.argv[1])
# Get the list of legal actions
legal_actions = ale.getLegalActionSet()
#--------------------------------------------------------
(screen_width, screen_height) = ale.getScreenDims()
#screen_data = np.zeros(screen_width*screen_height, dtype=np.uint32)
screen_data = np.zeros((screen_height,screen_width), dtype=np.uint8)
pooling_data = np.zeros((40,40), dtype=np.uint8)
#--------------------------------------------------------
# Play 10 episodes
for episode in xrange(1):
  total_reward = 0
  #while not ale.game_over():
  for i in xrange(200):




    a = legal_actions[randrange(len(legal_actions))]
    # Apply an action and get the resulting reward
    reward = ale.act(a);
    total_reward += reward

  ale.getScreenGrayscale(screen_data)
  trimmed_data = np.delete(screen_data, np.s_[195::], 0)
  trimmed_data2 = np.delete(trimmed_data, np.s_[0:35], 0)
  for I in range(0,156,4):
    for J in range(0,156,4):
      squareLoop(I,J)



  np.savetxt('dat1.txt', trimmed_data2, fmt='%2.2f')
  np.savetxt('dat1pooled.txt', pooling_data, fmt='%2.2f')
  print 'Episode', episode, 'ended with score:', total_reward
  print screen_data.shape
  print trimmed_data.shape
  print trimmed_data2.shape
  print pooling_data.shape
  ale.reset_game()