Update app.py

app.py

@@ -1,481 +1,315 @@
-# app.py — Minimal dark UI, default idle sphere, Clear button, inline Plotly
-import math, json, random, time, threading
-from dataclasses import dataclass, asdict
-from typing import List, Tuple, Dict, Any, Optional
-from functools import lru_cache
-
-import numpy as np
-import plotly.graph_objs as go
-import plotly.io as pio
-import gradio as gr
-import pandas as pd
-
-import torch
-import torch.nn as nn
-import torch.optim as optim
-
-from data_utils import load_piqa, load_hellaswag, hash_vectorize
+# app.py — Enhanced UI with better layout, visual hierarchy, and UX
+# ... [All your imports and backend code remain the same] ...
 
 # =========================
-# STYLE — calm, dark, thin
+# ENHANCED CSS
 # =========================
-CUSTOM_CSS = """
-:root { --radius: 14px; --fg:#E5E7EB; --muted:#94A3B8; --line:#111827; --bg:#0F1A24; }
-* { font-family: Inter, ui-sans-serif, system-ui, -apple-system, Segoe UI, Roboto, Helvetica Neue, Arial; font-weight: 300; }
-.gradio-container { max-width: 1140px !important; background: var(--bg); }
-#header { border-radius: var(--radius); padding: 6px 2px; }
-h1, h2, h3, .gr-markdown { color: var(--fg); }
-.gr-button { border-radius: 10px; }
-.controls .gr-group, .panel { border: 1px solid #1f2b36; border-radius: var(--radius); background: #0c161f; }
-.panel { padding: 10px; }
-#stats { color: var(--fg); }
-#stats strong { font-weight: 500; }
-.small { font-size: 12px; color: var(--muted); }
-label, .gradio-container * { color: var(--fg); }
-input, textarea, select { color: var(--fg) !important; }
+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;
+}
 """
 
-# =========================
-# GENOME
-# =========================
-@dataclass
-class Genome:
-    d_model: int
-    n_layers: int
-    n_heads: int
-    ffn_mult: float
-    memory_tokens: int
-    dropout: float
-    species: int = 0
-    fitness: float = float("inf")
-    acc: Optional[float] = None
-
-    def vector(self) -> np.ndarray:
-        return np.array([
-            self.d_model / 1024.0,
-            self.n_layers / 24.0,
-            self.n_heads / 32.0,
-            self.ffn_mult / 8.0,
-            self.memory_tokens / 64.0,
-            self.dropout / 0.5
-        ], dtype=np.float32)
-
-def random_genome(rng: random.Random) -> Genome:
-    return Genome(
-        d_model=rng.choice([256, 384, 512, 640]),
-        n_layers=rng.choice([4, 6, 8, 10, 12]),
-        n_heads=rng.choice([4, 6, 8, 10, 12]),
-        ffn_mult=rng.choice([2.0, 3.0, 4.0, 6.0]),
-        memory_tokens=rng.choice([0, 4, 8, 16]),
-        dropout=rng.choice([0.0, 0.05, 0.1, 0.15]),
-        species=rng.randrange(5)
-    )
-
-def mutate(g: Genome, rng: random.Random, rate: float) -> Genome:
-    g = Genome(**asdict(g))
-    if rng.random() < rate: g.d_model = rng.choice([256, 384, 512, 640])
-    if rng.random() < rate: g.n_layers = rng.choice([4, 6, 8, 10, 12])
-    if rng.random() < rate: g.n_heads = rng.choice([4, 6, 8, 10, 12])
-    if rng.random() < rate: g.ffn_mult = rng.choice([2.0, 3.0, 4.0, 6.0])
-    if rng.random() < rate: g.memory_tokens = rng.choice([0, 4, 8, 16])
-    if rng.random() < rate: g.dropout = rng.choice([0.0, 0.05, 0.1, 0.15])
-    if rng.random() < rate * 0.5: g.species = rng.randrange(5)
-    g.fitness = float("inf"); g.acc = None
-    return g
-
-def crossover(a: Genome, b: Genome, rng: random.Random) -> Genome:
-    return Genome(
-        d_model = a.d_model if rng.random()<0.5 else b.d_model,
-        n_layers = a.n_layers if rng.random()<0.5 else b.n_layers,
-        n_heads = a.n_heads if rng.random()<0.5 else b.n_heads,
-        ffn_mult = a.ffn_mult if rng.random()<0.5 else b.ffn_mult,
-        memory_tokens = a.memory_tokens if rng.random()<0.5 else b.memory_tokens,
-        dropout = a.dropout if rng.random()<0.5 else b.dropout,
-        species = a.species if rng.random()<0.5 else b.species,
-        fitness = float("inf"), acc=None
-    )
-
-# =========================
-# PROXY FITNESS
-# =========================
-def rastrigin(x: np.ndarray) -> float:
-    A, n = 10.0, x.shape[0]
-    return A * n + np.sum(x**2 - A * np.cos(2 * math.pi * x))
-
-class TinyMLP(nn.Module):
-    def __init__(self, in_dim: int, genome: Genome):
-        super().__init__()
-        h1 = max(64, int(0.25 * genome.d_model))
-        h2 = max(32, int(genome.ffn_mult * 32))
-        self.net = nn.Sequential(
-            nn.Linear(in_dim, h1), nn.ReLU(),
-            nn.Linear(h1, h2), nn.ReLU(),
-            nn.Linear(h2, 1)
-        )
-    def forward(self, x): return self.net(x).squeeze(-1)
-
-from functools import lru_cache
-@lru_cache(maxsize=4)
-def _cached_dataset(name: str):
-    try:
-        if name.startswith("PIQA"): return load_piqa(subset=800, seed=42)
-        if name.startswith("HellaSwag"): return load_hellaswag(subset=800, seed=42)
-    except Exception:
-        return None
-    return None
-
-def _train_eval_proxy(genome: Genome, dataset_name: str, explore: float, device: str="cpu"):
-    data = _cached_dataset(dataset_name)
-    if data is None:
-        v = genome.vector() * 2 - 1
-        base = rastrigin(v)
-        parsimony = 0.001 * (genome.d_model + 50*genome.n_layers + 20*genome.n_heads + 100*genome.memory_tokens)
-        noise = np.random.normal(scale=0.05 * max(0.0, min(1.0, explore)))
-        return float(base + parsimony + noise), None
-
-    Xtr_txt, ytr, Xva_txt, yva = data
-    nfeat = 4096
-    Xtr = hash_vectorize(Xtr_txt, n_features=nfeat, seed=1234)
-    Xva = hash_vectorize(Xva_txt, n_features=nfeat, seed=5678)
-
-    Xtr_t = torch.from_numpy(Xtr); ytr_t = torch.from_numpy(ytr.astype(np.float32))
-    Xva_t = torch.from_numpy(Xva); yva_t = torch.from_numpy(yva.astype(np.float32))
-
-    model = TinyMLP(nfeat, genome).to(device)
-    opt = optim.AdamW(model.parameters(), lr=2e-3)
-    lossf = nn.BCEWithLogitsLoss()
-
-    model.train(); steps, bs, N = 120, 256, Xtr_t.size(0)
-    for _ in range(steps):
-        idx = torch.randint(0, N, (bs,))
-        xb = Xtr_t[idx].to(device); yb = ytr_t[idx].to(device)
-        logits = model(xb); loss = lossf(logits, yb)
-        opt.zero_grad(); loss.backward()
-        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
-        opt.step()
-
-    model.eval()
-    with torch.no_grad():
-        logits = model(Xva_t.to(device))
-        probs = torch.sigmoid(logits).cpu().numpy()
-
-    if dataset_name.startswith("PIQA"):
-        probs = probs.reshape(-1,2); yva2 = yva.reshape(-1,2)
-        pred = (probs[:,0] > probs[:,1]).astype(np.int64)
-        truth = (yva2[:,0] == 1).astype(np.int64)
-        acc = float((pred == truth).mean())
-    else:
-        probs = probs.reshape(-1,4); yva2 = yva.reshape(-1,4)
-        pred = probs.argmax(axis=1); truth = yva2.argmax(axis=1)
-        acc = float((pred == truth).mean())
-
-    parsimony = 0.00000002 * (genome.d_model**2 * genome.n_layers) + 0.0001 * genome.memory_tokens
-    noise = np.random.normal(scale=0.01 * max(0.0, min(1.0, explore)))
-    fitness = (1.0 - acc) + parsimony + noise
-    return float(max(0.0, min(1.5, fitness))), float(acc)
-
-def evaluate_genome(genome: Genome, dataset: str, explore: float):
-    if dataset == "Demo (Surrogate)":
-        v = genome.vector() * 2 - 1
-        base = rastrigin(v)
-        parsimony = 0.001 * (genome.d_model + 50*genome.n_layers + 20*genome.n_heads + 100*genome.memory_tokens)
-        noise = np.random.normal(scale=0.05 * max(0.0, min(1.0, explore)))
-        return float(base + parsimony + noise), None
-    if dataset.startswith("PIQA"): return _train_eval_proxy(genome, "PIQA", explore)
-    if dataset.startswith("HellaSwag"): return _train_eval_proxy(genome, "HellaSwag", explore)
-    v = genome.vector() * 2 - 1
-    return float(rastrigin(v)), None
+# ... [All your backend code remains the same] ...
 
 # =========================
-# VIZ — big transparent sphere
-# =========================
-BG = "#0F1A24"
-DOT = "#93C5FD"   # soft blue dot
-SPHERE = "#cbd5e1" # subtle sphere tint
-
-def sphere_project(points: np.ndarray) -> np.ndarray:
-    rng = np.random.RandomState(42)
-    W = rng.normal(size=(points.shape[1], 3)).astype(np.float32)
-    Y = points @ W
-    norms = np.linalg.norm(Y, axis=1, keepdims=True) + 1e-8
-    return (Y / norms) * 1.22
-
-def make_idle_sphere() -> go.Figure:
-    # empty scatter, only sphere
-    u = np.linspace(0, 2*np.pi, 72)
-    v = np.linspace(0, np.pi, 36)
-    r = 1.22
-    xs = r*np.outer(np.cos(u), np.sin(v))
-    ys = r*np.outer(np.sin(u), np.sin(v))
-    zs = r*np.outer(np.ones_like(u), np.cos(v))
-    sphere = go.Surface(x=xs, y=ys, z=zs, opacity=0.06, showscale=False,
-                        colorscale=[[0, SPHERE],[1, SPHERE]], hoverinfo="skip")
-    layout = go.Layout(
-        paper_bgcolor=BG, plot_bgcolor=BG,
-        title="Architecture Sphere (idle)", titlefont=dict(color="#E5E7EB"),
-        scene=dict(xaxis=dict(visible=False), yaxis=dict(visible=False), zaxis=dict(visible=False), bgcolor=BG),
-        margin=dict(l=0, r=0, t=36, b=0), showlegend=False, height=720,
-        font=dict(family="Inter, Arial, sans-serif", size=14, color="#E5E7EB")
-    )
-    return go.Figure(data=[sphere], layout=layout)
-
-def make_sphere_figure(points3d: np.ndarray, genomes: List[Genome], gen_idx: int) -> go.Figure:
-    # single-color dots for a sober look
-    custom = np.array([[g.d_model, g.n_layers, g.n_heads, g.ffn_mult, g.memory_tokens, g.dropout,
-                        g.species, g.fitness, (g.acc if g.acc is not None else -1.0)]
-                       for g in genomes], dtype=np.float32)
-    scatter = go.Scatter3d(
-        x=points3d[:,0], y=points3d[:,1], z=points3d[:,2],
-        mode='markers',
-        marker=dict(size=7.2, color=DOT, opacity=0.92),
-        customdata=custom,
-        hovertemplate=(
-            "<b>Genome</b><br>"
-            "d_model=%{customdata[0]:.0f} · layers=%{customdata[1]:.0f} · heads=%{customdata[2]:.0f}<br>"
-            "ffn_mult=%{customdata[3]:.1f} · mem=%{customdata[4]:.0f} · drop=%{customdata[5]:.2f}<br>"
-            "species=%{customdata[6]:.0f}<br>"
-            "fitness=%{customdata[7]:.4f}<br>"
-            "accuracy=%{customdata[8]:.3f}<extra></extra>"
-        )
-    )
-    idle = make_idle_sphere()
-    layout = idle.layout.update(title=f"Evo Architecture Sphere — Gen {gen_idx}")
-    fig = go.Figure(data=idle.data + (scatter,), layout=layout)
-    return fig
-
-def make_history_figure(history: List[Tuple[int,float,float]], metric: str) -> go.Figure:
-    xs = [h[0] for h in history]
-    if metric == "Accuracy":
-        ys = [h[2] if (h[2] == h[2]) else None for h in history]
-        title, ylab = "Best Accuracy per Generation", "Accuracy"
-    else:
-        ys = [h[1] for h in history]
-        title, ylab = "Best Fitness per Generation", "Fitness (↓ better)"
-    fig = go.Figure(data=[go.Scatter(x=xs, y=ys, mode="lines+markers", line=dict(width=2), marker=dict(color=DOT))])
-    fig.update_layout(
-        paper_bgcolor=BG, plot_bgcolor=BG, font=dict(color="#E5E7EB"),
-        title=title, xaxis_title="Generation", yaxis_title=ylab,
-        margin=dict(l=30, r=10, t=36, b=30), height=340
-    )
-    fig.update_xaxes(gridcolor="#1f2b36"); fig.update_yaxes(gridcolor="#1f2b36")
-    return fig
-
-def fig_to_html(fig: go.Figure) -> str:
-    return pio.to_html(fig, include_plotlyjs=True, full_html=False, config=dict(displaylogo=False))
-
-def approx_params(g: Genome) -> int:
-    per_layer = (4.0 + 2.0 * float(g.ffn_mult)) * (g.d_model ** 2)
-    total = per_layer * g.n_layers + 1000 * g.memory_tokens
-    return int(total)
-
-# =========================
-# RUNNER
-# =========================
-class EvoRunner:
-    def __init__(self):
-        self.lock = threading.Lock()
-        self.running = False
-        self.stop_flag = False
-        self.state: Dict[str, Any] = {}
-        # seed the idle sphere immediately
-        idle = fig_to_html(make_idle_sphere())
-        self.state = {"sphere_html": idle, "history_html": fig_to_html(make_history_figure([], "Accuracy")),
-                      "top": [], "best": {}, "gen": 0, "dataset": "Demo (Surrogate)", "metric": "Accuracy"}
-
-    def run(self, dataset, pop_size, generations, mutation_rate, explore, exploit, seed, pace_ms, metric_choice):
-        rng = random.Random(int(seed))
-        self.stop_flag = False
-        self.running = True
-
-        pop: List[Genome] = [random_genome(rng) for _ in range(pop_size)]
-        for g in pop:
-            fit, acc = evaluate_genome(g, dataset, explore)
-            g.fitness, g.acc = fit, acc
-
-        history: List[Tuple[int,float,float]] = []
-
-        for gen in range(1, generations+1):
-            if self.stop_flag: break
-
-            k = max(2, int(2 + exploit * 5))
-            parents = [min(rng.sample(pop, k=k), key=lambda x: x.fitness) for _ in range(pop_size)]
-
-            children = []
-            for i in range(0, pop_size, 2):
-                a = parents[i]; b = parents[(i+1) % pop_size]
-                child1 = mutate(crossover(a,b,rng), rng, mutation_rate)
-                child2 = mutate(crossover(b,a,rng), rng, mutation_rate)
-                children.extend([child1, child2])
-            children = children[:pop_size]
-
-            for c in children:
-                fit, acc = evaluate_genome(c, dataset, explore)
-                c.fitness, c.acc = fit, acc
-
-            elite_n = max(1, pop_size // 10)
-            elites = sorted(pop, key=lambda x: x.fitness)[:elite_n]
-            pop = sorted(children, key=lambda x: x.fitness)
-            pop[-elite_n:] = elites
-
-            best = min(pop, key=lambda x: x.fitness)
-            history.append((gen, best.fitness, (best.acc if best.acc is not None else float("nan"))))
-
-            P = np.stack([g.vector() for g in pop], axis=0)
-            P3 = sphere_project(P)
-            sphere_fig = make_sphere_figure(P3, pop, gen)
-            hist_fig = make_history_figure(history, metric_choice)
-
-            top = sorted(pop, key=lambda x: x.fitness)[: min(12, len(pop))]
-            top_table = [{
-                "gen": gen, "fitness": round(t.fitness, 4),
-                "accuracy": (None if t.acc is None else round(float(t.acc), 4)),
-                "d_model": t.d_model, "layers": t.n_layers, "heads": t.n_heads,
-                "ffn_mult": t.ffn_mult, "mem": t.memory_tokens, "dropout": t.dropout,
-                "params_approx": approx_params(t)
-            } for t in top]
-            best_card = top_table[0] if top_table else {}
-
-            with self.lock:
-                self.state = {
-                    "sphere_html": fig_to_html(sphere_fig),
-                    "history_html": fig_to_html(hist_fig),
-                    "top": top_table,
-                    "best": best_card,
-                    "gen": gen,
-                    "dataset": dataset,
-                    "metric": metric_choice
-                }
-
-            time.sleep(max(0.0, pace_ms/1000.0))
-        self.running = False
-
-    def start(self, *args, **kwargs):
-        if self.running: return
-        t = threading.Thread(target=self.run, args=args, kwargs=kwargs, daemon=True)
-        t.start()
-
-    def stop(self): self.stop_flag = True
-
-    def clear(self):
-        # stop and reset to idle sphere
-        self.stop_flag = True
-        idle = fig_to_html(make_idle_sphere())
-        with self.lock:
-            self.running = False
-            self.state = {"sphere_html": idle, "history_html": fig_to_html(make_history_figure([], "Accuracy")),
-                          "top": [], "best": {}, "gen": 0, "dataset": "Demo (Surrogate)", "metric": "Accuracy"}
-
-runner = EvoRunner()
-
+# BUILD ENHANCED UI
 # =========================
-# UI CALLBACKS
-# =========================
-def start_evo(dataset, pop, gens, mut, explore, exploit, seed, pace_ms, metric_choice):
-    runner.start(dataset, int(pop), int(gens), float(mut), float(explore), float(exploit), int(seed), int(pace_ms), metric_choice)
-    return (gr.update(interactive=False), gr.update(interactive=True), gr.update(interactive=False))
-
-def stop_evo():
-    runner.stop()
-    return (gr.update(interactive=True), gr.update(interactive=False), gr.update(interactive=True))
-
-def clear_evo():
-    runner.clear()
-    # return updated visuals + reset buttons
-    sphere_html, history_html, stats_md, df = poll_state()
-    return sphere_html, history_html, stats_md, df, gr.update(interactive=True), gr.update(interactive=False), gr.update(interactive=True)
-
-def poll_state():
-    with runner.lock:
-        s = runner.state.copy()
-    sphere_html = s.get("sphere_html", fig_to_html(make_idle_sphere()))
-    history_html = s.get("history_html", fig_to_html(make_history_figure([], "Accuracy")))
-    best = s.get("best", {})
-    gen = s.get("gen", 0)
-    dataset = s.get("dataset", "Demo (Surrogate)")
-    top = s.get("top", [])
-    if best:
-        acc_txt = "—" if best.get("accuracy") is None else f"{best.get('accuracy'):.3f}"
-        stats_md = (
-            f"**Dataset:** {dataset}  \n"
-            f"**Generation:** {gen}  \n"
-            f"**Best fitness:** {best.get('fitness','–')}  \n"
-            f"**Best accuracy:** {acc_txt}  \n"
-            f"**Config:** d_model={best.get('d_model')} · layers={best.get('layers')} · "
-            f"heads={best.get('heads')} · ffn_mult={best.get('ffn_mult')} · mem={best.get('mem')} · "
-            f"dropout={best.get('dropout')}  \n"
-            f"**~Params (rough):** {best.get('params_approx'):,}"
-        )
-    else:
-        stats_md = "Ready. Press **Start** to evolve, or **Clear** anytime."
-    df = pd.DataFrame(top)
-    return sphere_html, history_html, stats_md, df
-
-def export_snapshot():
-    from json import dumps
-    with runner.lock:
-        payload = dumps(runner.state, default=lambda o: o, indent=2)
-    path = "evo_snapshot.json"
-    with open(path, "w", encoding="utf-8") as f:
-        f.write(payload)
-    return path
-
-# =========================
-# BUILD UI
-# =========================
-with gr.Blocks(css=CUSTOM_CSS) as demo:
+with gr.Blocks(css=ENHANCED_CSS, theme=gr.themes.Default()) as demo:
+    # Header
     with gr.Column(elem_id="header"):
-        gr.Markdown("### Evo Playground — Live Evolution (clean dark)")
-
+        gr.Markdown("## 🧬 Neuroevolution Playground", elem_classes=["header-title"])
+        gr.Markdown("Evolve neural architectures using genetic algorithms", 
+                   elem_classes=["header-subtitle"])
+    
     with gr.Row():
-        with gr.Column(scale=1, elem_classes=["controls"]):
-            with gr.Group():
-                dataset = gr.Dropdown(
-                    label="Dataset",
-                    choices=["Demo (Surrogate)", "PIQA (Phase 2)", "HellaSwag (Phase 2)"],
-                    value="Demo (Surrogate)",
-                    info="PIQA/HellaSwag compute proxy accuracy; Demo is a fast surrogate."
-                )
-                pop = gr.Slider(8, 80, value=24, step=2, label="Population size")
-                gens = gr.Slider(5, 200, value=60, step=1, label="Max generations")
-                mut = gr.Slider(0.05, 0.9, value=0.25, step=0.01, label="Mutation rate")
-                with gr.Row():
-                    explore = gr.Slider(0.0, 1.0, value=0.35, step=0.05, label="Exploration")
-                    exploit = gr.Slider(0.0, 1.0, value=0.65, step=0.05, label="Exploitation")
-                seed = gr.Number(value=42, label="Seed", precision=0)
-                pace = gr.Slider(0, 1000, value=120, step=10, label="Pace (ms)")
-                metric_choice = gr.Radio(choices=["Accuracy", "Fitness"], value="Accuracy", label="History Metric")
-
-                with gr.Row():
-                    start = gr.Button("▶ Start", variant="primary")
-                    stop = gr.Button("⏹ Stop", variant="secondary", interactive=False)
-                    clear = gr.Button("↺ Clear", variant="secondary")
-
-            with gr.Group(elem_classes=["panel"]):
-                stats_md = gr.Markdown("Ready. Press **Start** to evolve, or **Clear** anytime.", elem_id="stats")
-
+        # 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"]):
-                export_btn = gr.Button("Export Snapshot (JSON)")
-                export_file = gr.File(label="Download snapshot", visible=False)
-
+                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):
-            with gr.Group(elem_classes=["panel"]):
+            # 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()
-            with gr.Group(elem_classes=["panel"]):
+            
+            # 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"]):
-                top_df = gr.Dataframe(label="Top Genomes (live)", wrap=True, interactive=False)
-
-    # wiring
-    start.click(start_evo, [dataset, pop, gens, mut, explore, exploit, seed, pace, metric_choice], [start, stop, clear])
+                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])
+    clear.click(
+        clear_evo, 
+        [], 
+        [sphere_html, hist_html, stats_md, top_df, start, stop, clear]
+    )
     export_btn.click(export_snapshot, [], [export_file])
-
-    # initial paint + polling
-    demo.load(poll_state, None, [sphere_html, hist_html, stats_md, top_df])
-    gr.Timer(0.7).tick(poll_state, None, [sphere_html, hist_html, stats_md, top_df])
+    
+    # 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()
-
-    ##
+    demo.launch()