model_page.py

import matplotlib.pyplot as plt
import pandas as pd
from utils import generate_underlined_line, COLORS
from data import extract_model_data, find_failure_first_seen

# Figure dimensions
FIGURE_WIDTH_DUAL = 18
FIGURE_HEIGHT_DUAL = 9

# Colors imported from utils

# Styling constants
BLACK = '#000000'
LABEL_COLOR = '#AAAAAA'
TITLE_COLOR = '#FFFFFF'

# Font sizes
DEVICE_TITLE_FONT_SIZE = 28

# Layout constants
SEPARATOR_LINE_Y_END = 0.85
SUBPLOT_TOP = 0.85
SUBPLOT_WSPACE = 0.4
PIE_START_ANGLE = 90
BORDER_LINE_WIDTH = 0.5
SEPARATOR_ALPHA = 0.5
SEPARATOR_LINE_WIDTH = 1
DEVICE_TITLE_PAD = 2
MODEL_TITLE_Y = 1

# Processing constants
MAX_FAILURE_ITEMS = 10


def _create_pie_chart(ax: plt.Axes, device_label: str, filtered_stats: dict) -> None:
    """Create a pie chart for device statistics."""
    if not filtered_stats:
        ax.text(0.5, 0.5, 'No test results',
                horizontalalignment='center', verticalalignment='center',
                transform=ax.transAxes, fontsize=14, color='#888888',
                fontfamily='monospace', weight='normal')
        ax.set_title(device_label, fontsize=DEVICE_TITLE_FONT_SIZE, weight='bold',
                     pad=DEVICE_TITLE_PAD, color=TITLE_COLOR, fontfamily='monospace')
        ax.axis('off')
        return

    chart_colors = [COLORS[category] for category in filtered_stats.keys()]

    # Create minimal pie chart - full pie, no donut effect
    wedges, texts, autotexts = ax.pie(
        filtered_stats.values(),
        labels=[label.lower() for label in filtered_stats.keys()],  # Lowercase for minimal look
        colors=chart_colors,
        autopct=lambda pct: f'{round(pct * sum(filtered_stats.values()) / 100)}',
        startangle=PIE_START_ANGLE,
        explode=None,  # No separation
        shadow=False,
        wedgeprops=dict(edgecolor='#1a1a1a', linewidth=BORDER_LINE_WIDTH),  # Minimal borders
        textprops={'fontsize': 12, 'weight': 'normal',
                   'color': LABEL_COLOR, 'fontfamily': 'monospace'}
    )

    # Enhanced percentage text styling for better readability
    for autotext in autotexts:
        autotext.set_color(BLACK)  # Black text for better contrast
        autotext.set_weight('bold')
        autotext.set_fontsize(14)
        autotext.set_fontfamily('monospace')

    # Minimal category labels
    for text in texts:
        text.set_color(LABEL_COLOR)
        text.set_weight('normal')
        text.set_fontsize(13)
        text.set_fontfamily('monospace')

    # Device label closer to chart and bigger
    ax.set_title(device_label, fontsize=DEVICE_TITLE_FONT_SIZE, weight='normal',
                 pad=DEVICE_TITLE_PAD, color=TITLE_COLOR, fontfamily='monospace')


def plot_model_stats(df: pd.DataFrame, model_name: str, historical_df: pd.DataFrame = None) -> tuple[plt.Figure, str, str]:
    """Draws pie charts of model's passed, failed, skipped, and error stats for AMD and NVIDIA."""
    # Handle case where the dataframe is empty or the model name could not be found in it
    if df.empty or model_name not in df.index:
        # Create empty stats for both devices
        amd_filtered = {}
        nvidia_filtered = {}
        failures_amd = failures_nvidia = {}
    else:
        row = df.loc[model_name]

        # Extract and process model data
        amd_stats, nvidia_stats = extract_model_data(row)[:2]

        # Filter out categories with 0 values for cleaner visualization
        amd_filtered = {k: v for k, v in amd_stats.items() if v > 0}
        nvidia_filtered = {k: v for k, v in nvidia_stats.items() if v > 0}

        # Generate failure info directly from dataframe
        failures_amd = row.get('failures_amd', None)
        failures_amd = {} if (failures_amd is None or pd.isna(failures_amd)) else dict(failures_amd)
        failures_nvidia = row.get('failures_nvidia')
        failures_nvidia = {} if (failures_nvidia is None or pd.isna(failures_nvidia)) else dict(failures_nvidia)

    # failure_xxx = {"single": [test, ...], "multi": [...]}
    # test = {"line": test_name. "trace": error_msg}

    # Always create figure with two subplots side by side with padding
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(FIGURE_WIDTH_DUAL, FIGURE_HEIGHT_DUAL), facecolor=BLACK)
    ax1.set_facecolor(BLACK)
    ax2.set_facecolor(BLACK)

    # Create both pie charts with device labels
    _create_pie_chart(ax1, "amd", amd_filtered)
    _create_pie_chart(ax2, "nvidia", nvidia_filtered)

    # Add subtle separation line between charts - stops at device labels level
    line_x = 0.5
    fig.add_artist(plt.Line2D([line_x, line_x], [0.0, SEPARATOR_LINE_Y_END],
                             color='#333333', linewidth=SEPARATOR_LINE_WIDTH,
                             alpha=SEPARATOR_ALPHA, transform=fig.transFigure))

    # Add central shared title for model name
    fig.suptitle(f'{model_name.lower()}', fontsize=32, weight='bold',
                 color='#CCCCCC', fontfamily='monospace', y=MODEL_TITLE_Y)

    # Clean layout with padding and space for central title
    plt.tight_layout()
    plt.subplots_adjust(top=SUBPLOT_TOP, wspace=SUBPLOT_WSPACE)

    amd_failed_info = prepare_textbox_content(failures_amd, 'AMD', bool(amd_filtered), model_name, historical_df)
    nvidia_failed_info = prepare_textbox_content(failures_nvidia, 'NVIDIA', bool(nvidia_filtered), model_name, historical_df)

    return fig, amd_failed_info, nvidia_failed_info


def prepare_textbox_content(failures: dict[str, list], device: str, data_available: bool, model_name: str = None, historical_df: pd.DataFrame = None) -> str:
    """Extract failure information from failures object with first seen dates."""
    # Catch the case where there is no data
    if not data_available:
        return generate_underlined_line(f"No data for {device}")
    # Catch the case where there are no failures
    if not failures:
        return generate_underlined_line(f"No failures for {device}")

    # Summary of failures
    single_failures = failures.get("single", [])
    multi_failures = failures.get("multi", [])
    info_lines = [
        generate_underlined_line(f"Failure summary for {device}:"),
        f"Single GPU failures: {len(single_failures)}",
        f"Multi GPU failures: {len(multi_failures)}",
        ""
    ]

    # Helper function to format failure line with first seen date
    def format_failure_line(test: dict, gpu_type: str) -> str:
        full_name = test.get("line", "::*could not find name*")
        short_name = full_name.split("::")[-1]
        
        # Try to find first seen date if historical data is available
        if historical_df is not None and model_name is not None and not historical_df.empty:
            first_seen = find_failure_first_seen(
                historical_df, 
                model_name, 
                full_name, 
                device.lower(), 
                gpu_type
            )
            if first_seen:
                # Format date as MM-DD-YYYY
                try:
                    from datetime import datetime
                    date_obj = datetime.strptime(first_seen, "%Y-%m-%d")
                    formatted_date = date_obj.strftime("%m-%d-%Y")
                    return f"{short_name} (First seen: {formatted_date})"
                except:
                    return f"{short_name} (First seen: {first_seen})"
        
        return short_name

    # Add single-gpu failures
    if single_failures:
        info_lines.append(generate_underlined_line("Single GPU failures:"))
        for test in single_failures:
            info_lines.append(format_failure_line(test, "single"))
        info_lines.append("\n")

    # Add multi-gpu failures
    if multi_failures:
        info_lines.append(generate_underlined_line("Multi GPU failures:"))
        for test in multi_failures:
            info_lines.append(format_failure_line(test, "multi"))

    return "\n".join(info_lines)