import streamlit as st
import wandb
import pandas as pd
import os
# import matplotlib.pyplot as plt

from utils import fetch_runs_to_df, fetch_run, fetch_models_to_df

# Access the API key from the environment variable
wandb_api_key = os.getenv('WANDB_API_KEY')

# Log in to wandb using the API key
if wandb_api_key:
    wandb.login(key=wandb_api_key)
else:
    st.error("WANDB_API_KEY not found in environment variables.")

# Initialize W&B API
api = wandb.Api()


# Define available projects (bookmarks)
projects = {
    "Competition 1": {"entity": "urbaniak-bruno-safescanai", "project": "pytorch-intro"},
    "Competition 2": {"entity": "urbaniak-bruno-safescanai", "project": "basic-intro"},
    "Competition 3 (mamymodelexd)": {"entity": "urbaniak-bruno-safescanai", "project": "simple-cnn"},
    # Add more projects as needed
}
competitions = list(projects.keys())
competitions.append("All")

# Sidebar for project selection
st.sidebar.title("Competitions")
selected_project = st.sidebar.selectbox("Select a competition:", competitions)


# Sidebar with buttons
st.sidebar.title("Iluzja wyboru")
option = st.sidebar.radio(
    "Select an option:",
    ["General", "Researchers", "Validators", "Models"]
)

df = fetch_runs_to_df(api, projects, selected_project)

if option == "General":
    # Streamlit UI
    st.title("W&B Data in Streamlit")

    # Sidebar filter options
    st.sidebar.header("Filter Options")
    run_name_filter = st.sidebar.text_input("Filter by Run Name", "")
    state_filter = st.sidebar.selectbox("Filter by State", ["All"] + df["State"].unique().tolist())
    tag_filter = st.sidebar.text_input("Filter by Tags (comma separated)", "")

    # Apply filters
    if run_name_filter:
        df = df[df["Run Name"].str.contains(run_name_filter, case=False, na=False)]
    if state_filter != "All":
        df = df[df["State"] == state_filter]
    if tag_filter:
        tag_list = [tag.strip() for tag in tag_filter.split(",")]
        df = df[df["Tags"].apply(lambda tags: any(tag in tags for tag in tag_list))]

    # Display the filtered DataFrame
    st.dataframe(df)

    # Display details of selected run
    if not df.empty:
        selected_run_id = st.selectbox("Select a Run ID to see details", df["ID"].tolist())
        run = fetch_run(api, projects, selected_project, selected_run_id)

        if run:
            run_df = run.history()
            st.write(f"Details for run: {run.name}")
            st.dataframe(run_df)
    else:
        st.warning("No runs available to select.")



if option == "Models":
    st.subheader("Model Ranking")

    # Ensure the DataFrame is not empty
    if not df.empty:
        # Convert to DataFrame
        ranking_df = fetch_models_to_df(api, projects, selected_project, df)

        # Rank by Accuracy (or another metric)
        ranking_df = ranking_df.sort_values(by="Accuracy", ascending=False).reset_index(drop=True)
        ranking_df.index += 1  # Start ranking from 1

        # Display the ranking table
        st.dataframe(ranking_df)
    else:
        st.warning("No runs available for ranking.")