import gradio as gr
from evo_transformer import EvoTransformer
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns

def simulate_evolution(generations):
    model = EvoTransformer()
    model.evolve(generations)
    history = model.get_history()
    
    # Evaluation of final model
    final_result = model.evaluate()
    
    # Visualization
    fig, ax = plt.subplots(figsize=(10, 5))
    df = pd.DataFrame(history)
    sns.lineplot(data=df.drop("memory", axis=1), markers=True, ax=ax)
    ax.set_title("EvoTransformer Trait Evolution")
    ax.set_ylabel("Value")
    ax.set_xlabel("Generation")
    plt.xticks(range(len(df)))
    plt.tight_layout()
    
    return (
        final_result["accuracy"],
        round(final_result["params"], 2),
        fig
    )

# UI
with gr.Blocks(title="EvoTransformer: Adaptive Architecture Evolution") as demo:
    gr.Markdown("# 🧬 EvoTransformer Demo")
    gr.Markdown("Evolving architecture live — inspired by nature. Tune the number of generations below and visualize how traits change during evolution.")

    generations = gr.Slider(1, 20, value=5, step=1, label="Number of Generations")
    btn = gr.Button("Evolve EvoTransformer")

    accuracy = gr.Textbox(label="Estimated Accuracy", interactive=False)
    params = gr.Textbox(label="Estimated Parameter Count (M)", interactive=False)
    plot = gr.Plot(label="Architecture Trait Evolution")

    btn.click(fn=simulate_evolution, inputs=generations, outputs=[accuracy, params, plot])

demo.launch()