import gradio as gr
from agents import evolve_agent
from emotion import get_emotion_level
from shared_memory import MemoryStore
import matplotlib.pyplot as plt
import io
import base64

memory = MemoryStore()

def plot_scores(score_history):
    plt.figure(figsize=(6, 3))
    for agent, scores in score_history.items():
        plt.plot(scores, label=agent)
    plt.xlabel("Iteration")
    plt.ylabel("Score")
    plt.title("Agent Performance Over Time")
    plt.legend()
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)
    img_base64 = base64.b64encode(buf.read()).decode('utf-8')
    plt.close()
    return f"data:image/png;base64,{img_base64}"

def evolve_step(target_value: float, emotion_input: str):
    emotion_mod = get_emotion_level(emotion_input)
    result_a = evolve_agent('Agent A', target_value, emotion_mod, memory)
    result_b = evolve_agent('Agent B', target_value, emotion_mod, memory)
    best = max([result_a, result_b], key=lambda x: x['score'])
    chart = plot_scores(memory.score_history)
    return (f"Best result: {best['agent']} => Expression: {best['expression']} (Score: {best['score']})", chart)

demo = gr.Interface(
    fn=evolve_step,
    inputs=[
        gr.Number(label="Target Value"),
        gr.Radio(["Calm", "Frustrated", "Curious"], label="User Emotion")
    ],
    outputs=["text", gr.Image(type="filepath")],
    title="Multi-Agent Evolution Sandbox (with Score Graph)",
    description="Two agents evolve expressions to match a target number. Emotions influence exploration. See real-time performance graph."
)

demo.launch()