import streamlit as st
import wandb
import pandas as pd
import os
import time
import datetime
from utils import fetch_competition_summary, fetch_models_evaluation, get_all_competition_summary

st.set_page_config(layout="wide")

### WANDB

# 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()

### STREAMLIT APP

# Define available projects (bookmarks)
projects = {
    "Melanoma-1": {
        "entity": "safe-scan-ai",
        "project": "melanoma-1",
        "description": "This is a testnet competition for melanoma detection."
    },
    "Melanoma-testnet": {
        "entity": "safe-scan-ai",
        "project": "melanoma-testnet",
        "description": "This is a testnet competition for melanoma detection."
    },
    
    # Add more projects as needed
}

# Initialize leader info
@st.cache_data
def update_leader_info(leader_info, competition, best_model):
    if leader_info.get(competition) is None:
        leader_info[competition] = {
            "Username": best_model["Username"],
            "Model Name": best_model["Model Name"],
            "Hotkey": best_model["Hotkey"],
            "Date": time.strftime("%Y-%m-%d"),
            "UID": "xd",
            "Days on Top": 1
        }
    else:
        if leader_info[competition]["UID"] == "xd": #best_model["ID"]:
            # # update the days on top to the difference between the current date and the date of the last update
            # current_date = time.strftime("%Y-%m-%d")
            # last_update_date = leader_info[competition]["Date"]
            # # Convert the dates from string format to datetime objects
            # current_date_obj = datetime.strptime(current_date, "%Y-%m-%d")
            # last_update_date_obj = datetime.strptime(last_update_date, "%Y-%m-%d")

            # # Calculate the difference in days
            # days_on_top = (current_date_obj - last_update_date_obj).days
            # leader_info[competition]["Days on Top"] = days_on_top
            leader_info[competition]["Days on Top"] += 1
        else:
            leader_info[competition]["Username"] = best_model["Username"]
            leader_info[competition]["Model Name"] = best_model["Model Name"]
            leader_info[competition]["Hotkey"] = best_model["Hotkey"]
            leader_info[competition]["Date"] = time.strftime("%Y-%m-%d")
            leader_info[competition]["UID"] = best_model["ID"]
            leader_info[competition]["Days on Top"] = 1
    return leader_info[competition]

@st.cache_data()
def load_competition_data(last_update_time=None):
    competition_summaries = {}
    model_evaluations = {}
    for competition in projects:
        competition_summaries[competition] = fetch_competition_summary(api, projects, competition)
        model_evaluations[competition] = fetch_models_evaluation(api, projects, competition)
    
    last_update_time = time.time()
    return competition_summaries, model_evaluations, last_update_time

# Streamlit app main function
def main():
    st.title("SafeScan Dashboard")
    st.subheader("Welcome to the Competition Dashboard!")
    st.write("Explore the various AI competitions and their respective rankings. Select a competition to view more details and rankings.")

    update_interval = 20 * 60  # 20 minutes in seconds

    if 'last_update_time' not in st.session_state:
        st.session_state.last_update_time = None
    if "leader_info" not in st.session_state:
        st.session_state.leader_info = {}
    if 'selected_competition' not in st.session_state:
        st.session_state.selected_competition = None

    if st.session_state.last_update_time is None or (time.time() - st.session_state.last_update_time > update_interval):
        competition_summaries, model_evaluations, st.session_state.last_update_time = load_competition_data(st.session_state.last_update_time)

        for competition in projects:
            if not model_evaluations[competition].empty:
                best_model = model_evaluations[competition].iloc[0]
                st.session_state.leader_info[competition] = update_leader_info(st.session_state.leader_info, competition, best_model)
            else:
                st.session_state.leader_info[competition] = {
                    "Username": "N/A", "Model Name": "N/A", "Hotkey": "N/A",
                    "Date": "N/A", "UID": "N/A", "Days on Top": "N/A"
                }
    else:
        competition_summaries, model_evaluations, _ = load_competition_data(st.session_state.last_update_time)

    number_of_competitions,number_of_runs = get_all_competition_summary(api, projects)
    st.write(f"Number of competitions: {number_of_competitions}")
    st.write(f"Number of model: {number_of_runs}")
    
    #----- 

    st.write(f"Last updated: {time.ctime(st.session_state.last_update_time)}")

    st.markdown("### Competitions")

    st.write("### Select Competition")

    # Create a header for the table
    cols = st.columns([1, 3, 2, 2, 2, 2, 1, 2])
    cols[0].write("Index")
    cols[1].write("Competition Name")
    cols[2].write("Leader")
    cols[3].write("Model")
    cols[4].write("Date")
    cols[5].write("UID")
    cols[6].write("Hotkey")
    cols[7].write("Days on Top")

    for index, (competition, details) in enumerate(projects.items(), start=1):
        leader_info = st.session_state.get("leader_info", {}).get(competition, {})

        cols = st.columns([1, 3, 2, 2, 2, 2, 1, 2])
        cols[0].write(index)
        if cols[1].button(competition):
            st.session_state.selected_competition = competition
        cols[2].write(leader_info.get("Username", "N/A"))
        cols[3].write(leader_info.get("Model Name", "N/A"))
        cols[4].write(leader_info.get("Date", "N/A"))
        cols[5].write(leader_info.get("UID", "N/A"))
        cols[6].write(leader_info.get("Hotkey", "N/A"))
        cols[7].write(leader_info.get("Days on Top", "N/A"))

    if st.session_state.selected_competition:
        competition_name = st.session_state.selected_competition
        st.write(f"Selected competition: {competition_name}")
        competition_details = projects.get(competition_name, {})
        description = competition_details.get("description", "No description available.")
        
        st.header(f"Competition: {competition_name}")
        st.write(description)

        st.subheader("Competition Summary")
        competition_summary_df = competition_summaries.get(competition_name, pd.DataFrame())
        if not competition_summary_df.empty:
            st.dataframe(competition_summary_df)
        else:
            st.warning("No competition summary data available.")

        st.subheader("Models Evaluation")
        models_evaluation_df = model_evaluations.get(competition_name, pd.DataFrame())
        if not models_evaluation_df.empty:
            st.dataframe(models_evaluation_df)
        else:
            st.warning("No models evaluation data available.")
    else:
        st.write("Please select a competition to view details.")

# Run the app
if __name__ == "__main__":
    main()