import streamlit as st
import gym
import gym_super_mario_bros
import numpy as np
import cv2
from PIL import Image

# Initialize the environment
env = gym.make('SuperMarioBros-v0')
env.reset()

# Create a Streamlit slider for controlling Mario's actions
st.title("Super Mario Bros. Game!")
st.write("Use the arrow keys to control Mario!")
action = st.slider("Move Mario", 0, 3, 1, step=1)

# Function to convert the game frames to images
def convert_frame_to_image(frame):
    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)  # Convert BGR to RGB
    return Image.fromarray(frame)

# Run the game loop and show the game state
if st.button("Start Game"):
    st.write("Game is running. Use the slider to control Mario.")

    while True:
        # Step the game forward using the action chosen by the slider
        action_taken = np.array([action])
        state, reward, done, info = env.step(action_taken)
        
        # Show the game frame
        frame = convert_frame_to_image(state)
        st.image(frame, use_column_width=True)

        # If the game is over, reset
        if done:
            st.write("Game Over! Resetting...")
            env.reset()