import streamlit as st
import pandas as pd

def display_csv(file_path, columns_to_display):
    # Load the CSV file using pandas
    df = pd.read_csv(file_path)

    # Select only the specified columns
    df_selected_columns = df[columns_to_display].sort_values(by=['avg_score'], ascending=False).reset_index(drop=True)

    # Display the selected columns as a table
    st.dataframe(df_selected_columns, height=500, width=1000)

def main():
    # Hardcoded file paths
    file_path1 = "merged-averaged-model_timings_2.1.0_12.1_NVIDIA_A10G_False.csv"
    file_path2 = "merged-averaged-model_timings_2.1.0_12.1_Tesla_T4_False.csv"  # Replace with the path to your second CSV file

    # Columns to display
    columns_to_display = [
        "model_name", "pretrained", "avg_score", "image_time", "text_time",
        "image_shape", "text_shape",
        "output shape",
        "params (M)", "FLOPs (B)"
    ]

    # Add header and description
    st.header("CLIP benchmarks - retrieval and inference")
    st.write("CLIP benchmarks for inference and retrieval performance. Image size, context length and output dimensions are also included. Retrieval performance comes from https://github.com/mlfoundations/open_clip/blob/main/docs/openclip_retrieval_results.csv. Tested with T4 and A10G, CUDA 12.1, Torch 2.1.0, AMP, Open CLIP v2.24 .")

    # Add radio button to select the CSV file
    selected_file = st.radio("Select results for a specific GPU", ("GPU: A10g", "GPU: T4"))

    # Determine the file path based on the selected file
    if selected_file == "GPU: A10g":
        file_path = file_path1
    else:
        file_path = file_path2

    # Call the display_csv function with the selected file path and columns
    display_csv(file_path, columns_to_display)

    # Custom CSS to make the app full screen
    st.markdown("""
    <style>
    .reportview-container {
        width: 100%;
        height: 100%;
    }
    </style>
    """, unsafe_allow_html=True)

if __name__ == "__main__":
    main()