# app.py
import gradio as gr
from evo_transformer import EvoTransformer
import matplotlib.pyplot as plt

def run_evolution(generations):
    evo = EvoTransformer()
    evo.evolve(generations)
    history = evo.get_history()
    results = evo.evaluate()

    # Format history for output
    trait_logs = ""
    for i, config in enumerate(history):
        trait_logs += f"Generation {i}: {config}\n"

    # Plot number of attention heads over generations
    heads = [conf["attention_heads"] for conf in history]
    layers = [conf["layers"] for conf in history]
    ffn = [conf["ffn_dim"] for conf in history]

    fig, ax = plt.subplots()
    ax.plot(range(len(heads)), heads, label="Attention Heads")
    ax.plot(range(len(layers)), layers, label="Layers")
    ax.plot(range(len(ffn)), ffn, label="FFN Dim")
    ax.set_title("Evolution of Traits")
    ax.set_xlabel("Generation")
    ax.set_ylabel("Value")
    ax.legend()

    return trait_logs, results["accuracy"], round(results["params"], 2), fig

# Gradio UI
demo = gr.Interface(
    fn=run_evolution,
    inputs=gr.Slider(minimum=1, maximum=10, step=1, label="Generations"),
    outputs=[
        gr.Textbox(label="Evolution History"),
        gr.Number(label="Simulated Accuracy"),
        gr.Number(label="Estimated Parameters (M)"),
        gr.Plot(label="Trait Evolution Plot"),
    ],
    title="🧬 EvoTransformer Demo",
    description="An evolving Transformer that mutates architecture traits during training. Watch the architecture change in real time!"
)

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