import gradio as gr
import pandas as pd
import json
import os

DATA_FILE = "benchmark_results.json"

def load_data():
    if not os.path.exists(DATA_FILE):
        return pd.DataFrame([{"error": f"{DATA_FILE} not found!", "timestamp": "N/A"}])

    try:
        with open(DATA_FILE, "r") as f:
            data = json.load(f)
        df = pd.DataFrame(data)
        
        if "timestamp" in df.columns:
            df = df.sort_values(by="timestamp", ascending=False)
            
        return df
    except Exception as e:
        return pd.DataFrame([{"error": str(e)}])

df_initial = load_data()
all_columns = list(df_initial.columns) if not df_initial.empty else []

default_x = "seq_len" if "seq_len" in all_columns else (all_columns[0] if all_columns else None)
default_y = "avg_wall_time_ms" if "avg_wall_time_ms" in all_columns else (all_columns[1] if len(all_columns) > 1 else None)

def update_plot(x_col, y_col):
    df = load_data() 
    
    if df.empty or "error" in df.columns:
        return gr.ScatterPlot(visible=False)
        
    return gr.ScatterPlot(
        value=df,
        x=x_col,
        y=y_col,
        color="device" if "device" in df.columns else None,
        title=f"{y_col} vs {x_col}",
        tooltip=list(df.columns),
        height=400
    )

with gr.Blocks(title="Benchmark Leaderboard") as demo:
    gr.Markdown("# Benchmark Leaderboard")
    
    with gr.Row():
        experiment_count = len(df_initial)
        time_cols = [c for c in all_columns if "time" in c and "avg" in c]
        best_time_col = time_cols[0] if time_cols else None
        
        if best_time_col and not df_initial.empty and "error" not in df_initial.columns:
            best_val = df_initial[best_time_col].min()
            best_metric_label = f"{best_val:.4f} ms"
        else:
            best_metric_label = "N/A"

        gr.Label(experiment_count, label="Total Experiments")
        gr.Label(best_metric_label, label=f"Best {best_time_col or 'Time'}")

    gr.Markdown("### Analysis")

    with gr.Row():
        x_axis_dropdown = gr.Dropdown(
            choices=all_columns, 
            value=default_x, 
            label="X Axis"
        )
        y_axis_dropdown = gr.Dropdown(
            choices=all_columns, 
            value=default_y, 
            label="Y Axis"
        )
        refresh_btn = gr.Button("Refresh")

    plot_component = gr.ScatterPlot(
        value=df_initial,
        x=default_x,
        y=default_y,
        color="device" if "device" in df_initial.columns else None,
        tooltip=all_columns,
        height=400
    )

    gr.Markdown("#### Results")
    data_table = gr.DataFrame(value=df_initial, interactive=False)

    x_axis_dropdown.change(fn=update_plot, inputs=[x_axis_dropdown, y_axis_dropdown], outputs=plot_component)
    y_axis_dropdown.change(fn=update_plot, inputs=[x_axis_dropdown, y_axis_dropdown], outputs=plot_component)
    
    def refresh_all(x, y):
        df = load_data()
        cols = list(df.columns)
        new_plot = update_plot(x, y)
        return new_plot, df, gr.Dropdown(choices=cols), gr.Dropdown(choices=cols)

    refresh_btn.click(
        fn=refresh_all, 
        inputs=[x_axis_dropdown, y_axis_dropdown], 
        outputs=[plot_component, data_table, x_axis_dropdown, y_axis_dropdown]
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft())