import pickle
import sys
import numpy as np
import pygame
from ale_py import ALEInterface

# Initialize the ALE interface
ale = ALEInterface()

# Load the ROM (game) file
rom_file = "montezuma_revenge.bin"  # Replace with the path to your Space Invaders ROM
ale.loadROM(rom_file)

print(ale.getLegalActionSet())

room_10 = pickle.load(open("room_10.pkl", "rb"))
ale.restoreState(room_10)
# Set up Pygame for keyboard input
pygame.init()

# Mapping of actions to keyboard keys (you can modify this based on your preference)
# Mapping of ALE actions to keyboard keys (WSAD controls)
key_mapping = {
    pygame.K_w: 2,           # Action.UP
    pygame.K_d: 3,           # Action.RIGHT
    pygame.K_a: 4,           # Action.LEFT
    pygame.K_s: 5,           # Action.DOWN
    pygame.K_w + pygame.K_d: 6,   # Action.UPRIGHT
    pygame.K_w + pygame.K_a: 7,   # Action.UPLEFT
    pygame.K_s + pygame.K_d: 8,   # Action.DOWNRIGHT
    pygame.K_s + pygame.K_a: 9,   # Action.DOWNLEFT
    pygame.K_SPACE: 1,       # Action.FIRE
    pygame.K_w + pygame.K_SPACE: 10,  # Action.UPFIRE
    pygame.K_d + pygame.K_SPACE: 11,  # Action.RIGHTFIRE
    pygame.K_a + pygame.K_SPACE: 12,  # Action.LEFTFIRE
    pygame.K_s + pygame.K_SPACE: 13,  # Action.DOWNFIRE
    pygame.K_w + pygame.K_d + pygame.K_SPACE: 14,  # Action.UPRIGHTFIRE
    pygame.K_w + pygame.K_a + pygame.K_SPACE: 15,  # Action.UPLEFTFIRE
    pygame.K_s + pygame.K_d + pygame.K_SPACE: 16,  # Action.DOWNRIGHTFIRE
    pygame.K_s + pygame.K_a + pygame.K_SPACE: 17  # Action.DOWNLEFTFIRE
}


WIDTH=600
HEIGHT=480
screen = pygame.display.set_mode((WIDTH, HEIGHT))


# Main loop
target_fps = 20  # Set your target frame rate here

# Initialize action to None
action = None
running = True
last_key = None

while running:
    for event in pygame.event.get():

        keys = pygame.key.get_pressed()

        if event.type == pygame.QUIT:
            running = False
        elif event.type == pygame.KEYDOWN:
            if event.key in key_mapping:
                action = key_mapping[event.key]
                last_key = event.key
        # If no key is currently pressed, set action to None
        if pygame.key.get_focused() and last_key is not None and not pygame.key.get_pressed()[last_key]:
            action = 0

    if action is not None:
        reward = ale.act(action)
        screen_data = ale.getScreenRGB()

        # Convert screen_data to a Pygame surface
        screen_surface = pygame.surfarray.make_surface(screen_data)

        screen_surface = pygame.transform.scale(screen_surface, (HEIGHT, WIDTH))

        # Rotate the screen by 90 degrees clockwise
        screen_surface = pygame.transform.rotate(screen_surface, -90)

        # Mirror the screen by the Y-axis (flip horizontally)
        screen_surface = pygame.transform.flip(screen_surface, True, False)

        # Blit the screen surface onto the display
        screen.blit(screen_surface, (0, 0))

        # Process the screen_data as needed
        # For example, you can display the game screen using a library like OpenCV
        # print("Action:", action, "Reward:", reward)
        if ale.game_over():
            ale.reset_game()

    # Refresh the screen
    pygame.display.flip()

    # Control the frame rate to achieve 20 FPS
    pygame.time.delay(1000 // target_fps)

# Clean up
ale.reset_game()
pygame.quit()
sys.exit()
