#!/usr/bin/env python3
"""
Create a visualization showing performance across Python versions
Minimal aesthetic matching the site design
"""

import json
import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np

# Configure matplotlib for beautiful output
try:
    import matplotlib.font_manager as fm
    fonts = fm.findSystemFonts(fontpaths=None, fontext='ttf')
    noto_fonts = [f for f in fonts if 'NotoSansMono' in f or 'Noto Sans Mono' in f]
    if noto_fonts:
        fm.fontManager.addfont(noto_fonts[0])
        mpl.rcParams['font.family'] = 'Noto Sans Mono'
    else:
        mpl.rcParams['font.family'] = 'monospace'
except:
    mpl.rcParams['font.family'] = 'monospace'

mpl.rcParams['font.size'] = 10
mpl.rcParams['axes.linewidth'] = 0.5
mpl.rcParams['xtick.major.width'] = 0.5
mpl.rcParams['ytick.major.width'] = 0.5

def load_summary():
    """Load the summary data"""
    with open('summary.json', 'r') as f:
        return json.load(f)

def create_version_visualization():
    """Create the Python version comparison visualization"""
    data = load_summary()
    
    # Extract Python versions and sort them
    versions = sorted(data['python_versions'])
    version_labels = [f"{v.split('.')[0]}.{v.split('.')[1]}" for v in versions]
    
    # Prepare data for each operation
    operations = ['creation_per_object', 'modification_per_operation', 'serialization_per_object']
    operation_labels = ['Object Creation', 'Modification', 'Serialization']
    
    # Colors
    color_dataclass = '#e0e0e0'
    color_pydantic = '#404040'
    
    # Create figure with 3 subplots
    fig, axes = plt.subplots(3, 1, figsize=(10, 10), gridspec_kw={'hspace': 0.4})
    fig.subplots_adjust(top=0.90)  # Make room for title and legend
    
    for idx, (op, op_label) in enumerate(zip(operations, operation_labels)):
        ax = axes[idx]
        
        # Extract data for this operation
        dataclass_times = []
        pydantic_times = []
        
        for version in versions:
            # Get absolute times in microseconds
            dc_time = data['raw_results'][version]['benchmarks']['dataclass'][op] * 1_000_000
            py_time = data['raw_results'][version]['benchmarks']['pydantic'][op] * 1_000_000
            
            dataclass_times.append(dc_time)
            pydantic_times.append(py_time)
        
        # Create grouped bar chart
        x = np.arange(len(versions))
        width = 0.35
        
        bars1 = ax.bar(x - width/2, dataclass_times, width, 
                       label='dataclasses', color=color_dataclass, edgecolor='black', linewidth=0.5)
        bars2 = ax.bar(x + width/2, pydantic_times, width,
                       label='pydantic', color=color_pydantic, edgecolor='black', linewidth=0.5)
        
        # Customize subplot
        ax.set_ylabel('time (μs)', fontsize=10)
        ax.set_title(op_label, fontsize=11, pad=10)
        ax.set_xticks(x)
        ax.set_xticklabels(version_labels)
        if idx == 2:  # Only add xlabel to bottom subplot
            ax.set_xlabel('Python version', fontsize=10)
        
        # Remove top and right spines
        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)
        
        # Set y-limits with proper padding
        max_val = max(max(dataclass_times), max(pydantic_times))
        ax.set_ylim(0, max_val * 1.2)  # 20% padding
        
        # Add grid for better readability
        ax.grid(axis='y', alpha=0.3, linestyle=':', linewidth=0.5)
        ax.set_axisbelow(True)
        
        # Highlight interesting patterns
        if op == 'serialization_per_object':
            # Add annotation about pydantic's advantage
            # Position it more carefully
            mid_idx = len(versions) // 2
            y_pos = max(max(dataclass_times), max(pydantic_times)) * 0.9
            ax.annotate('pydantic faster', 
                       xy=(mid_idx, pydantic_times[mid_idx]), 
                       xytext=(mid_idx, y_pos),
                       arrowprops=dict(arrowstyle='->', color='black', lw=0.5),
                       fontsize=8, ha='center')
    
    # Overall title
    fig.suptitle('dataclasses vs pydantic: performance across python versions', fontsize=12, y=0.96)
    
    # Add a single legend for the entire figure
    # Get handles and labels from the first subplot
    handles, labels = axes[0].get_legend_handles_labels()
    fig.legend(handles, labels, loc='upper center', bbox_to_anchor=(0.5, 0.92), 
              ncol=2, frameon=False, fontsize=10, columnspacing=3)
    
    # Add note at bottom
    fig.text(0.5, 0.02, 'lower is better • 10,000 iterations per test', 
             ha='center', fontsize=8, style='italic', color='#666666')
    
    # Save with high quality and transparent background
    plt.savefig('benchmark_versions.png', dpi=150, bbox_inches='tight', 
                facecolor='none', transparent=True, edgecolor='none')
    plt.savefig('benchmark_versions.svg', format='svg', bbox_inches='tight',
                facecolor='none', transparent=True, edgecolor='none')
    
    print("Version comparison saved as benchmark_versions.png and .svg")

if __name__ == "__main__":
    create_version_visualization()