import streamlit as st
import pandas as pd
import io
import re

# Constants
GITHUB_URL = "https://github.com/Sartify/STEL"
NON_BENCHMARK_COLS = ["Open?", "Publisher", "Basemodel", "Matryoshka", "Dimension"]

def extract_table_from_markdown(markdown_text, table_start):
    """Extract table content from markdown text."""
    lines = markdown_text.split('\n')
    table_content = []
    capture = False
    for line in lines:
        if line.startswith(table_start):
            capture = True
            continue
        if capture and line.strip() == '':
            break
        if capture:
            table_content.append(line)
    return '\n'.join(table_content)

def markdown_table_to_df(table_content):
    """Convert markdown table to pandas DataFrame."""
    df = pd.read_csv(io.StringIO(table_content), sep='|', skipinitialspace=True)
    df.columns = df.columns.str.strip()
    df = df.applymap(lambda x: x.strip() if isinstance(x, str) else x)
    df = df.dropna(axis=1, how='all')
    return df

def setup_page():
    """Set up the Streamlit page."""
    st.set_page_config(page_title="Swahili Text Embeddings Leaderboard", page_icon="⚡", layout="wide")
    st.title("⚡ Swahili Text Embeddings Leaderboard (STEL)")
    st.image("STEL.jpg", width=300)

def display_leaderboard(df):
    """Display the leaderboard."""
    st.header("📊 Leaderboard")
    
    # Add filters
    columns_to_filter = [col for col in df.columns if col not in NON_BENCHMARK_COLS]
    selected_columns = st.multiselect("Select benchmarks to display:", columns_to_filter, default=columns_to_filter)
    
    # Filter dataframe
    df_display = df[NON_BENCHMARK_COLS + selected_columns]
    
    # Display dataframe
    st.dataframe(df_display)
    
    # Download buttons
    csv = df_display.to_csv(index=False)
    st.download_button(label="Download as CSV", data=csv, file_name="leaderboard.csv", mime="text/csv")

def display_evaluation():
    """Display the evaluation section."""
    st.header("🧪 Evaluation")
    st.markdown("""
    To evaluate a model on the Swahili Embeddings Text Benchmark, you can use the following Python script:
    ```python
    pip install mteb
    pip install sentence-transformers
    import mteb
    from sentence_transformers import SentenceTransformer

    models = ["sartifyllc/MultiLinguSwahili-bert-base-sw-cased-nli-matryoshka"]

    for model_name in models:
        truncate_dim = 768
        language = "swa"
        
        device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
        model = SentenceTransformer(model_name, device=device, trust_remote_code=True)
        
        tasks = [
            mteb.get_task("AfriSentiClassification", languages=["swa"]),
            mteb.get_task("AfriSentiLangClassification", languages=["swa"]),
            mteb.get_task("MasakhaNEWSClassification", languages=["swa"]),
            mteb.get_task("MassiveIntentClassification", languages=["swa"]),
            mteb.get_task("MassiveScenarioClassification", languages=["swa"]),
            mteb.get_task("SwahiliNewsClassification", languages=["swa"]),
        ]
        
        evaluation = mteb.MTEB(tasks=tasks)
        results = evaluation.run(model, output_folder=f"{model_name}")
        
        tasks = mteb.get_tasks(task_types=["PairClassification", "Reranking", "BitextMining", "Clustering", "Retrieval"], languages=["swa"])
        
        evaluation = mteb.MTEB(tasks=tasks)
        results = evaluation.run(model, output_folder=f"{model_name}")
    ```
    """)

def display_contribution():
    """Display the contribution section."""
    st.header("🤝 How to Contribute")
    st.markdown("""
    We welcome and appreciate all contributions! You can help by:

    ### Table Work

    - Filling in missing entries.
    - New models are added as new rows to the leaderboard (maintaining descending order).
    - Add new benchmarks as new columns in the leaderboard and include them in the benchmarks table (maintaining descending order).

    ### Code Work

    - Improving the existing code.
    - Requesting and implementing new features.
    """)

def display_sponsorship():
    """Display the sponsorship section."""
    st.header("🤝 Sponsorship")
    st.markdown("""
    This benchmark is Swahili-based, and we need support translating and curating more tasks into Swahili. 
    Sponsorships are welcome to help advance this endeavour. Your sponsorship will facilitate essential 
    translation efforts, bridge language barriers, and make the benchmark accessible to a broader audience. 
    We are grateful for the dedication shown by our collaborators and aim to extend this impact further 
    with the support of sponsors committed to advancing language technologies.
    """)

def main():
    setup_page()
    
    # Read README content
    with open("README.md", "r") as f:
        readme_content = f.read()
    
    # Extract and process leaderboard table
    leaderboard_table = extract_table_from_markdown(readme_content, "| Model Name")
    df_leaderboard = markdown_table_to_df(leaderboard_table)
    
    display_leaderboard(df_leaderboard)
    display_evaluation()
    display_contribution()
    display_sponsorship()
    
    st.markdown("---")
    st.markdown("Thank you for being part of this effort to advance Swahili language technologies!")

if __name__ == "__main__":
    main()