# app.py — Enhanced UI with better layout, visual hierarchy, and UX
# ... [All your imports and backend code remain the same] ...

# =========================
# ENHANCED CSS
# =========================
ENHANCED_CSS = """
:root {
    --radius: 14px;
    --fg: #E5E7EB;
    --muted: #94A3B8;
    --line: #1f2b36;
    --bg: #0F1A24;
    --panel-bg: #0c161f;
    --accent: #3B82F6;
    --accent-hover: #2563EB;
    --danger: #EF4444;
}

.gradio-container {
    max-width: 1400px !important;
    background: var(--bg);
    padding: 16px !important;
}

#header {
    padding: 16px 0;
    margin-bottom: 16px;
    border-bottom: 1px solid var(--line);
}

h1, h2, h3, .gr-markdown {
    color: var(--fg);
}

.gr-button {
    border-radius: 8px;
    padding: 8px 16px;
    transition: all 0.2s ease;
    font-weight: 500 !important;
}

.btn-primary {
    background: var(--accent) !important;
    border: 1px solid var(--accent) !important;
}

.btn-primary:hover {
    background: var(--accent-hover) !important;
}

.btn-secondary {
    background: transparent !important;
    border: 1px solid var(--line) !important;
}

.btn-danger {
    background: var(--danger) !important;
    border: 1px solid var(--danger) !important;
}

.control-group {
    border: 1px solid var(--line);
    border-radius: var(--radius);
    background: var(--panel-bg);
    padding: 20px;
    margin-bottom: 20px;
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.05);
}

.panel {
    border: 1px solid var(--line);
    border-radius: var(--radius);
    background: var(--panel-bg);
    padding: 20px;
    margin-bottom: 20px;
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.05);
}

.stats-panel {
    background: linear-gradient(145deg, #0a121b, #0c161f);
    border-left: 3px solid var(--accent);
}

#stats {
    color: var(--fg);
    line-height: 1.6;
}

#stats strong {
    font-weight: 500;
    color: var(--accent);
}

.param-slider {
    margin-bottom: 12px;
}

.visualization-container {
    display: flex;
    flex-direction: column;
    gap: 20px;
    height: 100%;
}

.viz-panel {
    flex: 1;
    min-height: 300px;
}

.viz-header {
    display: flex;
    justify-content: space-between;
    align-items: center;
    margin-bottom: 12px;
    padding-bottom: 8px;
    border-bottom: 1px solid var(--line);
}

.viz-title {
    font-size: 1.1rem;
    font-weight: 500;
    color: var(--accent);
}

.gen-counter {
    font-size: 0.9rem;
    background: rgba(59, 130, 246, 0.15);
    padding: 4px 10px;
    border-radius: 12px;
}

.slider-info {
    display: flex;
    justify-content: space-between;
    font-size: 0.85rem;
    color: var(--muted);
    margin-top: 4px;
}

.controls-grid {
    display: grid;
    grid-template-columns: 1fr 1fr;
    gap: 16px;
}

@media (max-width: 1200px) {
    .controls-grid {
        grid-template-columns: 1fr;
    }
}

.data-table {
    max-height: 400px;
    overflow-y: auto;
}

.data-table table {
    width: 100%;
    border-collapse: collapse;
}

.data-table th {
    background: rgba(15, 26, 36, 0.8);
    position: sticky;
    top: 0;
    text-align: left;
    padding: 10px 12px;
    font-weight: 500;
    color: var(--accent);
    border-bottom: 1px solid var(--line);
}

.data-table td {
    padding: 8px 12px;
    border-bottom: 1px solid rgba(31, 43, 54, 0.5);
}

.data-table tr:hover {
    background: rgba(31, 43, 54, 0.3);
}

.action-buttons {
    display: flex;
    gap: 12px;
    margin-top: 20px;
}

.footer {
    margin-top: 20px;
    padding-top: 20px;
    border-top: 1px solid var(--line);
    font-size: 0.85rem;
    color: var(--muted);
    text-align: center;
}
"""

# ... [All your backend code remains the same] ...

# =========================
# BUILD ENHANCED UI
# =========================
with gr.Blocks(css=ENHANCED_CSS, theme=gr.themes.Default()) as demo:
    # Header
    with gr.Column(elem_id="header"):
        gr.Markdown("## 🧬 Neuroevolution Playground", elem_classes=["header-title"])
        gr.Markdown("Evolve neural architectures using genetic algorithms", 
                   elem_classes=["header-subtitle"])
    
    with gr.Row():
        # Left Panel - Controls
        with gr.Column(scale=1):
            # Parameters Group
            with gr.Group(elem_classes=["control-group"]):
                gr.Markdown("### 🛠 Evolution Parameters")
                
                with gr.Column():
                    dataset = gr.Dropdown(
                        label="Evaluation Dataset",
                        choices=["Demo (Surrogate)", "PIQA (Phase 2)", "HellaSwag (Phase 2)"],
                        value="Demo (Surrogate)",
                        info="Dataset used for fitness evaluation"
                    )
                    
                    with gr.Row():
                        with gr.Column():
                            pop = gr.Slider(8, 80, value=24, step=2, label="Population Size", 
                                          elem_classes=["param-slider"])
                            gens = gr.Slider(5, 200, value=60, step=1, label="Max Generations", 
                                           elem_classes=["param-slider"])
                            mut = gr.Slider(0.05, 0.9, value=0.25, step=0.01, label="Mutation Rate", 
                                          elem_classes=["param-slider"])
                        with gr.Column():
                            explore = gr.Slider(0.0, 1.0, value=0.35, step=0.05, label="Exploration", 
                                              elem_classes=["param-slider"])
                            exploit = gr.Slider(0.0, 1.0, value=0.65, step=0.05, label="Exploitation", 
                                              elem_classes=["param-slider"])
                            seed = gr.Number(value=42, label="Random Seed", precision=0)
                    
                    pace = gr.Slider(0, 1000, value=120, step=10, label="Simulation Speed (ms)", 
                                   elem_classes=["param-slider"])
                    metric_choice = gr.Radio(choices=["Accuracy", "Fitness"], value="Accuracy", 
                                          label="History Metric Display")
            
            # Status Panel
            with gr.Group(elem_classes=["panel", "stats-panel"]):
                gr.Markdown("### 📊 Current Status")
                stats_md = gr.Markdown("Ready. Press **Start** to begin evolution.", elem_id="stats")
            
            # Action Buttons
            with gr.Row(elem_classes=["action-buttons"]):
                start = gr.Button("▶ Start Evolution", variant="primary", elem_classes=["btn-primary"])
                stop = gr.Button("⏹ Stop", variant="stop", elem_classes=["btn-danger"], interactive=False)
                clear = gr.Button("↻ Reset", elem_classes=["btn-secondary"])
            
            # Export
            with gr.Group(elem_classes=["panel"]):
                gr.Markdown("### 💾 Export Results")
                with gr.Row():
                    export_btn = gr.Button("Save Snapshot (JSON)")
                    export_file = gr.File(label="Download snapshot", visible=False)
        
        # Right Panel - Visualizations
        with gr.Column(scale=2):
            # 3D Visualization
            with gr.Group(elem_classes=["panel", "viz-panel"]):
                with gr.Column(elem_classes=["viz-header"]):
                    with gr.Row():
                        gr.Markdown("### 🌐 Architecture Space", elem_classes=["viz-title"])
                        gen_counter = gr.Markdown("", elem_classes=["gen-counter"])
                sphere_html = gr.HTML()
            
            # History Visualization
            with gr.Group(elem_classes=["panel", "viz-panel"]):
                with gr.Column(elem_classes=["viz-header"]):
                    gr.Markdown("### 📈 Performance History", elem_classes=["viz-title"])
                hist_html = gr.HTML()
            
            # Results Table
            with gr.Group(elem_classes=["panel"]):
                gr.Markdown("### 🏆 Top Genomes")
                with gr.Column(elem_classes=["data-table"]):
                    top_df = gr.Dataframe(
                        label="",
                        headers=["Fitness", "Accuracy", "d_model", "Layers", "Heads", "FFN", "Mem", "Dropout", "Params"],
                        datatype=["number", "number", "number", "number", "number", "number", "number", "number", "number"],
                        wrap=True,
                        interactive=False
                    )
    
    # Footer
    with gr.Column(elem_classes=["footer"]):
        gr.Markdown("Evotransformer Playground v1.0 • Using Plotly and Gradio")
    
    # Wiring
    start.click(
        start_evo, 
        [dataset, pop, gens, mut, explore, exploit, seed, pace, metric_choice], 
        [start, stop, clear]
    )
    stop.click(stop_evo, [], [start, stop, clear])
    clear.click(
        clear_evo, 
        [], 
        [sphere_html, hist_html, stats_md, top_df, start, stop, clear]
    )
    export_btn.click(export_snapshot, [], [export_file])
    
    # State polling
    demo.load(poll_state, None, [sphere_html, hist_html, stats_md, top_df, gen_counter])
    gr.Timer(0.7).tick(poll_state, None, [sphere_html, hist_html, stats_md, top_df, gen_counter])

if __name__ == "__main__":
    demo.launch()