import streamlit as st
import numpy as np
import matplotlib.pyplot as plt
import plotly.graph_objects as go
from PIL import Image
import io
import base64
import pandas as pd

# ======================
# 🎨 STYLING & LAYOUT
# ======================
st.set_page_config(
    page_title="RL Training Studio",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS
st.markdown("""
<style>
    /* Main container */
    .main {
        background: linear-gradient(135deg, #0f2027 0%, #203a43 50%, #2c5364 100%);
        color: white;
    }
    
    /* Headers */
    h1, h2, h3 {
        color: #00d4ff;
        font-family: 'Arial', sans-serif;
    }
    
    /* Sidebar */
    [data-testid="stSidebar"] {
        background: linear-gradient(195deg, #000000 0%, #0f2027 100%) !important;
        border-right: 1px solid #00d4ff33;
    }
    
    /* Buttons */
    .stButton>button {
        background: linear-gradient(to right, #00d4ff, #0083b0);
        color: white;
        border: none;
        border-radius: 8px;
        padding: 10px 24px;
        font-weight: bold;
        transition: all 0.3s ease;
    }
    .stButton>button:hover {
        transform: scale(1.05);
        box-shadow: 0 0 15px #00d4ff80;
    }
    
    /* Tabs */
    .stTabs>div>div>button {
        color: white !important;
        font-weight: bold;
    }
    .stTabs>div>div>button[aria-selected="true"] {
        border-bottom: 2px solid #00d4ff !important;
    }
    
    /* Cards */
    .card {
        background: rgba(0, 212, 255, 0.1) !important;
        border-radius: 10px;
        padding: 20px;
        border: 1px solid #00d4ff33;
        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
        transition: all 0.3s ease;
    }
    .card:hover {
        transform: translateY(-5px);
        box-shadow: 0 10px 20px rgba(0, 212, 255, 0.2);
    }
    
    /* Sliders */
    .stSlider>div>div>div>div {
        background: #00d4ff !important;
    }
    
    /* Metrics */
    [data-testid="metric-container"] {
        background: rgba(0, 212, 255, 0.15) !important;
        border-radius: 10px;
        padding: 15px !important;
        border: 1px solid #00d4ff33;
    }
</style>
""", unsafe_allow_html=True)

# ======================
# 🧠 SIMULATION DATA
# ======================
@st.cache_data
def generate_training_data():
    # Generate sample training curves
    episodes = np.arange(1, 101)
    reward = np.cumsum(np.random.normal(5, 2, 100))
    loss = np.abs(np.random.normal(0, 0.5, 100).cumsum())
    epsilon = np.concatenate([np.linspace(1, 0.1, 90), np.ones(10)*0.1])
    
    return pd.DataFrame({
        'Episode': episodes,
        'Reward': reward,
        'Loss': loss,
        'Epsilon': epsilon
    })

df = generate_training_data()

# ======================
# 🎛️ SIDEBAR CONTROLS
# ======================
with st.sidebar:
    st.title("⚙️ RL Controls")
    st.markdown("---")
    
    st.subheader("Environment")
    env_type = st.selectbox(
        "Select Environment",
        ["OpenAI Gym", "Unity ML-Agents", "Custom Robotics"],
        index=0
    )
    
    st.subheader("Algorithm")
    algorithm = st.selectbox(
        "RL Algorithm",
        ["PPO", "DQN", "A2C", "SAC", "TD3"],
        index=0
    )
    
    st.subheader("Hyperparameters")
    learning_rate = st.slider("Learning Rate", 1e-5, 1e-2, 3e-4, format="%.0e")
    gamma = st.slider("Discount Factor (γ)", 0.8, 0.999, 0.99, 0.001)
    batch_size = st.slider("Batch Size", 32, 1024, 128, 32)
    
    if st.button("🚀 Start Training", use_container_width=True):
        st.toast("Training session started!", icon="🤖")
    
    st.markdown("---")
    st.subheader("📊 Training Metrics")
    st.metric("Current Episode", "42", "+3")
    st.metric("Mean Reward", "187.5", "12.3%")
    st.metric("Exploration Rate", "23%", "-5%")

# ======================
# 🚀 MAIN INTERFACE
# ======================
st.title("🤖 Reinforcement Learning Studio")
st.caption("Train AI agents for games and robotics with state-of-the-art RL algorithms")

# ======================
# 📊 DASHBOARD
# ======================
tab1, tab2, tab3 = st.tabs(["📈 Training Progress", "🎮 Simulation", "🧠 Model Details"])

with tab1:
    col1, col2 = st.columns(2)
    
    with col1:
        st.markdown("### Reward Curve")
        fig = go.Figure()
        fig.add_trace(go.Scatter(
            x=df['Episode'], 
            y=df['Reward'],
            line=dict(color='#00d4ff', width=3),
            name="Episode Reward"
        ))
        fig.update_layout(
            template="plotly_dark",
            paper_bgcolor='rgba(0,0,0,0)',
            plot_bgcolor='rgba(0,0,0,0)',
            margin=dict(l=20, r=20, t=30, b=20),
            height=300
        )
        st.plotly_chart(fig, use_container_width=True)
        
        st.markdown("### Loss Curve")
        fig = go.Figure()
        fig.add_trace(go.Scatter(
            x=df['Episode'], 
            y=df['Loss'],
            line=dict(color='#ff4d4d', width=3),
            name="Training Loss"
        ))
        fig.update_layout(
            template="plotly_dark",
            paper_bgcolor='rgba(0,0,0,0)',
            plot_bgcolor='rgba(0,0,0,0)',
            margin=dict(l=20, r=20, t=30, b=20),
            height=300
        )
        st.plotly_chart(fig, use_container_width=True)
    
    with col2:
        st.markdown("### Exploration Rate")
        fig = go.Figure()
        fig.add_trace(go.Scatter(
            x=df['Episode'], 
            y=df['Epsilon'],
            line=dict(color='#f39c12', width=3),
            name="Epsilon"
        ))
        fig.update_layout(
            template="plotly_dark",
            paper_bgcolor='rgba(0,0,0,0)',
            plot_bgcolor='rgba(0,0,0,0)',
            margin=dict(l=20, r=20, t=30, b=20),
            height=300
        )
        st.plotly_chart(fig, use_container_width=True)
        
        st.markdown("### Action Distribution")
        actions = np.random.dirichlet(np.ones(5), size=1)[0]
        fig = go.Figure()
        fig.add_trace(go.Bar(
            x=["Left", "Right", "Up", "Down", "Stay"],
            y=actions,
            marker_color=['#00d4ff', '#f39c12', '#2ecc71', '#e74c3c', '#9b59b6']
        ))
        fig.update_layout(
            template="plotly_dark",
            paper_bgcolor='rgba(0,0,0,0)',
            plot_bgcolor='rgba(0,0,0,0)',
            margin=dict(l=20, r=20, t=30, b=20),
            height=300
        )
        st.plotly_chart(fig, use_container_width=True)

with tab2:
    st.markdown("### 🎮 Live Agent Visualization")
    
    # Placeholder for simulation
    col1, col2 = st.columns([3, 1])
    
    with col1:
        # Create a sample "simulation" frame
        fig, ax = plt.subplots(figsize=(8, 6))
        ax.set_facecolor('#2c3e50')
        ax.grid(color='#34495e', linestyle='--', linewidth=0.5)
        
        # Draw a simple robot/agent
        agent = plt.Circle((0.5, 0.5), 0.1, color='#00d4ff')
        ax.add_patch(agent)
        
        # Draw some obstacles
        for _ in range(5):
            x, y = np.random.rand(2) * 0.8 + 0.1
            obs = plt.Rectangle((x, y), 0.1, 0.1, color='#e74c3c')
            ax.add_patch(obs)
        
        # Draw a goal
        goal = plt.Rectangle((0.8, 0.8), 0.1, 0.1, color='#2ecc71')
        ax.add_patch(goal)
        
        ax.set_xlim(0, 1)
        ax.set_ylim(0, 1)
        ax.set_xticks([])
        ax.set_yticks([])
        ax.set_title("Agent Simulation", color='white', pad=20)
        
        # Convert to image
        buf = io.BytesIO()
        plt.savefig(buf, format='png', dpi=100, facecolor='#2c3e50', bbox_inches='tight')
        buf.seek(0)
        img = Image.open(buf)
        st.image(img, use_column_width=True)
        
    with col2:
        st.markdown("### Agent State")
        st.json({
            "position": {"x": 0.45, "y": 0.52},
            "velocity": {"x": 0.02, "y": -0.01},
            "sensors": [0.87, 0.45, 0.23, 0.12],
            "last_action": "Move Right",
            "current_reward": 2.45
        })
        
        if st.button("▶️ Run Episode", use_container_width=True):
            st.toast("Running simulation episode...", icon="🎮")

with tab3:
    st.markdown("### 🧠 Model Architecture")
    
    col1, col2 = st.columns(2)
    
    with col1:
        st.markdown("#### Neural Network Layers")
        st.table(pd.DataFrame({
            "Layer": ["Input", "Dense", "LSTM", "Dense", "Output"],
            "Units": ["State Size (24)", "256", "128", "64", "Action Size (5)"],
            "Activation": ["-", "ReLU", "Tanh", "ReLU", "Linear"]
        }))
        
        st.markdown("#### Training Configuration")
        st.table(pd.DataFrame({
            "Parameter": ["Algorithm", "Optimizer", "Learning Rate", "Gamma", "Batch Size"],
            "Value": [algorithm, "Adam", f"{learning_rate:.1e}", gamma, batch_size]
        }))
    
    with col2:
        st.markdown("#### Policy Network Diagram")
        # Simple neural network visualization
        fig = go.Figure()
        
        # Add nodes
        fig.add_trace(go.Scatter(
            x=[1, 2, 2, 3, 3, 4],
            y=[0, -1, 1, -1.5, -0.5, 0.5, 1.5, 0],
            mode="markers",
            marker=dict(size=20, color='#00d4ff'),
            name="Neurons"
        ))
        
        # Add connections
        for x in [1]:
            for y in [2, 2]:
                fig.add_trace(go.Scatter(
                    x=[x, y],
                    y=[0, -1],
                    mode="lines",
                    line=dict(width=1, color='rgba(0, 212, 255, 0.3)'),
                    hoverinfo='none'
                ))
        
        fig.update_layout(
            template="plotly_dark",
            paper_bgcolor='rgba(0,0,0,0)',
            plot_bgcolor='rgba(0,0,0,0)',
            showlegend=False,
            xaxis=dict(showgrid=False, zeroline=False, visible=False),
            yaxis=dict(showgrid=False, zeroline=False, visible=False),
            height=400,
            margin=dict(l=20, r=20, t=20, b=20)
        )
        st.plotly_chart(fig, use_container_width=True)

# ======================
# 📝 FOOTER
# ======================
st.markdown("---")
st.markdown("""
<div style="text-align: center;">
    <p>🤖 Training <strong>{}</strong> in <strong>{}</strong> environment | ⚡ Algorithm: <strong>{}</strong></p>
</div>
""".format(algorithm, env_type, algorithm), unsafe_allow_html=True)