app.py

import matplotlib.pyplot as plt
import matplotlib
import pandas as pd
import gradio as gr

from data import CIResults
from utils import logger
from summary_page import create_summary_page
from model_page import plot_model_stats
from time_series_gradio import create_time_series_summary_gradio, create_model_time_series_gradio


# Configure matplotlib to prevent memory warnings and set dark background
matplotlib.rcParams['figure.facecolor'] = '#000000'
matplotlib.rcParams['axes.facecolor'] = '#000000'
matplotlib.rcParams['savefig.facecolor'] = '#000000'
plt.ioff()  # Turn off interactive mode to prevent figure accumulation


# Load data once at startup
Ci_results = CIResults()
Ci_results.load_data()
# Start the auto-reload scheduler
Ci_results.schedule_data_reload()


# Function to check if a model has failures
def model_has_failures(model_name):
    """Check if a model has any failures (AMD or NVIDIA)."""
    if Ci_results.df is None or Ci_results.df.empty:
        return False
    
    # Normalize model name to match DataFrame index
    model_name_lower = model_name.lower()
    
    # Check if model exists in DataFrame
    if model_name_lower not in Ci_results.df.index:
        return False
    row = Ci_results.df.loc[model_name_lower]
    
    # Check for failures in both AMD and NVIDIA
    amd_multi_failures = row.get('failed_multi_no_amd', 0)
    amd_single_failures = row.get('failed_single_no_amd', 0)
    nvidia_multi_failures = row.get('failed_multi_no_nvidia', 0)
    nvidia_single_failures = row.get('failed_single_no_nvidia', 0)
    return any([
        amd_multi_failures > 0,
        amd_single_failures > 0,
        nvidia_multi_failures > 0,
        nvidia_single_failures > 0,
    ])


# Function to get current description text
def get_description_text():
    """Get description text with integrated last update time."""
    msg = [
        "Transformer CI Dashboard",
        "-",
        "AMD runs on MI325",
        "NVIDIA runs on A10",
    ]
    msg = ["**" + x + "**" for x in msg] + [""]
    if Ci_results.latest_update_msg:
        msg.append(f"*({Ci_results.latest_update_msg})*")
    else:
        msg.append("*(loading...)*")
    return "<br>".join(msg)

# Load CSS from external file
def load_css():
    try:
        with open("styles.css", "r") as f:
            css_content = f.read()
        
        return css_content
    except FileNotFoundError:
        logger.warning("styles.css not found, using minimal default styles")
        return "body { background: #000; color: #fff; }"


# Create the Gradio interface with sidebar and dark theme
with gr.Blocks(title="Model Test Results Dashboard", css=load_css()) as demo:


    with gr.Row():
        # Sidebar for model selection
        with gr.Column(scale=1, elem_classes=["sidebar"]):
            gr.Markdown("# 🤖 TCID", elem_classes=["sidebar-title"])

            # Description with integrated last update time
            description_text = get_description_text()
            description_display = gr.Markdown(description_text, elem_classes=["sidebar-description"])

            # View toggle buttons
            with gr.Row(elem_classes=["view-toggle-row"]):
                current_view_button = gr.Button(
                    "current\n📊",
                    variant="primary",
                    size="lg",
                    elem_classes=["view-toggle-button", "view-toggle-active"]
                )
                historical_view_button = gr.Button(
                    "history\n📈",
                    variant="secondary",
                    size="lg",
                    elem_classes=["view-toggle-button"]
                )

            # Date selection toggle button (initially hidden)
            date_toggle_button = gr.Button(
                "► Date Selection",
                variant="secondary",
                elem_classes=["date-header"],
                visible=False
            )

            # Date selection container (collapsible) - start folded
            with gr.Column(visible=True, elem_classes=["date-selection", "date-selection-hidden"]) as date_selection:
                start_date = gr.Dropdown(
                    choices=Ci_results.available_dates,
                    value=Ci_results.available_dates[-1] if Ci_results.available_dates else None,  # Last date (oldest)
                    label="Start Date",
                    elem_classes=["date-dropdown"]
                )
                end_date = gr.Dropdown(
                    choices=Ci_results.available_dates,
                    value=Ci_results.available_dates[0] if Ci_results.available_dates else None,  # First date (newest)
                    label="End Date",
                    elem_classes=["date-dropdown"]
                )
                load_historical_button = gr.Button(
                    "Reload Historical Data",
                    variant="primary",
                    size="sm",
                    elem_classes=["load-historical-button"]
                )

            # Summary button (for current view)
            summary_button = gr.Button(
                "summary\n📊",
                variant="primary",
                size="lg",
                elem_classes=["summary-button"]
            )

            # Model selection header (clickable toggle)
            model_toggle_button = gr.Button(
                f"► Select model ({len(Ci_results.available_models)})",
                variant="secondary",
                elem_classes=["model-header"]
            )

            # Model buttons container (collapsible) - start folded
            with gr.Column(elem_classes=["model-list", "model-list-hidden"]) as model_list_container:
                # Create individual buttons for each model
                model_buttons = []
                model_choices = [model.lower() for model in Ci_results.available_models] if Ci_results.available_models else ["auto", "bert", "clip", "llama"]
                
                print(f"Creating {len(model_choices)} model buttons: {model_choices}")

                for model_name in model_choices:
                    # Check if model has failures to determine styling
                    has_failures = model_has_failures(model_name)
                    button_classes = ["model-button"]
                    if has_failures:
                        button_classes.append("model-button-failed")
                    
                    btn = gr.Button(
                        model_name,
                        variant="secondary",
                        size="sm",
                        elem_classes=button_classes
                    )
                    model_buttons.append(btn)

            # CI job links at bottom of sidebar
            ci_links_display = gr.Markdown("🔗 **CI Jobs:** *Loading...*", elem_classes=["sidebar-links"])

        # Main content area
        with gr.Column(scale=4, elem_classes=["main-content"]):
            # Current view components
            with gr.Column(visible=True, elem_classes=["current-view"]) as current_view:
                # Summary display (default view)
                summary_display = gr.Plot(
                    value=create_summary_page(Ci_results.df, Ci_results.available_models),
                    label="",
                    format="png",
                    elem_classes=["plot-container"],
                    visible=True
                )

                # Detailed view components (hidden by default)
                with gr.Column(visible=False, elem_classes=["detail-view"]) as detail_view:
                    # Back button for current view detail
                    back_to_summary_current_button = gr.Button(
                        "← Back to Summary",
                        variant="secondary",
                        size="sm",
                        elem_classes=["back-button"]
                    )
                    
                    # Create the plot output
                    plot_output = gr.Plot(
                        label="",
                        format="png",
                        elem_classes=["plot-container"]
                    )

                    # Create two separate failed tests displays in a row layout
                    with gr.Row():
                        with gr.Column(scale=1):
                            amd_failed_tests_output = gr.Textbox(
                                value="",
                                lines=8,
                                max_lines=8,
                                interactive=False,
                                container=False,
                                elem_classes=["failed-tests"]
                            )
                        with gr.Column(scale=1):
                            nvidia_failed_tests_output = gr.Textbox(
                                value="",
                                lines=8,
                                max_lines=8,
                                interactive=False,
                                container=False,
                                elem_classes=["failed-tests"]
                            )

            # Historical view components (hidden by default)
            with gr.Column(visible=False, elem_classes=["historical-view"]) as historical_view:
                
                
                # Time-series summary displays (multiple Gradio plots)
                time_series_failure_rates = gr.LinePlot(
                    label="",
                    elem_classes=["plot-container"]
                )
                
                time_series_amd_tests = gr.LinePlot(
                    label="",
                    elem_classes=["plot-container"]
                )
                
                time_series_nvidia_tests = gr.LinePlot(
                    label="",
                    elem_classes=["plot-container"]
                )

                # Time-series model view (hidden by default)
                with gr.Column(visible=False, elem_classes=["time-series-detail-view"]) as time_series_detail_view:
                    # Back button for time-series model view
                    back_to_summary_button = gr.Button(
                        "← Back to Summary",
                        variant="secondary",
                        size="sm",
                        elem_classes=["back-button"]
                    )
                    
                    # Time-series plots for specific model (with spacing)
                    time_series_amd_model_plot = gr.LinePlot(
                        label="",
                        elem_classes=["plot-container"]
                    )
                    
                    time_series_nvidia_model_plot = gr.LinePlot(
                        label="",
                        elem_classes=["plot-container"]
                    )

    # Set up click handlers for model buttons
    for i, btn in enumerate(model_buttons):
        model_name = model_choices[i]
        btn.click(
            fn=lambda selected_model=model_name: plot_model_stats(Ci_results.df, selected_model),
            outputs=[plot_output, amd_failed_tests_output, nvidia_failed_tests_output]
        ).then(
            fn=lambda: [gr.update(visible=False), gr.update(visible=True)],
            outputs=[summary_display, detail_view]
        )

    # Model toggle functionality
    def toggle_model_list(current_visible):
        """Toggle the visibility of the model list."""
        new_visible = not current_visible
        arrow = "▼" if new_visible else "►"
        button_text = f"{arrow} Select model ({len(Ci_results.available_models)})"
        
        # Use CSS classes instead of Gradio visibility
        css_classes = ["model-list"]
        if new_visible:
            css_classes.append("model-list-visible")
        else:
            css_classes.append("model-list-hidden")
            
        return gr.update(value=button_text), gr.update(elem_classes=css_classes), new_visible

    # Track model list visibility state
    model_list_visible = gr.State(False)

    model_toggle_button.click(
        fn=toggle_model_list,
        inputs=[model_list_visible],
        outputs=[model_toggle_button, model_list_container, model_list_visible]
    )

    # Date toggle functionality
    def toggle_date_selection(current_visible):
        """Toggle the visibility of the date selection."""
        new_visible = not current_visible
        arrow = "▼" if new_visible else "►"
        button_text = f"{arrow} Date Selection"
        
        # Use CSS classes instead of Gradio visibility
        css_classes = ["date-selection"]
        if new_visible:
            css_classes.append("date-selection-visible")
        else:
            css_classes.append("date-selection-hidden")
            
        return gr.update(value=button_text), gr.update(elem_classes=css_classes), new_visible

    # Track date selection visibility state
    date_selection_visible = gr.State(False)

    date_toggle_button.click(
        fn=toggle_date_selection,
        inputs=[date_selection_visible],
        outputs=[date_toggle_button, date_selection, date_selection_visible]
    )

    # Summary button click handler
    def show_summary_and_update_links():
        """Show summary page and update CI links."""
        return create_summary_page(Ci_results.df, Ci_results.available_models), get_description_text(), get_ci_links()

    summary_button.click(
        fn=show_summary_and_update_links,
        outputs=[summary_display, description_display, ci_links_display]
    ).then(
        fn=lambda: [gr.update(visible=True), gr.update(visible=False)],
        outputs=[summary_display, detail_view]
    )

    # Function to get CI job links
    def get_ci_links():
        """Get CI job links from the most recent data."""
        try:
            # Check if df exists and is not empty
            if Ci_results.df is None or Ci_results.df.empty:
                return "🔗 **CI Jobs:** *Loading...*"

            # Get links from any available model (they should be the same for all models in a run)
            amd_multi_link = None
            amd_single_link = None
            nvidia_multi_link = None
            nvidia_single_link = None

            for model_name in Ci_results.df.index:
                row = Ci_results.df.loc[model_name]

                # Extract AMD links
                if pd.notna(row.get('job_link_amd')) and (not amd_multi_link or not amd_single_link):
                    amd_link_raw = row.get('job_link_amd')
                    if isinstance(amd_link_raw, dict):
                        if 'multi' in amd_link_raw and not amd_multi_link:
                            amd_multi_link = amd_link_raw['multi']
                        if 'single' in amd_link_raw and not amd_single_link:
                            amd_single_link = amd_link_raw['single']

                # Extract NVIDIA links
                if pd.notna(row.get('job_link_nvidia')) and (not nvidia_multi_link or not nvidia_single_link):
                    nvidia_link_raw = row.get('job_link_nvidia')
                    if isinstance(nvidia_link_raw, dict):
                        if 'multi' in nvidia_link_raw and not nvidia_multi_link:
                            nvidia_multi_link = nvidia_link_raw['multi']
                        if 'single' in nvidia_link_raw and not nvidia_single_link:
                            nvidia_single_link = nvidia_link_raw['single']

                # Break if we have all links
                if amd_multi_link and amd_single_link and nvidia_multi_link and nvidia_single_link:
                    break


            # Add FAQ link at the bottom
            links_md = "❓ [**FAQ**](https://huggingface.co/spaces/transformers-community/transformers-ci-dashboard/blob/main/README.md)\n\n"
            links_md += "🔗 **CI Jobs:**\n\n"

            # AMD links
            if amd_multi_link or amd_single_link:
                links_md += "**AMD:**\n"
                if amd_multi_link:
                    links_md += f"• [Multi GPU]({amd_multi_link})\n"
                if amd_single_link:
                    links_md += f"• [Single GPU]({amd_single_link})\n"
                links_md += "\n"

            # NVIDIA links
            if nvidia_multi_link or nvidia_single_link:
                links_md += "**NVIDIA:**\n"
                if nvidia_multi_link:
                    links_md += f"• [Multi GPU]({nvidia_multi_link})\n"
                if nvidia_single_link:
                    links_md += f"• [Single GPU]({nvidia_single_link})\n"

            if not (amd_multi_link or amd_single_link or nvidia_multi_link or nvidia_single_link):
                links_md += "*No links available*"

            return links_md
        except Exception as e:
            logger.error(f"getting CI links: {e}")
            return "🔗 **CI Jobs:** *Error loading links*\n\n❓ **[FAQ](README.md)**"


    # View toggle functionality
    def toggle_to_current_view():
        """Switch to current view."""
        return [
            gr.update(visible=True),   # current_view
            gr.update(visible=False),  # historical_view
            gr.update(visible=False),  # date_toggle_button
            gr.update(visible=True),   # summary_button
            gr.update(variant="primary", elem_classes=["view-toggle-button", "view-toggle-active"]),  # current_view_button
            gr.update(variant="secondary", elem_classes=["view-toggle-button"])  # historical_view_button
        ]

    def toggle_to_historical_view():
        """Switch to historical view first, then auto-load data."""
        # First, just switch the view
        return [
            gr.update(visible=False),  # current_view
            gr.update(visible=True),   # historical_view
            gr.update(visible=True),   # date_toggle_button
            gr.update(visible=False),  # summary_button
            gr.update(variant="secondary", elem_classes=["view-toggle-button"]),  # current_view_button
            gr.update(variant="primary", elem_classes=["view-toggle-button", "view-toggle-active"]),  # historical_view_button
            gr.update(),  # time_series_failure_rates
            gr.update(),  # time_series_amd_tests
            gr.update(),  # time_series_nvidia_tests
        ]

    def auto_load_historical_data():
        """Auto-load data for preselected dates after view switch."""
        # Get the preselected dates
        start_date_val = Ci_results.available_dates[-1] if Ci_results.available_dates else None
        end_date_val = Ci_results.available_dates[0] if Ci_results.available_dates else None
        
        # Check if we already have data for these dates
        if (hasattr(Ci_results, 'cached_start_date') and hasattr(Ci_results, 'cached_end_date') and
            Ci_results.cached_start_date == start_date_val and Ci_results.cached_end_date == end_date_val and
            not Ci_results.historical_df.empty):
            # Use cached data - no loading indicator needed
            plots = create_time_series_summary_gradio(Ci_results.historical_df)
            return plots['failure_rates'], plots['amd_tests'], plots['nvidia_tests']
        
        # Auto-load historical data if dates are available
        if start_date_val and end_date_val:
            try:
                Ci_results.load_historical_data(start_date_val, end_date_val)
                
                if not Ci_results.historical_df.empty:
                    # Cache the loaded data
                    Ci_results.cached_start_date = start_date_val
                    Ci_results.cached_end_date = end_date_val
                    
                    plots = create_time_series_summary_gradio(Ci_results.historical_df)
                    return plots['failure_rates'], plots['amd_tests'], plots['nvidia_tests']
                else:
                    return gr.update(), gr.update(), gr.update()
            except Exception as e:
                logger.error(f"Error auto-loading historical data: {e}")
                return gr.update(), gr.update(), gr.update()
        else:
            return gr.update(), gr.update(), gr.update()

    current_view_button.click(
        fn=toggle_to_current_view,
        outputs=[current_view, historical_view, date_toggle_button, summary_button, current_view_button, historical_view_button]
    )

    historical_view_button.click(
        fn=toggle_to_historical_view,
        outputs=[current_view, historical_view, date_toggle_button, summary_button, current_view_button, historical_view_button, time_series_failure_rates, time_series_amd_tests, time_series_nvidia_tests]
    ).then(
        fn=auto_load_historical_data,
        outputs=[time_series_failure_rates, time_series_amd_tests, time_series_nvidia_tests]
    )

    # Historical data loading functionality
    def load_historical_data(start_date, end_date):
        """Load and display historical data."""
        if not start_date or not end_date:
            logger.error("No start or end date provided")
            return gr.update(), gr.update(), gr.update()
        
        try:
            Ci_results.load_historical_data(start_date, end_date)
            
            if Ci_results.historical_df.empty:
                logger.error("No historical data found for the selected date range")
                return gr.update(), gr.update(), gr.update()
            
            # Create time-series summary plots
            plots = create_time_series_summary_gradio(Ci_results.historical_df)
            
            # Cache the loaded data
            Ci_results.cached_start_date = start_date
            Ci_results.cached_end_date = end_date
            
            return plots['failure_rates'], plots['amd_tests'], plots['nvidia_tests']
            
        except Exception as e:
            logger.error(f"Error loading historical data: {e}")
            return gr.update(), gr.update(), gr.update()

    load_historical_button.click(
        fn=load_historical_data,
        inputs=[start_date, end_date],
        outputs=[time_series_failure_rates, time_series_amd_tests, time_series_nvidia_tests]
    )

    # Time-series model selection functionality
    def show_time_series_model(selected_model):
        """Show time-series view for a specific model."""
        if Ci_results.historical_df.empty:
            return gr.update(), gr.update()
        
        try:
            plots = create_model_time_series_gradio(Ci_results.historical_df, selected_model)
            return plots['amd_plot'], plots['nvidia_plot']
        except Exception as e:
            logger.error(f"Error creating time-series for model {selected_model}: {e}")
            return gr.update(), gr.update()

    # Back button functionality
    def back_to_summary():
        """Return from model time-series view to summary time-series view."""
        return [
            gr.update(visible=True),   # time_series_failure_rates
            gr.update(visible=True),   # time_series_amd_tests
            gr.update(visible=True),   # time_series_nvidia_tests
            gr.update(visible=False)   # time_series_detail_view
        ]

    back_to_summary_button.click(
        fn=back_to_summary,
        outputs=[time_series_failure_rates, time_series_amd_tests, time_series_nvidia_tests, time_series_detail_view]
    )

    # Back button functionality for current view
    def back_to_summary_current():
        """Return from model detail view to summary view in current view."""
        return [
            gr.update(visible=True),   # summary_display
            gr.update(visible=False)   # detail_view
        ]

    back_to_summary_current_button.click(
        fn=back_to_summary_current,
        outputs=[summary_display, detail_view]
    )

    # Update model button handlers to work with both views
    for i, btn in enumerate(model_buttons):
        model_name = model_choices[i]
        
        # Current view handler (existing functionality)
        btn.click(
            fn=lambda selected_model=model_name: plot_model_stats(Ci_results.df, selected_model),
            outputs=[plot_output, amd_failed_tests_output, nvidia_failed_tests_output]
        ).then(
            fn=lambda: [gr.update(visible=False), gr.update(visible=True)],
            outputs=[summary_display, detail_view]
        )
        
        # Historical view handler (new functionality)
        btn.click(
            fn=lambda selected_model=model_name: show_time_series_model(selected_model),
            outputs=[time_series_amd_model_plot, time_series_nvidia_model_plot]
        ).then(
            fn=lambda: [gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)],
            outputs=[time_series_failure_rates, time_series_amd_tests, time_series_nvidia_tests, time_series_detail_view]
        )

    # Auto-update CI links when the interface loads
    demo.load(
        fn=get_ci_links,
        outputs=[ci_links_display]
    )


# Gradio entrypoint
if __name__ == "__main__":
    demo.launch()