import matplotlib.pyplot as plt
import numpy as np

import gradio as gr


# Sample stats
MODELS = {
    "llama" : {"passed": 14, "failed": 1, "skipped": 6, "error": 0},
    "gemma3" : {"passed": 42, "failed": 6, "skipped": 12, "error": 0},
    "csm" : {"passed": 0, "failed": 0, "skipped": 0, "error": 1},
}

def plot_model_stats(model_name: str) -> plt.Figure:
    """Draws a pie chart of model's passed, failed, skipped, and error stats."""
    model_stats = MODELS[model_name]
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.pie(model_stats.values(), labels=model_stats.keys())
    return fig

demo = gr.Interface(
    plot_model_stats,
    [
        gr.Radio(list(MODELS.keys()), label="Model"),
    ],
    gr.Plot(label="Model Stats", format="png"),
)

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