app.py

import streamlit as st
import time
import os
from web_game_rl import SnakeGameRL
import numpy as np
from PIL import Image

# Set page configuration
st.set_page_config(
    page_title="Snake AI Reinforcement Learning",
    page_icon="🐍",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for styling
st.markdown("""
<style>
    .main-header {
        font-size: 3rem;
        color: #4CAF50;
        text-align: center;
        margin-bottom: 2rem;
        font-weight: bold;
    }
    .sub-header {
        font-size: 1.5rem;
        color: #2E7D32;
        margin-bottom: 1rem;
    }
    .info-text {
        font-size: 1.1rem;
        line-height: 1.6;
    }
    .stButton>button {
        background-color: #4CAF50;
        color: white;
        font-size: 1.2rem;
        padding: 0.5rem 2rem;
        border-radius: 8px;
        border: none;
        transition: all 0.3s;
    }
    .stButton>button:hover {
        background-color: #2E7D32;
        transform: scale(1.05);
    }
</style>
""", unsafe_allow_html=True)

# Initialize session state variables
if 'initialized' not in st.session_state:
    st.session_state['initialized'] = True
    st.session_state['game'] = None
    st.session_state['game_active'] = False
    st.session_state['score'] = 0
    st.session_state['game_over'] = False
    st.session_state['auto_refresh'] = True
    st.session_state['games_played'] = 0
    st.session_state['record'] = 0

# Main app layout
st.markdown("<h1 class='main-header'>🐍 Snake AI Reinforcement Learning</h1>", unsafe_allow_html=True)

# Create a two-column layout
left_col, right_col = st.columns([2, 3])

with left_col:
    # Control section
    st.markdown("<h2 class='sub-header'>Snake Game Controls</h2>", unsafe_allow_html=True)
    
    # Start/Stop button
    if not st.session_state['game_active']:
        if st.button("Start Snake AI"):
            # Initialize the game with reinforcement learning
            st.session_state['game'] = SnakeGameRL(480, 360)
            st.session_state['game_active'] = True
            st.session_state['score'] = 0
            st.session_state['game_over'] = False
    else:
        if st.button("Stop Snake AI"):
            st.session_state['game_active'] = False
    
    # Reset button
    if st.button("Reset Game"):
        if st.session_state['game'] is not None:
            st.session_state['game'].reset()
            st.session_state['score'] = 0
            st.session_state['game_over'] = False
    
    # Train button
    if st.button("Train Network"):
        if st.session_state['game'] is not None:
            st.session_state['game'].train_long_memory()
            st.info("Training completed on current memory batch")
    
    # Auto-refresh toggle
    auto_refresh = st.checkbox("Auto-refresh Game", value=st.session_state['auto_refresh'])
    st.session_state['auto_refresh'] = auto_refresh
    
    # Speed slider
    speed = st.slider("Game Speed", min_value=10, max_value=60, value=20, step=5)
    if st.session_state['game'] is not None:
        # This won't actually change the speed in the current game instance
        # But we'll use it as a reference for the next game
        st.session_state['game_speed'] = speed
    else:
        st.session_state['game_speed'] = speed
    
    # Game metrics
    st.markdown("<h2 class='sub-header'>Game Stats</h2>", unsafe_allow_html=True)
    
    col1, col2 = st.columns(2)
    with col1:
        st.metric("Current Score", st.session_state['score'])
        st.metric("Games Played", st.session_state['games_played'])
    with col2:
        status = "Game Over" if st.session_state['game_over'] else "Playing" if st.session_state['game_active'] else "Not Started"
        st.metric("Game Status", status)
        st.metric("Record Score", st.session_state['record'])
    
    # Project information
    st.markdown("<h2 class='sub-header'>Project Information</h2>", unsafe_allow_html=True)
    
    st.markdown("""
    <p class='info-text'>
    This Snake AI project uses reinforcement learning to train an agent to play the Snake game.
    The agent learns through experience, improving its performance over time.
    </p>
    
    <p class='info-text'>
    <b>How Reinforcement Learning Works:</b>
    <ul>
    <li><b>State:</b> The AI sees the game state (danger directions, food location, etc.)</li>
    <li><b>Action:</b> Based on the state, the AI chooses an action (move straight, turn left/right)</li>
    <li><b>Reward:</b> The AI gets rewards for eating food (+10) and penalties for crashing (-10)</li>
    <li><b>Learning:</b> The AI updates its neural network based on these experiences</li>
    <li><b>Memory:</b> The AI remembers past experiences and learns from them in batches</li>
    </ul>
    </p>
    """, unsafe_allow_html=True)

with right_col:
    # Game visualization
    st.markdown("<h2 class='sub-header'>Snake Game</h2>", unsafe_allow_html=True)
    
    # Create a placeholder for the game
    game_placeholder = st.empty()
    
    # Update game state and display
    if st.session_state['game_active']:
        # Play one step of the game
        game_over, score = st.session_state['game'].play_step()
        
        # Update session state
        st.session_state['score'] = score
        st.session_state['game_over'] = game_over
        
        # Update games played and record
        st.session_state['games_played'] = st.session_state['game'].agent.n_games
        st.session_state['record'] = st.session_state['game'].record
        
        # Get the current frame and display it
        frame = st.session_state['game'].get_frame()
        game_placeholder.image(frame, caption="Snake AI Game", use_column_width=False, width=480)
    elif st.session_state['game'] is not None:
        # Display the last frame if game is paused
        frame = st.session_state['game'].get_frame()
        game_placeholder.image(frame, caption="Snake AI Game", use_column_width=False, width=480)
    else:
        # Display a placeholder if game hasn't started
        game_placeholder.info("Click 'Start Snake AI' to begin playing!")

# Auto-refresh mechanism
if st.session_state['auto_refresh'] and st.session_state['game_active']:
    time.sleep(0.05)  # Reduced wait time for faster refresh
    st.rerun()