# app.py
import math, json, random, time, threading, io, os
from dataclasses import dataclass, asdict
from typing import List, Tuple, Dict, Any, Optional
import numpy as np
import plotly.graph_objs as go
import gradio as gr

# =========================
# UX THEME & STYLES
# =========================
CUSTOM_CSS = """
:root {
  --radius-2xl: 20px;
}
.gradio-container {max-width: 1400px !important}
#header-card {border-radius: var(--radius-2xl); box-shadow: 0 6px 24px rgba(0,0,0,0.08)}
#viz-card, #right-card, #table-card {border-radius: var(--radius-2xl); box-shadow: 0 6px 24px rgba(0,0,0,0.06)}
#stats {display:flex; gap:16px; flex-wrap:wrap}
.stat {flex:1; min-width:180px; background:#0b1220; color:white; border-radius:16px; padding:14px 16px}
.stat .k {font-size:14px; opacity:0.8}
.stat .v {font-size:22px; font-weight:700}
.gr-button {border-radius:14px}
"""

# =========================
# 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")

    def vector(self) -> np.ndarray:
        # Normalized structural vector (0..1)
        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")
    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")
    )

# =========================
# FITNESS HOOK (Phase 1: fast surrogate)
# Swap this later for real PIQA/HellaSwag evaluation
# =========================
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))

def fitness_hook(genome: Genome, dataset: str, explore: float) -> float:
    """
    Phase 1 (demo, fast):
    - Build vector v in [-1,1] from genome params and score via Rastrigin.
    - Add small parsimony penalty and exploration noise.
    Phase 2 (real):
    - Replace with tiny train/eval steps on chosen dataset (PIQA/HellaSwag/WikiText-ppl).
    """
    v = genome.vector() * 2 - 1  # [-1,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)

# =========================
# PROJECTION & VIZ
# =========================
def sphere_project(points: np.ndarray) -> np.ndarray:
    # Fixed random projection 6D -> 3D then normalize to unit sphere
    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

def make_sphere_figure(points3d: np.ndarray, genomes: List[Genome], gen_idx: int) -> go.Figure:
    species = np.array([g.species for g in genomes])
    tooltip = [
        json.dumps({k:v for k,v in asdict(g).items() if k!="fitness"}) + f"\nfitness={g.fitness:.3f}"
        for g in genomes
    ]

    scatter = go.Scatter3d(
        x=points3d[:,0], y=points3d[:,1], z=points3d[:,2],
        mode='markers',
        marker=dict(size=6, color=species, opacity=0.9),
        text=tooltip, hoverinfo='text'
    )

    # Sphere mesh
    u = np.linspace(0, 2*np.pi, 48)
    v = np.linspace(0, np.pi, 24)
    xs = np.outer(np.cos(u), np.sin(v))
    ys = np.outer(np.sin(u), np.sin(v))
    zs = np.outer(np.ones_like(u), np.cos(v))
    sphere = go.Surface(x=xs, y=ys, z=zs, opacity=0.15, showscale=False)

    layout = go.Layout(
        title=f"Evo Sphere — Generation {gen_idx}",
        scene=dict(xaxis=dict(visible=False), yaxis=dict(visible=False), zaxis=dict(visible=False)),
        margin=dict(l=0, r=0, t=40, b=0),
        showlegend=False
    )
    return go.Figure(data=[sphere, scatter], layout=layout)

def make_history_figure(history: List[Tuple[int,float]]) -> go.Figure:
    xs = [h[0] for h in history]
    ys = [h[1] for h in history]
    fig = go.Figure(data=[go.Scatter(x=xs, y=ys, mode="lines+markers")])
    fig.update_layout(title="Best Fitness per Generation", xaxis_title="Generation",
                      yaxis_title="Fitness (lower is better)",
                      margin=dict(l=30,r=10,t=40,b=30))
    return fig

def approx_params(g: Genome) -> int:
    # Very rough estimate ignoring embeddings/vocab:
    # per-layer ~ (4 + 2*ffn_mult) * d_model^2
    per_layer = (4.0 + 2.0 * float(g.ffn_mult)) * (g.d_model ** 2)
    total = per_layer * g.n_layers
    # tiny bump for memory tokens pathways (illustrative only)
    total += 1000 * g.memory_tokens
    return int(total)

# =========================
# ORCHESTRATOR
# =========================
class EvoRunner:
    def __init__(self):
        self.lock = threading.Lock()
        self.running = False
        self.stop_flag = False
        self.state: Dict[str, Any] = {}

    def run(self, dataset, pop_size, generations, mutation_rate, explore, exploit, seed, pace_ms):
        rng = random.Random(int(seed))
        self.stop_flag = False
        self.running = True

        pop: List[Genome] = [random_genome(rng) for _ in range(pop_size)]
        # initial eval
        for g in pop:
            g.fitness = fitness_hook(g, dataset, explore)

        history: List[Tuple[int,float]] = []
        best_overall: Optional[Genome] = None

        for gen in range(1, generations+1):
            if self.stop_flag: break

            # Selection: tournament size depends on exploitation
            k = max(2, int(2 + exploit * 5))
            parents = []
            for _ in range(pop_size):
                sample = rng.sample(pop, k=k)
                parents.append(min(sample, key=lambda x: x.fitness))

            # Reproduce
            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]

            # Evaluate kids
            for c in children:
                c.fitness = fitness_hook(c, dataset, explore)

            # Elitism
            elite_n = max(1, pop_size // 10)
            elites = sorted(pop, key=lambda x: x.fitness)[:elite_n]

            # Next pop
            pop = sorted(children, key=lambda x: x.fitness)
            pop[-elite_n:] = elites

            best = min(pop, key=lambda x: x.fitness)
            if best_overall is None or best.fitness < best_overall.fitness:
                best_overall = best

            history.append((gen, best.fitness))

            # Viz snapshot
            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)
            top = sorted(pop, key=lambda x: x.fitness)[: min(12, len(pop))]
            top_table = [
                {
                    "gen": gen,
                    "fitness": round(t.fitness, 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,
                    "species": t.species,
                    "params_approx": approx_params(t)
                } for t in top
            ]
            best_card = top_table[0] if len(top_table) else {}

            with self.lock:
                self.state = {
                    "sphere": sphere_fig,
                    "history": hist_fig,
                    "top": top_table,
                    "best": best_card,
                    "gen": gen,
                    "dataset": dataset
                }

            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

runner = EvoRunner()

# =========================
# GRADIO UI CALLBACKS
# =========================
def start_evo(dataset, pop, gens, mut, explore, exploit, seed, pace_ms):
    runner.start(dataset, int(pop), int(gens), float(mut), float(explore), float(exploit), int(seed), int(pace_ms))
    return (gr.update(interactive=False), gr.update(interactive=True))

def stop_evo():
    runner.stop()
    return (gr.update(interactive=True), gr.update(interactive=False))

def poll_state():
    with runner.lock:
        s = runner.state.copy()
    sphere = s.get("sphere", go.Figure())
    history = s.get("history", go.Figure())
    best = s.get("best", {})
    gen = s.get("gen", 0)
    dataset = s.get("dataset", "Demo (Surrogate)")
    top = s.get("top", [])
    if best:
        stats_md = (
            f"**Dataset:** {dataset}  \n"
            f"**Generation:** {gen}  \n"
            f"**Best fitness:** {best.get('fitness','–')}  \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 = "Waiting… click **Start Evolution**."
    import pandas as pd
    df = pd.DataFrame(top)
    return sphere, history, stats_md, df

def export_snapshot():
    with runner.lock:
        payload = json.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(theme=gr.themes.Soft(), css=CUSTOM_CSS) as demo:
    with gr.Column(elem_id="header-card"):
        gr.Markdown(
            "# Evo Playground — Live Evolving Transformer Architectures\n"
            "Watch the population **mutate, recombine, and converge** in real time. "
            "Choose a dataset and search behavior; the 3D sphere shows the architecture landscape (species = colors)."
        )

    with gr.Row():
        # LEFT: Controls
        with gr.Column(scale=1):
            with gr.Group():
                dataset = gr.Dropdown(
                    label="Dataset",
                    choices=["Demo (Surrogate)", "PIQA (Phase 2)", "HellaSwag (Phase 2)", "WikiText Perplexity (Phase 2)"],
                    value="Demo (Surrogate)",
                    info="Demo is instant. Phase 2 datasets will do tiny train/eval steps per genome."
                )
                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 between gens)")
                with gr.Row():
                    start = gr.Button("▶ Start Evolution", variant="primary")
                    stop = gr.Button("⏹ Stop", variant="secondary")

            with gr.Group(elem_id="right-card"):
                stats_md = gr.Markdown("Waiting…")

                export_btn = gr.Button("Export Snapshot (JSON)")
                export_file = gr.File(label="Download snapshot", visible=False)

        # RIGHT: Viz + Table
        with gr.Column(scale=2):
            with gr.Group(elem_id="viz-card"):
                sphere_plot = gr.Plot(label="Evolution Sphere")
            with gr.Group(elem_id="viz-card"):
                hist_plot = gr.Plot(label="Best Fitness History")
            with gr.Group(elem_id="table-card"):
                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], [start, stop])
    stop.click(stop_evo, [], [start, stop])
    export_btn.click(export_snapshot, [], [export_file])

    # Initial paint once when app loads
    demo.load(poll_state, None, [sphere_plot, hist_plot, stats_md, top_df])

    # Continuous polling (every 0.7s)
    poller = gr.Timer(0.7)
    poller.tick(poll_state, None, [sphere_plot, hist_plot, stats_md, top_df])

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