summary_page.py

import matplotlib.pyplot as plt
import pandas as pd
from matplotlib.offsetbox import OffsetImage, AnnotationBbox
from matplotlib.patches import FancyBboxPatch
import matplotlib.image as mpimg
import os
from data import extract_model_data

# Layout parameters
COLUMNS = 3

# Derived constants
COLUMN_WIDTH = 100 / COLUMNS  
BAR_WIDTH = COLUMN_WIDTH * 0.8
BAR_MARGIN = COLUMN_WIDTH * 0.1

# Figure dimensions
FIGURE_WIDTH = 22
MAX_HEIGHT = 14
MIN_HEIGHT_PER_ROW = 2.8
FIGURE_PADDING = 1

# Bar styling
BAR_HEIGHT_RATIO = 0.22
VERTICAL_SPACING_RATIO = 0.2
AMD_BAR_OFFSET = 0.25
NVIDIA_BAR_OFFSET = 0.54

# Colors
COLORS = {
    'passed': '#4CAF50',
    'failed': '#E53E3E',
    'skipped': '#FFD54F',
    'error': '#8B0000',
    'empty': "#5B5B5B"
}

# Font styling
MODEL_NAME_FONT_SIZE = 16
LABEL_FONT_SIZE = 14
LABEL_OFFSET = 1 
FAILURE_RATE_FONT_SIZE = 28

# Logo settings
LOGO_BOX_WIDTH = 4.5
LOGO_BOX_HEIGHT = 0.43
LOGO_ZOOM = 0.09

# Load logos once at module level
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
try:
    AMD_LOGO = mpimg.imread(os.path.join(SCRIPT_DIR, 'logos/amd_logo.png'))
except:
    AMD_LOGO = None
try:
    NVIDIA_LOGO = mpimg.imread(os.path.join(SCRIPT_DIR, 'logos/nvidia_logo.png'))
except:
    NVIDIA_LOGO = None


def calculate_overall_failure_rates(df: pd.DataFrame, available_models: list[str]) -> tuple[float, float]:
    """Calculate overall failure rates for AMD and NVIDIA across all models."""
    if df.empty or not available_models:
        return 0.0, 0.0
    
    total_amd_tests = 0
    total_amd_failures = 0
    total_nvidia_tests = 0
    total_nvidia_failures = 0
    
    for model_name in available_models:
        if model_name not in df.index:
            continue
            
        row = df.loc[model_name]
        amd_stats, nvidia_stats = extract_model_data(row)[:2]
        
        # AMD totals
        amd_total = amd_stats['passed'] + amd_stats['failed'] + amd_stats['error']
        if amd_total > 0:
            total_amd_tests += amd_total
            total_amd_failures += amd_stats['failed'] + amd_stats['error']
        
        # NVIDIA totals
        nvidia_total = nvidia_stats['passed'] + nvidia_stats['failed'] + nvidia_stats['error']
        if nvidia_total > 0:
            total_nvidia_tests += nvidia_total
            total_nvidia_failures += nvidia_stats['failed'] + nvidia_stats['error']
    
    amd_failure_rate = (total_amd_failures / total_amd_tests * 100) if total_amd_tests > 0 else 0.0
    nvidia_failure_rate = (total_nvidia_failures / total_nvidia_tests * 100) if total_nvidia_tests > 0 else 0.0
    
    return amd_failure_rate, nvidia_failure_rate


def draw_logo_and_bar(
    label: str,
    stats: dict[str, int],
    y_bar: float,
    column_left_position: float,
    bar_height: float,
    ax: plt.Axes,
) -> None:
    """Draw a horizontal bar chart for given stats with a logo box on the left."""
    # Determine if there are failures
    failures_present = any(stats[category] > 0 for category in ['failed', 'error'])
    
    # Select the appropriate logo
    logo = AMD_LOGO if label.lower() == "amd" else NVIDIA_LOGO
    
    # Calculate box position (centered on the bar vertically)
    box_x = column_left_position - LABEL_OFFSET - LOGO_BOX_WIDTH
    box_y = y_bar - LOGO_BOX_HEIGHT / 2
    
    # Draw the colored box
    if failures_present:
        box_color = COLORS['failed']  # Red for failures
        box_alpha = 0.6
    else:
        box_color = '#2a2a2a'  # Dark gray for no failures
        box_alpha = 0.5
    
    box = FancyBboxPatch(
        (box_x, box_y), 
        LOGO_BOX_WIDTH, 
        LOGO_BOX_HEIGHT,
        boxstyle="round,pad=0.05",
        facecolor=box_color,
        edgecolor='#444444',
        linewidth=1,
        alpha=box_alpha
    )
    ax.add_patch(box)
    
    # Add logo image inside the box if available
    if logo is not None:
        try:
            imagebox = OffsetImage(logo, zoom=LOGO_ZOOM)
            ab = AnnotationBbox(
                imagebox, 
                (box_x + LOGO_BOX_WIDTH / 2, y_bar),
                frameon=False,
                box_alignment=(0.5, 0.5)
            )
            ax.add_artist(ab)
        except:
            # Fallback to text if logo doesn't work
            ax.text(
                box_x + LOGO_BOX_WIDTH / 2, y_bar, 
                label.upper(), 
                ha='center', va='center', 
                color='#FFFFFF', 
                fontsize=10,
                fontfamily='monospace', 
                fontweight='bold'
            )
    else:
        # Fallback to text if logo not loaded
        ax.text(
            box_x + LOGO_BOX_WIDTH / 2, y_bar, 
            label.upper(), 
            ha='center', va='center', 
            color='#FFFFFF', 
            fontsize=10,
            fontfamily='monospace', 
            fontweight='bold'
        )
    
    # Draw the bar
    total = sum(stats.values())
    if total > 0:
        left = column_left_position
        for category in ['passed', 'failed', 'skipped', 'error']:
            if stats[category] > 0:
                width = stats[category] / total * BAR_WIDTH
                ax.barh(y_bar, width, left=left, height=bar_height, color=COLORS[category], alpha=0.9)
                left += width
    else:
        ax.barh(y_bar, BAR_WIDTH, left=column_left_position, height=bar_height, color=COLORS['empty'], alpha=0.9)

def create_summary_page(df: pd.DataFrame, available_models: list[str]) -> plt.Figure:
    """Create a summary page with model names and both AMD/NVIDIA test stats bars."""
    if df.empty:
        fig, ax = plt.subplots(figsize=(16, 8), facecolor='#000000')
        ax.set_facecolor('#000000')
        ax.text(0.5, 0.5, 'No data available',
                horizontalalignment='center', verticalalignment='center',
                transform=ax.transAxes, fontsize=20, color='#888888',
                fontfamily='monospace', weight='normal')
        ax.axis('off')
        return fig
    
    # Calculate overall failure rates
    amd_failure_rate, nvidia_failure_rate = calculate_overall_failure_rates(df, available_models)

    # Calculate dimensions for N-column layout
    model_count = len(available_models)
    rows = (model_count + COLUMNS - 1) // COLUMNS  # Ceiling division

    # Figure dimensions - wider for columns, height based on rows
    height_per_row = min(MIN_HEIGHT_PER_ROW, MAX_HEIGHT / max(rows, 1))
    figure_height = min(MAX_HEIGHT, rows * height_per_row + FIGURE_PADDING)

    fig, ax = plt.subplots(figsize=(FIGURE_WIDTH, figure_height), facecolor='#000000')
    ax.set_facecolor('#000000')
    
    # Add overall failure rates at the top as a proper title
    failure_text = f"Overall Failure Rates: AMD {amd_failure_rate:.1f}%  |  NVIDIA {nvidia_failure_rate:.1f}%"
    ax.text(50, -1.25, failure_text, ha='center', va='top', 
           color='#FFFFFF', fontsize=FAILURE_RATE_FONT_SIZE,
           fontfamily='monospace', fontweight='bold')

    visible_model_count = 0
    max_y = 0
    
    # Initialize counters for total tests
    amd_totals = {'passed': 0, 'failed': 0, 'skipped': 0}
    nvidia_totals = {'passed': 0, 'failed': 0, 'skipped': 0}

    for i, model_name in enumerate(available_models):
        if model_name not in df.index:
            continue

        row = df.loc[model_name]

        # Extract and process model data
        amd_stats, nvidia_stats = extract_model_data(row)[:2]
        
        # Accumulate totals
        amd_totals['passed'] += amd_stats['passed']
        amd_totals['failed'] += amd_stats['failed'] + amd_stats['error']
        amd_totals['skipped'] += amd_stats['skipped']
        
        nvidia_totals['passed'] += nvidia_stats['passed']
        nvidia_totals['failed'] += nvidia_stats['failed'] + nvidia_stats['error']
        nvidia_totals['skipped'] += nvidia_stats['skipped']

        # Calculate position in 4-column grid
        col = visible_model_count % COLUMNS
        row = visible_model_count // COLUMNS

        # Calculate horizontal position for this column
        col_left = col * COLUMN_WIDTH + BAR_MARGIN
        col_center = col * COLUMN_WIDTH + COLUMN_WIDTH / 2

        # Calculate vertical position for this row - start from top
        vertical_spacing = height_per_row
        y_base = (VERTICAL_SPACING_RATIO + row) * vertical_spacing
        y_model_name = y_base    # Model name above AMD bar
        y_amd_bar = y_base + vertical_spacing * AMD_BAR_OFFSET       # AMD bar
        y_nvidia_bar = y_base + vertical_spacing * NVIDIA_BAR_OFFSET    # NVIDIA bar
        max_y = max(max_y, y_nvidia_bar + vertical_spacing * 0.3)

        # Model name centered above the bars in this column
        ax.text(col_center, y_model_name, model_name.lower(),
               ha='center', va='center', color='#FFFFFF',
               fontsize=MODEL_NAME_FONT_SIZE, fontfamily='monospace', fontweight='bold')

        # AMD label and bar in this column
        bar_height = min(0.4, vertical_spacing * BAR_HEIGHT_RATIO)
        # Draw AMD bar with logo
        draw_logo_and_bar("amd", amd_stats, y_amd_bar, col_left, bar_height, ax)
        # Draw NVIDIA bar with logo
        draw_logo_and_bar("nvidia", nvidia_stats, y_nvidia_bar, col_left, bar_height, ax)

        # Increment counter for next visible model
        visible_model_count += 1


    # Add legend horizontally in bottom right corner
    patch_height = 0.3
    patch_width = 3

    legend_start_x = 68.7
    legend_y = max_y + 1
    legend_spacing = 10
    legend_font_size = 15
    
    # Add AMD and NVIDIA test totals in the bottom left
    # Calculate line spacing to align middle with legend
    line_height = 0.4  # Height between lines
    
    # Position the two lines so their middle aligns with legend_y
    amd_y = legend_y - line_height / 2
    nvidia_y = legend_y + line_height / 2
    
    amd_totals_text = f"AMD Tests    - Passed: {amd_totals['passed']}, Failed: {amd_totals['failed']}, Skipped: {amd_totals['skipped']}"
    nvidia_totals_text = f"NVIDIA Tests - Passed: {nvidia_totals['passed']}, Failed: {nvidia_totals['failed']}, Skipped: {nvidia_totals['skipped']}"
    
    ax.text(0, amd_y, amd_totals_text,
           ha='left', va='bottom', color='#CCCCCC',
           fontsize=14, fontfamily='monospace')
    
    ax.text(0, nvidia_y, nvidia_totals_text,
           ha='left', va='bottom', color='#CCCCCC',
           fontsize=14, fontfamily='monospace')
    
    # Legend entries
    legend_items = [
        ('passed', 'Passed'),
        ('failed', 'Failed'), 
        ('skipped', 'Skipped'),
    ]
    
    for i, (status, label) in enumerate(legend_items):
        x_pos = legend_start_x + i * legend_spacing
        # Small colored square
        ax.add_patch(plt.Rectangle((x_pos - 0.6, legend_y), patch_width, -patch_height,
                                 facecolor=COLORS[status], alpha=0.9))
        # Status label
        ax.text(x_pos + patch_width, legend_y, label, 
               ha='left', va='bottom', color='#CCCCCC', 
               fontsize=legend_font_size, fontfamily='monospace')

    # Style the axes to be completely invisible and span full width
    ax.set_xlim(-5, 105)  # Slightly wider to accommodate labels
    ax.set_ylim(0, max_y + 1)  # Add some padding at the top for title
    ax.set_xlabel('')
    ax.set_ylabel('')
    ax.spines['bottom'].set_visible(False)
    ax.spines['left'].set_visible(False)
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.set_xticks([])
    ax.set_yticks([])
    ax.yaxis.set_inverted(True)

    # Remove all margins to make figure stick to top
    plt.tight_layout()
    return fig