time_series.py

import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
from datetime import datetime
from data import extract_model_data

# Colors matching the existing theme
COLORS = {
    'passed': '#4CAF50',
    'failed': '#E53E3E',
    'skipped': '#FFD54F',
    'error': '#8B0000'
}

# Figure dimensions
FIGURE_WIDTH = 20
FIGURE_HEIGHT = 12

# Styling constants
BLACK = '#000000'
LABEL_COLOR = '#CCCCCC'
TITLE_COLOR = '#FFFFFF'
GRID_COLOR = '#333333'

# Font sizes
TITLE_FONT_SIZE = 24
LABEL_FONT_SIZE = 14
LEGEND_FONT_SIZE = 12


def create_time_series_summary(historical_df: pd.DataFrame) -> plt.Figure:
    """Create time-series visualization for overall failure rates over time."""
    if historical_df.empty or 'date' not in historical_df.columns:
        fig, ax = plt.subplots(figsize=(FIGURE_WIDTH, FIGURE_HEIGHT), facecolor=BLACK)
        ax.set_facecolor(BLACK)
        ax.text(0.5, 0.5, 'No historical data available',
                horizontalalignment='center', verticalalignment='center',
                transform=ax.transAxes, fontsize=20, color='#888888',
                fontfamily='monospace', weight='normal')
        ax.axis('off')
        return fig
    
    # Convert date column to datetime
    historical_df['date_dt'] = pd.to_datetime(historical_df['date'])
    historical_df = historical_df.sort_values('date_dt')
    
    # Group by date and calculate overall statistics
    daily_stats = []
    dates = []
    
    for date in historical_df['date_dt'].unique():
        date_data = historical_df[historical_df['date_dt'] == date]
        
        total_amd_passed = 0
        total_amd_failed = 0
        total_amd_skipped = 0
        total_nvidia_passed = 0
        total_nvidia_failed = 0
        total_nvidia_skipped = 0
        
        for _, row in date_data.iterrows():
            amd_stats, nvidia_stats = extract_model_data(row)[:2]
            
            total_amd_passed += amd_stats['passed']
            total_amd_failed += amd_stats['failed']
            total_amd_skipped += amd_stats['skipped']
            
            total_nvidia_passed += nvidia_stats['passed']
            total_nvidia_failed += nvidia_stats['failed']
            total_nvidia_skipped += nvidia_stats['skipped']
        
        # Calculate failure rates
        amd_total = total_amd_passed + total_amd_failed
        nvidia_total = total_nvidia_passed + total_nvidia_failed
        
        amd_failure_rate = (total_amd_failed / amd_total * 100) if amd_total > 0 else 0
        nvidia_failure_rate = (total_nvidia_failed / nvidia_total * 100) if nvidia_total > 0 else 0
        
        daily_stats.append({
            'amd_failure_rate': amd_failure_rate,
            'nvidia_failure_rate': nvidia_failure_rate,
            'amd_passed': total_amd_passed,
            'amd_failed': total_amd_failed,
            'amd_skipped': total_amd_skipped,
            'nvidia_passed': total_nvidia_passed,
            'nvidia_failed': total_nvidia_failed,
            'nvidia_skipped': total_nvidia_skipped
        })
        dates.append(date)
    
    # Create the plot with 3 subplots: failure rates, AMD stacked, NVIDIA stacked
    fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(FIGURE_WIDTH, FIGURE_HEIGHT + 4), facecolor=BLACK)
    ax1.set_facecolor(BLACK)
    ax2.set_facecolor(BLACK)
    ax3.set_facecolor(BLACK)
    
    # Plot 1: Failure rates over time
    dates_array = np.array(dates)
    amd_rates = [stat['amd_failure_rate'] for stat in daily_stats]
    nvidia_rates = [stat['nvidia_failure_rate'] for stat in daily_stats]
    
    ax1.plot(dates_array, amd_rates, color='#FF6B6B', linewidth=3, label='AMD', marker='o', markersize=6)
    ax1.plot(dates_array, nvidia_rates, color='#4ECDC4', linewidth=3, label='NVIDIA', marker='s', markersize=6)
    
    ax1.set_title('Overall Failure Rates Over Time', fontsize=TITLE_FONT_SIZE, color=TITLE_COLOR, 
                  fontfamily='monospace', fontweight='bold', pad=20)
    ax1.set_ylabel('Failure Rate (%)', fontsize=LABEL_FONT_SIZE, color=LABEL_COLOR, fontfamily='monospace')
    ax1.grid(True, color=GRID_COLOR, alpha=0.3, linestyle='-', linewidth=0.5)
    ax1.legend(fontsize=LEGEND_FONT_SIZE, loc='upper right', frameon=False, 
               labelcolor=LABEL_COLOR, prop={'family': 'monospace'})
    
    # Format x-axis
    ax1.tick_params(colors=LABEL_COLOR, labelsize=LABEL_FONT_SIZE)
    ax1.xaxis.label.set_color(LABEL_COLOR)
    ax1.yaxis.label.set_color(LABEL_COLOR)
    
    # Plot 2: AMD Test counts over time (stacked area chart)
    amd_passed = [stat['amd_passed'] for stat in daily_stats]
    amd_failed = [stat['amd_failed'] for stat in daily_stats]
    amd_skipped = [stat['amd_skipped'] for stat in daily_stats]
    
    ax2.fill_between(dates_array, 0, amd_passed, color=COLORS['passed'], alpha=0.7, label='Passed')
    ax2.fill_between(dates_array, amd_passed, np.array(amd_passed) + np.array(amd_failed), 
                     color=COLORS['failed'], alpha=0.7, label='Failed')
    ax2.fill_between(dates_array, np.array(amd_passed) + np.array(amd_failed), 
                     np.array(amd_passed) + np.array(amd_failed) + np.array(amd_skipped),
                     color=COLORS['skipped'], alpha=0.7, label='Skipped')
    
    ax2.set_title('AMD Test Results Over Time', fontsize=TITLE_FONT_SIZE, color=TITLE_COLOR, 
                  fontfamily='monospace', fontweight='bold', pad=20)
    ax2.set_ylabel('Number of Tests', fontsize=LABEL_FONT_SIZE, color=LABEL_COLOR, fontfamily='monospace')
    ax2.grid(True, color=GRID_COLOR, alpha=0.3, linestyle='-', linewidth=0.5)
    ax2.legend(fontsize=LEGEND_FONT_SIZE, loc='upper right', frameon=False, 
               labelcolor=LABEL_COLOR, prop={'family': 'monospace'})
    
    # Format x-axis
    ax2.tick_params(colors=LABEL_COLOR, labelsize=LABEL_FONT_SIZE)
    ax2.xaxis.label.set_color(LABEL_COLOR)
    ax2.yaxis.label.set_color(LABEL_COLOR)
    
    # Plot 3: NVIDIA Test counts over time (stacked area chart)
    nvidia_passed = [stat['nvidia_passed'] for stat in daily_stats]
    nvidia_failed = [stat['nvidia_failed'] for stat in daily_stats]
    nvidia_skipped = [stat['nvidia_skipped'] for stat in daily_stats]
    
    ax3.fill_between(dates_array, 0, nvidia_passed, color=COLORS['passed'], alpha=0.7, label='Passed')
    ax3.fill_between(dates_array, nvidia_passed, np.array(nvidia_passed) + np.array(nvidia_failed), 
                     color=COLORS['failed'], alpha=0.7, label='Failed')
    ax3.fill_between(dates_array, np.array(nvidia_passed) + np.array(nvidia_failed), 
                     np.array(nvidia_passed) + np.array(nvidia_failed) + np.array(nvidia_skipped),
                     color=COLORS['skipped'], alpha=0.7, label='Skipped')
    
    ax3.set_title('NVIDIA Test Results Over Time', fontsize=TITLE_FONT_SIZE, color=TITLE_COLOR, 
                  fontfamily='monospace', fontweight='bold', pad=20)
    ax3.set_ylabel('Number of Tests', fontsize=LABEL_FONT_SIZE, color=LABEL_COLOR, fontfamily='monospace')
    ax3.set_xlabel('Date', fontsize=LABEL_FONT_SIZE, color=LABEL_COLOR, fontfamily='monospace')
    ax3.grid(True, color=GRID_COLOR, alpha=0.3, linestyle='-', linewidth=0.5)
    ax3.legend(fontsize=LEGEND_FONT_SIZE, loc='upper right', frameon=False, 
               labelcolor=LABEL_COLOR, prop={'family': 'monospace'})
    
    # Format x-axis
    ax3.tick_params(colors=LABEL_COLOR, labelsize=LABEL_FONT_SIZE)
    ax3.xaxis.label.set_color(LABEL_COLOR)
    ax3.yaxis.label.set_color(LABEL_COLOR)
    
    # Rotate x-axis labels for better readability
    for ax in [ax1, ax2, ax3]:
        ax.tick_params(axis='x', rotation=45)
    
    plt.tight_layout()
    
    # Close any existing figures to prevent memory issues
    plt.close('all')
    
    return fig


def create_model_time_series(historical_df: pd.DataFrame, model_name: str) -> plt.Figure:
    """Create time-series visualization for a specific model."""
    if historical_df.empty or 'date' not in historical_df.columns:
        fig, ax = plt.subplots(figsize=(FIGURE_WIDTH, FIGURE_HEIGHT), facecolor=BLACK)
        ax.set_facecolor(BLACK)
        ax.text(0.5, 0.5, f'No historical data available for {model_name}',
                horizontalalignment='center', verticalalignment='center',
                transform=ax.transAxes, fontsize=20, color='#888888',
                fontfamily='monospace', weight='normal')
        ax.axis('off')
        return fig
    
    # Filter data for the specific model
    model_data = historical_df[historical_df.index.str.lower() == model_name.lower()]
    
    if model_data.empty:
        fig, ax = plt.subplots(figsize=(FIGURE_WIDTH, FIGURE_HEIGHT), facecolor=BLACK)
        ax.set_facecolor(BLACK)
        ax.text(0.5, 0.5, f'No data found for model: {model_name}',
                horizontalalignment='center', verticalalignment='center',
                transform=ax.transAxes, fontsize=20, color='#888888',
                fontfamily='monospace', weight='normal')
        ax.axis('off')
        return fig
    
    # Convert date column to datetime and sort
    model_data = model_data.copy()
    model_data['date_dt'] = pd.to_datetime(model_data['date'])
    model_data = model_data.sort_values('date_dt')
    
    # Extract statistics for each date
    dates = model_data['date_dt'].values
    amd_stats_list = []
    nvidia_stats_list = []
    
    for _, row in model_data.iterrows():
        amd_stats, nvidia_stats = extract_model_data(row)[:2]
        amd_stats_list.append(amd_stats)
        nvidia_stats_list.append(nvidia_stats)
    
    # Create the plot
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(FIGURE_WIDTH, FIGURE_HEIGHT), facecolor=BLACK)
    ax1.set_facecolor(BLACK)
    ax2.set_facecolor(BLACK)
    
    # Plot 1: AMD results over time
    amd_passed = [stats['passed'] for stats in amd_stats_list]
    amd_failed = [stats['failed'] for stats in amd_stats_list]
    amd_skipped = [stats['skipped'] for stats in amd_stats_list]
    
    ax1.plot(dates, amd_passed, color=COLORS['passed'], linewidth=3, label='Passed', marker='o', markersize=6)
    ax1.plot(dates, amd_failed, color=COLORS['failed'], linewidth=3, label='Failed', marker='s', markersize=6)
    ax1.plot(dates, amd_skipped, color=COLORS['skipped'], linewidth=3, label='Skipped', marker='^', markersize=6)
    
    ax1.set_title(f'{model_name.upper()} - AMD Results Over Time', fontsize=TITLE_FONT_SIZE, color=TITLE_COLOR, 
                  fontfamily='monospace', fontweight='bold', pad=20)
    ax1.set_ylabel('Number of Tests', fontsize=LABEL_FONT_SIZE, color=LABEL_COLOR, fontfamily='monospace')
    ax1.grid(True, color=GRID_COLOR, alpha=0.3, linestyle='-', linewidth=0.5)
    ax1.legend(fontsize=LEGEND_FONT_SIZE, loc='upper right', frameon=False, 
               labelcolor=LABEL_COLOR, prop={'family': 'monospace'})
    
    # Plot 2: NVIDIA results over time
    nvidia_passed = [stats['passed'] for stats in nvidia_stats_list]
    nvidia_failed = [stats['failed'] for stats in nvidia_stats_list]
    nvidia_skipped = [stats['skipped'] for stats in nvidia_stats_list]
    
    ax2.plot(dates, nvidia_passed, color=COLORS['passed'], linewidth=3, label='Passed', marker='o', markersize=6)
    ax2.plot(dates, nvidia_failed, color=COLORS['failed'], linewidth=3, label='Failed', marker='s', markersize=6)
    ax2.plot(dates, nvidia_skipped, color=COLORS['skipped'], linewidth=3, label='Skipped', marker='^', markersize=6)
    
    ax2.set_title(f'{model_name.upper()} - NVIDIA Results Over Time', fontsize=TITLE_FONT_SIZE, color=TITLE_COLOR, 
                  fontfamily='monospace', fontweight='bold', pad=20)
    ax2.set_ylabel('Number of Tests', fontsize=LABEL_FONT_SIZE, color=LABEL_COLOR, fontfamily='monospace')
    ax2.set_xlabel('Date', fontsize=LABEL_FONT_SIZE, color=LABEL_COLOR, fontfamily='monospace')
    ax2.grid(True, color=GRID_COLOR, alpha=0.3, linestyle='-', linewidth=0.5)
    ax2.legend(fontsize=LEGEND_FONT_SIZE, loc='upper right', frameon=False, 
               labelcolor=LABEL_COLOR, prop={'family': 'monospace'})
    
    # Format axes
    for ax in [ax1, ax2]:
        ax.tick_params(colors=LABEL_COLOR, labelsize=LABEL_FONT_SIZE)
        ax.xaxis.label.set_color(LABEL_COLOR)
        ax.yaxis.label.set_color(LABEL_COLOR)
        ax.tick_params(axis='x', rotation=45)
    
    plt.tight_layout()
    
    # Close any existing figures to prevent memory issues
    plt.close('all')
    
    return fig