import gradio as gr
from evo_transformer import EvoTransformer
from plot import plot_radar_chart
import io
from PIL import Image

evo = EvoTransformer()

def evolve_model(generations):
    evo.reset()
    evo.evolve(generations)

    best_config = evo.get_best_config()
    accuracy = f"{best_config['accuracy']:.2f}%"
    params = f"{best_config['parameters']:.2f}M params"
    summary = str(best_config['traits'])

    # Radar chart
    radar_img = plot_radar_chart(best_config['traits'])

    # Files
    csv_file = ("evo_history.csv", io.BytesIO(evo.export_csv().encode()))
    json_file = ("evo_history.json", io.BytesIO(evo.export_json().encode()))

    # History
    history_text = [f"Gen {i+1} Config: {c['traits']}" for i, c in enumerate(evo.history)]

    return (
        accuracy, params, summary,
        radar_img,
        *history_text,
        csv_file, json_file
    )

with gr.Blocks(title="EvoTransformer – Evolving Transformer Architectures") as demo:
    gr.Markdown("🧬 **EvoTransformer – Evolving Transformer Architectures**")
    gr.Markdown("Simulate trait mutation and adaptive architecture generation.")

    generations = gr.Slider(minimum=1, maximum=20, step=1, label="Number of Generations", value=5)
    run_btn = gr.Button("🧬 Evolve Architecture")

    acc = gr.Textbox(label="Simulated Accuracy")
    params = gr.Textbox(label="Estimated Parameters")
    summary = gr.Textbox(label="Current Config Summary")
    radar = gr.Image(label="Final Generation Trait Radar")

    history_outputs = [gr.Textbox(label=f"Gen {i+1} Config") for i in range(10)]
    csv_download = gr.File(label="Download CSV History")
    json_download = gr.File(label="Download JSON History")

    run_btn.click(
        fn=evolve_model,
        inputs=[generations],
        outputs=[acc, params, summary, radar] + history_outputs + [csv_download, json_download]
    )

demo.launch()