import gradio as gr
import plotly.graph_objects as go
import numpy as np
import pandas as pd

def create_sota_plot(df):
    """
    Create a plot showing model performance evolution over time.
    
    Parameters:
    df: DataFrame with columns ['model_name', 'release_date', 'score']
    """
    # Sort by release date to ensure chronological order
    df_sorted = df.sort_values('release_date').copy()
    
    # Calculate cumulative best (SOTA) for each point
    df_sorted['cumulative_best'] = df_sorted['score'].cummax()
    
    # Identify which models are SOTA (where score equals cumulative best)
    df_sorted['is_sota'] = df_sorted['score'] == df_sorted['cumulative_best']
    
    # Get SOTA models for the line
    sota_df = df_sorted[df_sorted['is_sota']].copy()
    
    # Create the plot
    fig = go.Figure()
    
    # Add all models as scatter points (gray for non-SOTA, cyan for SOTA)
    fig.add_trace(go.Scatter(
        x=df_sorted['release_date'],
        y=df_sorted['score'],
        mode='markers',
        name='All models',
        marker=dict(
            color=['#00CED1' if is_sota else 'lightgray' 
                   for is_sota in df_sorted['is_sota']],
            size=8,
            opacity=0.7
        ),
        text=df_sorted['model_name'],
        hovertemplate='<b>%{text}</b><br>Date: %{x}<br>Score: %{y:.2f}%<extra></extra>'
    ))
    
    # Add SOTA line (cumulative best)
    fig.add_trace(go.Scatter(
        x=df_sorted['release_date'],
        y=df_sorted['cumulative_best'],
        mode='lines',
        name='State-of-the-art (cumulative best)',
        line=dict(color='#00CED1', width=2, dash='solid'),
        hovertemplate='SOTA Score: %{y:.2f}%<br>Date: %{x}<extra></extra>'
    ))
    
    # Add labels for SOTA models (models that improved the best score)
    for _, row in sota_df.iterrows():
        fig.add_annotation(
            x=row['release_date'],
            y=row['score'],
            text=row['model_name'],
            showarrow=True,
            arrowhead=2,
            arrowsize=1,
            arrowwidth=1,
            arrowcolor='gray',
            ax=0,
            ay=-30,
            font=dict(size=10)
        )
    
    # Update layout
    fig.update_layout(
        title='Evolution of Model Performance Over Time',
        xaxis_title='Release Date',
        yaxis_title='Score (%)',
        xaxis=dict(
            showgrid=True,
            gridcolor='lightgray'
        ),
        yaxis=dict(
            showgrid=True,
            gridcolor='lightgray'
        ),
        plot_bgcolor='white',
        paper_bgcolor='white',
        height=600,
        legend=dict(
            yanchor="top",
            y=0.99,
            xanchor="left",
            x=0.01
        ),
        hovermode='closest'
    )
    
    return fig

def create_sample_dataframe():
    """
    Create a sample DataFrame with model performance data.
    """
    # Create sample data
    data = {
        'model_name': [
            'SIFT + FVs', 'AlexNet', 'VGG-16', 'GoogLeNet', 'ResNet-50', 
            'SPPNet', 'Inception V2', 'Inception V3', 'ResNet-152', 'DenseNet',
            'MobileNet', 'NASNET-A(6)', 'EfficientNet', 'Vision Transformer',
            'CoAtNet-7', 'CLIP', 'DALL-E', 'GPT-Vision', 'Model-X', 'Model-Y',
            # Add some models that don't improve SOTA
            'SmallNet-1', 'SmallNet-2', 'BasicCNN', 'SimpleDNN', 'QuickNet',
            'FastNet', 'LiteModel', 'CompactNet', 'MiniVGG', 'TinyResNet'
        ],
        'release_date': pd.to_datetime([
            '2012-01-15', '2012-09-30', '2014-04-10', '2014-09-17', '2015-12-10',
            '2014-06-18', '2015-02-11', '2015-12-02', '2016-05-11', '2016-08-25',
            '2017-04-17', '2017-11-04', '2019-05-28', '2020-10-22',
            '2021-06-09', '2021-01-05', '2021-01-05', '2022-03-14', '2022-07-20', '2022-11-15',
            # Dates for non-SOTA models
            '2013-03-10', '2013-07-22', '2014-01-15', '2015-03-20', '2016-02-14',
            '2017-06-30', '2018-09-12', '2019-02-28', '2020-04-15', '2021-08-30'
        ]),
        'score': [
            53.0, 65.0, 71.5, 74.8, 76.0,
            74.0, 78.0, 81.0, 77.8, 79.2,
            70.6, 82.7, 84.3, 85.2,
            90.88, 86.5, 87.0, 87.79, 87.73, 88.1,
            # Scores for non-SOTA models (below SOTA at their time)
            58.0, 62.0, 68.0, 72.0, 73.5,
            75.0, 78.5, 80.0, 82.0, 84.0
        ]
    }
    
    return pd.DataFrame(data)

# Create Gradio interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# State-of-the-Art Models Timeline with Cumulative Best")
    gr.Markdown("""
    This visualization shows the evolution of model performance over time.
    The line represents the cumulative best (SOTA) score achieved up to each point in time.
    """)
    
    plot = gr.Plot(label="Model Performance Evolution")
    
    # Create the main DataFrame inline
    df_main = create_sample_dataframe()
    
    # Display data info
    with gr.Row():
        with gr.Column():
            gr.Markdown(f"**Total models in dataset:** {len(df_main)}")
            gr.Markdown(f"**Date range:** {df_main['release_date'].min().date()} to {df_main['release_date'].max().date()}")
            gr.Markdown(f"**Best score achieved:** {df_main['score'].max():.2f}%")
    
    # Create plot on load
    demo.load(fn=lambda: create_sota_plot(df_main), outputs=plot)
    
    # Add interactive controls
    with gr.Row():
        refresh_btn = gr.Button("Refresh Plot")
        show_data_btn = gr.Button("Show DataFrame")
    
    # DataFrame display (initially hidden)
    df_display = gr.Dataframe(
        value=df_main,
        label="Model Performance Data",
        visible=False
    )
    
    refresh_btn.click(fn=lambda: create_sota_plot(df_main), outputs=plot)
    
    def toggle_dataframe(current_df):
        return gr.Dataframe(value=current_df, visible=True)
    
    show_data_btn.click(
        fn=lambda: toggle_dataframe(df_main),
        outputs=df_display
    )
    
    gr.Markdown("""
    ### About this visualization:
    - **Cyan line**: Cumulative best (SOTA) score over time - shows the highest score achieved up to each date
    - **Cyan dots**: Models that achieved a new SOTA when released
    - **Gray dots**: Other models that didn't beat the existing SOTA
    - **Hover over points**: See model names, release dates, and scores
    - The SOTA line only increases or stays flat, never decreases (cumulative maximum)
    """)

demo.launch()