src/streamlit_app.py

import streamlit as st
import pandas as pd
import os
import hashlib
import sqlite3 # Kept for potential local debugging, though we use CSV for persistence
from datetime import datetime
from pathlib import Path
from huggingface_hub import CommitScheduler
from localization_eval import evaluate_submission
from PIL import Image

# --- CONFIGURATION & SETUP ---
st.set_page_config(
    page_title="AI Benchmark Arena",
    page_icon="🏆",
    layout="wide",
    initial_sidebar_state="expanded"
)

# --- HUGGING FACE PERSISTENCE SETUP ---
DATA_DIR = Path("data")
DATA_DIR.mkdir(exist_ok=True)
SUBMISSIONS_CSV = DATA_DIR / "submissions.csv"
USERS_CSV = DATA_DIR / "users.csv"

# Change 'your-username/your-dataset-name' to your actual repo ID
repo_id = "your-username/your-private-dataset" 

scheduler = CommitScheduler(
    repo_id=repo_id,
    repo_type="dataset",
    folder_path=DATA_DIR,
    path_in_repo="data",
    every=5,
    token=os.getenv("HF_TOKEN")
)



def init_db():
    """Initializes the CSV files if they do not exist in the data directory."""
    if not USERS_CSV.exists():
        pd.DataFrame(columns=["username", "password"]).to_csv(USERS_CSV, index=False)
    
    if not SUBMISSIONS_CSV.exists():
        pd.DataFrame(columns=["username", "bbox_mAP", "bbox_AP50", "segm_mAP", "segm_AP50", "timestamp"]).to_csv(SUBMISSIONS_CSV, index=False)

def make_hashes(password):
    return hashlib.sha256(str.encode(password)).hexdigest()

def add_user(username, password):
    with scheduler.lock:
        df = pd.read_csv(USERS_CSV)
        if username in df['username'].values:
            return False
        new_user = pd.DataFrame([{"username": username, "password": make_hashes(password)}])
        df = pd.concat([df, new_user], ignore_index=True)
        df.to_csv(USERS_CSV, index=False)
        return True

def login_user(username, password):
    if not USERS_CSV.exists():
        return []
    df = pd.read_csv(USERS_CSV)
    user_match = df[(df['username'] == username) & (df['password'] == make_hashes(password))]
    return user_match.values.tolist()

def save_submission(username, bbox_mAP, bbox_AP50, segm_mAP, segm_AP50):
    with scheduler.lock:
        df = pd.read_csv(SUBMISSIONS_CSV)
        new_row = {
            "username": username,
            "bbox_mAP": bbox_mAP,
            "bbox_AP50": bbox_AP50,
            "segm_mAP": segm_mAP,
            "segm_AP50": segm_AP50,
            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        }
        df = pd.concat([df, pd.DataFrame([new_row])], ignore_index=True)
        df.to_csv(SUBMISSIONS_CSV, index=False)

def get_leaderboard_data():
    if not SUBMISSIONS_CSV.exists():
        return pd.DataFrame()
    
    df = pd.read_csv(SUBMISSIONS_CSV)
    if df.empty:
        return df

    # Logic: Get the highest segm_mAP per user, then the earliest timestamp if tied
    df['timestamp'] = pd.to_datetime(df['timestamp'])
    df = df.sort_values(by=['segm_mAP', 'timestamp'], ascending=[False, True])
    df_best = df.drop_duplicates(subset='username', keep='first')
    
    df_best = df_best.rename(columns={'segm_mAP': 'Best_segm_mAP', 'timestamp': 'last_submission'})
    return df_best

# --- User Interface ---

def ui_login_signup():
    st.title("Welcome to Benchmark Arena 🏆")
    
    tab1, tab2 = st.tabs(["Login", "Sign Up"])

    with tab1:
        st.subheader("Sign In")
        username = st.text_input("Username", key="login_user")
        password = st.text_input("Password", type='password', key="login_pass")
        
        if st.button("Login"):
            if login_user(username, password):
                st.session_state['logged_in'] = True
                st.session_state['username'] = username
                st.rerun()
            else:
                st.error("Username or Password incorrect")

    with tab2:
        st.subheader("Create New Account")
        new_user = st.text_input("Username", key="new_user")
        new_pass = st.text_input("Password", type='password', key="new_pass")
        
        if st.button("Sign Up"):
            if add_user(new_user, new_pass):
                st.success("Account created! Please navigate to Login.")
            else:
                st.warning("Username already exists.")

def main_app():
    # Sidebar Navigation
    st.sidebar.title(f"Hi, {st.session_state['username']}!")
    menu = ["Submit Model", "Leaderboard"]
    choice = st.sidebar.radio("Navigation", menu)
    
    st.sidebar.markdown("---")
    if st.sidebar.button("Logout"):
        st.session_state['logged_in'] = False
        st.session_state['username'] = None
        st.rerun()

    with st.expander("ℹ️ Overview of the AI Benchmark Arena"):
        st.markdown(
            """
            A natural application of computer vision is to assist blind people, whether that may be to overcome their daily visual challenges or break down their social accessibility barriers. BIV-Priv is proposed to preserve a blind person's visual privacy to ensure they can access visual-related tools safely.
            VizWiz Challenge 2025 is the 1th edition of the Few-Shot Private Object Localization Challenge on the BIV-Priv dataset. To participate in the challenge, you can find instructions on the Challenge website.
            Note: All the timings on the EvalAI platform are local to your timezone.
            """
        )
        try:
            overview_image = Image.open("src/overview_image.png").resize((600, 600))
            st.image(overview_image, caption="Example of an object localization task")
        except FileNotFoundError:
            st.warning("Overview image not found in src/ folder.")

    with st.expander("📐 How is the Score Calculated?"):
        st.markdown(
            """
            **Terms and Conditions**
            The images and annotations in this dataset belong to the VizWiz team and are licensed under a Commons Attribution 4.0 International License.
            """
        )

    st.markdown("---") 

    if choice == "Submit Model":
        st.header("🚀 Submit your Predictions")
        
        col1, col2 = st.columns([2, 1])
        
        with col1:
            uploaded_file = st.file_uploader("Choose a JSON file", type="json")
            
            if uploaded_file is not None:
                save_path = f"./{uploaded_file.name}"  

                with open(save_path, "wb") as f:
                    f.write(uploaded_file.getbuffer())

                if st.button("Evaluate"):
                    with st.spinner('Calculating score against Ground Truth...'):
                        # Using your custom evaluation function
                        bbox_mAP, bbox_AP50, segm_mAP, segm_AP50 = evaluate_submission("src/biv_query.json", save_path)
                        
                        if all(v is not None for v in [bbox_mAP, bbox_AP50, segm_mAP, segm_AP50]):
                            st.success(f"Results: bbox_mAP: {bbox_mAP:.2f}, bbox_AP50: {bbox_AP50:.2f}, segm_mAP: {segm_mAP:.2f}, segm_AP50: {segm_AP50:.2f}")
                            save_submission(st.session_state['username'], bbox_mAP, bbox_AP50, segm_mAP, segm_AP50)
                            st.balloons()
                            st.success("Submission Successful!")
                        else:
                            st.error("Evaluation failed. Please check your JSON format.")

    elif choice == "Leaderboard":
        st.header("🏆 Leaderboard")
        st.write("Rankings based on the highest segmentation mAP score achieved.")
        
        df_leaderboard = get_leaderboard_data()
        
        if not df_leaderboard.empty:
            df_leaderboard.insert(0, 'Rank', range(1, len(df_leaderboard) + 1))
            
            st.dataframe(
                df_leaderboard,
                column_config={
                    "Rank": st.column_config.Column("Rank", width="small"),
                    "username": "Participant",
                    "Best_segm_mAP": st.column_config.NumberColumn("segm_mAP (Primary)", format="%.4f"),
                    "bbox_mAP": st.column_config.NumberColumn("bbox_mAP", format="%.4f"),
                    "bbox_AP50": st.column_config.NumberColumn("bbox_AP50", format="%.4f"),
                    "segm_AP50": st.column_config.NumberColumn("segm_AP50", format="%.4f"),
                    "last_submission": st.column_config.DatetimeColumn("Last Active Submission", format="D MMM YYYY, h:mm a"),
                },
                use_container_width=True,
                hide_index=True,
            )
        else:
            st.info("No submissions yet. Be the first to submit your model!")

if __name__ == '__main__':
    init_db()
    
    if 'logged_in' not in st.session_state:
        st.session_state['logged_in'] = False
        st.session_state['username'] = None

    if not st.session_state['logged_in']:
        ui_login_signup()
    else:
        main_app()