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]
    
    # Define appropriate colors for each category
    colors = {
        'passed': '#4CAF50',    # Green
        'failed': '#F44336',    # Red
        'skipped': '#FF9800',   # Orange
        'error': '#9C27B0'      # Purple
    }
    
    # Filter out categories with 0 values for cleaner visualization
    filtered_stats = {k: v for k, v in model_stats.items() if v > 0}
    
    if not filtered_stats:
        # Handle case where all values are 0 - dark theme styling
        fig, ax = plt.subplots(figsize=(10, 8), facecolor='#000000')
        ax.set_facecolor('#000000')
        ax.text(0.5, 0.5, 'No test results available', 
                horizontalalignment='center', verticalalignment='center',
                transform=ax.transAxes, fontsize=16, color='white')
        ax.set_xlim(0, 1)
        ax.set_ylim(0, 1)
        ax.axis('off')
        return fig
    
    # Create the pie chart with dark theme
    fig, ax = plt.subplots(figsize=(10, 8), facecolor='#000000')
    ax.set_facecolor('#000000')
    
    # Get colors for filtered categories
    chart_colors = [colors[category] for category in filtered_stats.keys()]
    
    # Create pie chart with custom styling for dark theme
    wedges, texts, autotexts = ax.pie(
        filtered_stats.values(), 
        labels=filtered_stats.keys(),
        colors=chart_colors,
        autopct='%1.1f%%',
        startangle=90,
        explode=[0.08] * len(filtered_stats),  # Slightly larger separation for dark theme
        shadow=False,  # Remove shadow for cleaner dark look
        textprops={'fontsize': 13, 'weight': 'bold', 'color': 'white'}
    )
    
    # Enhance text styling for dark theme
    for autotext in autotexts:
        autotext.set_color('white')
        autotext.set_weight('bold')
        autotext.set_fontsize(12)
    
    # Set title with white text
    total_tests = sum(model_stats.values())
    ax.set_title(f'{model_name.upper()} Test Results\n({total_tests} total tests)', 
                fontsize=18, weight='bold', pad=30, color='white')
    
    # Make the chart more minimalist
    plt.tight_layout()
    
    return fig

# Custom CSS for dark theme
dark_theme_css = """
/* Global dark theme */
.gradio-container {
    background-color: #000000 !important;
    color: white !important;
}

/* Remove borders from all components */
.gr-box, .gr-form, .gr-panel {
    border: none !important;
    background-color: #000000 !important;
}

/* Sidebar styling */
.sidebar {
    background-color: #111111 !important;
    border: none !important;
    padding: 20px !important;
}

/* Button styling */
.gr-button {
    background-color: #222222 !important;
    color: white !important;
    border: 1px solid #444444 !important;
    margin: 5px 0 !important;
    border-radius: 8px !important;
    transition: all 0.3s ease !important;
}

.gr-button:hover {
    background-color: #333333 !important;
    border-color: #666666 !important;
}

/* Plot container */
.plot-container {
    background-color: #000000 !important;
    border: none !important;
}

/* Text elements */
h1, h2, h3, p, .markdown {
    color: white !important;
}

/* Remove all borders globally */
* {
    border-color: transparent !important;
}

/* Main content area */
.main-content {
    background-color: #000000 !important;
    padding: 20px !important;
}
"""

# Create the Gradio interface with sidebar and dark theme
with gr.Blocks(title="Model Test Results Dashboard", css=dark_theme_css) as demo:
    
    with gr.Row():
        # Sidebar for model selection
        with gr.Column(scale=1, elem_classes=["sidebar"]):
            gr.Markdown("# 🎯 Models")
            gr.Markdown("Select a model to view test results")
            
            # Model selection buttons in sidebar
            model_buttons = []
            for model_name in MODELS.keys():
                btn = gr.Button(
                    f"📊 {model_name.upper()}", 
                    variant="secondary",
                    size="lg",
                    elem_classes=["model-button"]
                )
                model_buttons.append(btn)
        
        # Main content area
        with gr.Column(scale=4, elem_classes=["main-content"]):
            gr.Markdown("# 📈 Test Results Dashboard")
            
            # Create the plot output
            plot_output = gr.Plot(
                label="", 
                format="png",
                elem_classes=["plot-container"]
            )
    
    # Set up click handlers for each button
    for i, (model_name, button) in enumerate(zip(MODELS.keys(), model_buttons)):
        button.click(
            fn=lambda name=model_name: plot_model_stats(name),
            outputs=plot_output
        )
    
    # Initialize with the first model
    demo.load(
        fn=lambda: plot_model_stats(list(MODELS.keys())[0]),
        outputs=plot_output
    )

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