import streamlit as st
import pandas as pd
from PIL import Image, ImageDraw, ImageFont
import io

def main():
    # Inject custom CSS to change the color of selected tasks
    st.markdown(
        """
        <style>
            /* Change background color of selected items */
            .stMultiSelect [data-baseweb="tag"] {
                background-color: #3fa45bff !important; /* Custom green */
                color: white !important; /* White text */
                font-weight: medium;
                border-radius: 5px;
                padding: 5px 10px;
            }

            /* Change hover effect */
            .stMultiSelect [data-baseweb="tag"]:hover {
                background-color: #358d4d !important;
            }

            /* Style the dropdown input field */
            .stMultiSelect input {
                color: black !important;
            }
        </style>
        """,
        unsafe_allow_html=True,
    )

    # Sidebar logo and title
    with st.sidebar:
        col1, col2 = st.columns([1, 5])

        with col1:
            logo = Image.open("logo.png")
            resized_logo = logo.resize((50, 50))
            st.image(resized_logo)

        with col2:
            st.markdown(
                """
                <div style="
                    display: flex; 
                    align-items: center; 
                    gap: 10px; 
                    margin: 0; 
                    padding: 0; 
                    font-family: 'Inter', sans-serif; 
                    font-size: 26px; 
                    font-weight: medium;">
                    AI Energy Score
                </div>
                """,
                unsafe_allow_html=True,
            )

    st.sidebar.markdown("<hr style='border: 1px solid gray; margin: 15px 0;'>", unsafe_allow_html=True)
    st.sidebar.write("### Generate Label:")
    # Define the ordered list of tasks.
    task_order = [
        "Text Generation", 
        "Image Generation", 
        "Text Classification", 
        "Image Classification", 
        "Image Captioning", 
        "Summarization", 
        "Speech-to-Text (ASR)", 
        "Object Detection", 
        "Question Answering", 
        "Sentence Similarity"
    ]
    # Task selection
    st.sidebar.write("#### 1. Select task(s) to view models")
    selected_tasks = st.sidebar.multiselect("", options=task_order, default=["Text Generation"])
    # Mapping from task to CSV file name.
    task_to_file = {
        "Text Generation": "text_gen_energyscore.csv",
        "Image Generation": "image_generation_energyscore.csv",
        "Text Classification": "text_classification_energyscore.csv",
        "Image Classification": "image_classification_energyscore.csv",
        "Image Captioning": "image_caption_energyscore.csv",
        "Summarization": "summarization_energyscore.csv",
        "Speech-to-Text (ASR)": "asr_energyscore.csv",
        "Object Detection": "object_detection_energyscore.csv",
        "Question Answering": "question_answering_energyscore.csv",
        "Sentence Similarity": "sentence_similarity_energyscore.csv"
    }
    st.sidebar.write("#### 2. Select a model to generate label")
    default_model_data = {
        'provider': "AI Provider",
        'model': "Model Name",
        'full_model': "AI Provider/Model Name",
        'date': "",
        'task': "",
        'hardware': "",
        'energy': "?",
        'score': 5
    }
    if not selected_tasks:
        model_data = default_model_data
    else:
        dfs = []
        for task in selected_tasks:
            file_name = task_to_file[task]
            try:
                df = pd.read_csv(file_name)
            except FileNotFoundError:
                st.sidebar.error(f"Could not find '{file_name}' for task {task}!")
                continue
            except Exception as e:
                st.sidebar.error(f"Error reading '{file_name}' for task {task}: {e}")
                continue

            df['full_model'] = df['model']
            df[['provider', 'model']] = df['model'].str.split(pat='/', n=1, expand=True)
            # Multiply raw energy by 1000 to convert to Wh, then round to 2 decimals
            df['energy'] = (df['total_gpu_energy'] * 1000).round(2)
            df['score'] = df['energy_score'].fillna(1).astype(int)
            df['date'] = "February 2025"
            df['hardware'] = "NVIDIA H100-80GB"
            df['task'] = task

            dfs.append(df)

        if not dfs:
            model_data = default_model_data
        else:
            data_df = pd.concat(dfs, ignore_index=True)
            if data_df.empty:
                model_data = default_model_data
            else:
                model_options = data_df["full_model"].unique().tolist()
                selected_model = st.sidebar.selectbox(
                    "Scored Models",
                    model_options,
                    help="Start typing to search for a model"
                )
                model_data = data_df[data_df["full_model"] == selected_model].iloc[0]

    st.sidebar.write("#### 3. Download the label")
    
    try:
        score = int(model_data["score"])
        background_path = f"{score}.png"
        background = Image.open(background_path).convert("RGBA")
    except FileNotFoundError:
        st.sidebar.error(f"Could not find background image '{score}.png'. Using default background.")
        background = Image.open("default_background.png").convert("RGBA")
    except ValueError:
        st.sidebar.error(f"Invalid score '{model_data['score']}'. Score must be an integer.")
        return

    final_size = (520, 728)
    generated_label = create_label_single_pass(background, model_data, final_size)

    st.image(generated_label, caption="Generated Label Preview", width=520)

    img_buffer = io.BytesIO()
    generated_label.save(img_buffer, format="PNG")
    img_buffer.seek(0)

    st.sidebar.download_button(
        label="Download",
        data=img_buffer,
        file_name="AIEnergyScore.png",
        mime="image/png"
    )

    st.sidebar.write("#### 4. Share your label!")
    st.sidebar.write("[Guidelines](https://huggingface.github.io/AIEnergyScore/#transparency-and-guidelines-for-label-use)")    
    st.sidebar.markdown("<hr style='border: 1px solid gray; margin: 15px 0;'>", unsafe_allow_html=True)
    st.sidebar.write("### Key Links")
    st.sidebar.markdown(
    """
    <ul style="margin-top: 0; margin-bottom: 0; padding-left: 20px;">
        <li><a href="https://huggingface.co/spaces/AIEnergyScore/Leaderboard" target="_blank">Leaderboard</a></li>
        <li><a href="https://huggingface.co/spaces/AIEnergyScore/submission_portal" target="_blank">Submission Portal</a></li>
        <li><a href="https://huggingface.github.io/AIEnergyScore/#faq" target="_blank">FAQ</a></li>
        <li><a href="https://huggingface.github.io/AIEnergyScore/#documentation" target="_blank">Documentation</a></li>
    </ul>
    """,
    unsafe_allow_html=True,
    )

def create_label_single_pass(background_image, model_data, final_size=(520, 728)):
    bg_resized = background_image.resize(final_size, Image.Resampling.LANCZOS)

    # If no task is selected (i.e. using default model_data), return the background without drawing any text.
    if not model_data.get("task"):
        return bg_resized

    draw = ImageDraw.Draw(bg_resized)

    try:
        title_font = ImageFont.truetype("Inter_24pt-Bold.ttf", size=27)
        details_font = ImageFont.truetype("Inter_18pt-Regular.ttf", size=23)
        energy_font = ImageFont.truetype("Inter_18pt-Medium.ttf", size=24)
    except Exception as e:
        st.error(f"Font loading failed: {e}")
        return bg_resized

    title_x, title_y = 33, 150
    details_x, details_y = 480, 256
    energy_x = 480  # Right margin for the energy value
    energy_y = 472  

    # Capitalize only the first letter of the first word while keeping the rest as is
    def smart_capitalize(text):
        """Capitalizes the first letter of a string only if it's not already capitalized."""
        if not text:
            return text  # Return unchanged if empty
        return text if text[0].isupper() else text[0].upper() + text[1:]

    # Apply smart capitalization
    provider_text = smart_capitalize(str(model_data['provider']))
    model_text = smart_capitalize(str(model_data['model']))
    
    draw.text((title_x, title_y), provider_text, font=title_font, fill="black")
    draw.text((title_x, title_y + 38), model_text, font=title_font, fill="black")

    details_lines = [str(model_data['date']), str(model_data['task']), str(model_data['hardware'])]
    for i, line in enumerate(details_lines):
        bbox = draw.textbbox((0, 0), line, font=details_font)
        text_width = bbox[2] - bbox[0]  # Get text width
        draw.text((details_x - text_width, details_y + i * 47), line, font=details_font, fill="black")

    # Format the energy value to 2 decimal places and right-align the text
    energy_text = f"{model_data['energy']:.2f}"
    energy_bbox = draw.textbbox((0, 0), energy_text, font=energy_font)
    energy_text_width = energy_bbox[2] - energy_bbox[0]

    draw.text((energy_x - energy_text_width, energy_y), energy_text, font=energy_font, fill="black")

    return bg_resized

if __name__ == "__main__":
    main()