import streamlit as st
import pandas as pd
import numpy as np
import plotly.express as px

# Function to generate example data (replace with your data)
def generate_example_data():
    data = {
        'Model': [],
        'Version': [],
        'Dataset': [],
        'Epoch': [],
        'FLOPS': [],
        'Params': [],
        'mAP': []
    }
    
    # Define FLOPS, Params, and mAP for each model, version, dataset, and epoch
    epochs = range(1, 6)
    versions = ['v1', 'v2', 'v3']
    for epoch in epochs:
        for version in versions:
            for model in ['SSD300', 'SSD512', 'DETR']:
                for dataset in ['Dataset1', 'Dataset2']:
                    data['Model'].append(model)
                    data['Version'].append(version)
                    data['Dataset'].append(dataset)
                    data['Epoch'].append(epoch)
                    data['FLOPS'].append(np.random.randint(100, 300))
                    data['Params'].append(np.random.randint(10, 30))
                    data['mAP'].append(np.random.uniform(0.7, 0.9))
    
    return pd.DataFrame(data)

# Function to get final FLOPS, Params, and mAP for each model version
def get_final_metrics(data):
    final_metrics = data.groupby(['Model', 'Version']).agg({'FLOPS': 'last', 'Params': 'last', 'mAP': 'last'}).reset_index()
    return final_metrics

# Streamlit app
st.title("VIB Compression Results")

# Generate example data (replace with your actual data)
data = generate_example_data()

# Sidebar menu
model = st.sidebar.selectbox("Select Model", ["SSD300", "SSD512", "DETR"])
dataset = st.sidebar.selectbox("Select Dataset", ["Dataset1", "Dataset2"])

# Display final FLOPS, Params, and mAP for each model version
st.subheader("Final Metrics for Each Model Version")
filtered_metrics = get_final_metrics(data[(data['Model'] == model) & (data['Dataset'] == dataset)])
st.table(filtered_metrics)

# Allow the user to view evolution graphs for FLOPS, Params, and mAP of all model versions
factor = st.radio("Select Factor to Display", ["FLOPS", "Params", "mAP"])

# Filter data based on user selection of factor
filtered_data = data[(data['Model'] == model) & (data['Dataset'] == dataset)]

# Plot evolution graphs for all model versions
if factor == "FLOPS":
    fig = px.line(filtered_data, x='Epoch', y='FLOPS', color='Version', title='FLOPS Evolution')
elif factor == "Params":
    fig = px.line(filtered_data, x='Epoch', y='Params', color='Version', title='Params Evolution')
else:  # mAP
    fig = px.line(filtered_data, x='Epoch', y='mAP', color='Version', title='mAP Evolution')

# Display plot
st.plotly_chart(fig)