# app.py — Enhanced UI + stable backend (idle sphere, Clear, inline Plotly, accuracy)
import math, 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

# =========================
# 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;
}
"""

# =========================
# GENOME + EVOLUTION CORE
# =========================
@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)

@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:
        # Fallback to surrogate so UI still runs
        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

# =========================
# VIZ — idle sphere, big transparent surface
# =========================
BG = "#0F1A24"
DOT = "#93C5FD"
SPHERE = "#cbd5e1"

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:
    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=dict(text="Architecture Space (idle)", font=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:
    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.0, 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>"
            "fitness=%{customdata[7]:.4f} · acc=%{customdata[8]:.3f}<extra></extra>"
        )
    )
    idle = make_idle_sphere()
    fig = go.Figure(data=idle.data + (scatter,), layout=idle.layout)
    fig.update_layout(title=dict(text=f"Evo Architecture Space — Gen {gen_idx}"))
    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=dict(text=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:
    # Inline Plotly JS so it renders even without CDN
    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 idle visuals
        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()

# =========================
# 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()
    sphere_html, history_html, stats_md, df, gen_counter_md = poll_state()
    return sphere_html, history_html, stats_md, df, gen_counter_md, 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", "")
    history_html = s.get("history_html", "")
    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 begin evolution."

    df = pd.DataFrame(top)
    gen_counter_md = f"Gen **{gen}**"
    return sphere_html, history_html, stats_md, df, gen_counter_md

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 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")
        gr.Markdown("Evolve neural architectures using genetic algorithms")

    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="secondary", 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="", wrap=True, interactive=False)

    # Footer
    with gr.Column(elem_classes=["footer"]):
        gr.Markdown("Neuroevolution Playground v1.0 • Plotly + 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, gen_counter, 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()