import gradio as gr
import pandas as pd
from huggingface_hub import hf_hub_download
import re

# --- Constants ---
REPO_ID = "DontPlanToEnd/UGI-Leaderboard"
FILENAME = "ugi-leaderboard-data.csv"

def make_clickable_model(model_name, link):
    """Wraps model name in HTML link with neon styling."""
    if pd.isna(link) or not link or str(link).lower() == "nan":
        return model_name
    return f'<a target="_blank" href="{link}" class="neon-link">{model_name}</a>'

def get_data():
    """
    Downloads, processes, and returns the full leaderboard dataframe + status message.
    """
    print("🔄 Starting download...")
    try:
        # 1. Download file
        file_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME, repo_type="space")
        
        # 2. Load CSV
        df = pd.read_csv(file_path, encoding='utf-8-sig')
        df.columns = df.columns.str.strip()
        
        # 3. Exact & Fuzzy Column Matching
        def get_col(keyword):
            matches = [c for c in df.columns if keyword.lower() in c.lower()]
            return matches[0] if matches else None

        # Prioritize exact names from the CSV, fallback to fuzzy matching
        model_col = "author/model_name" if "author/model_name" in df.columns else (get_col("author") or get_col("model"))
        link_col = "Model Link" if "Model Link" in df.columns else get_col("link")
        ugi_col = "UGI 🏆" if "UGI 🏆" in df.columns else get_col("ugi")
        natint_col = "NatInt 💡" if "NatInt 💡" in df.columns else get_col("natint")
        writing_col = "Writing ✍️" if "Writing ✍️" in df.columns else get_col("writing")
        w10_col = "W/10 👍" if "W/10 👍" in df.columns else get_col("w/10")
        
        param_col = "Total Parameters" if "Total Parameters" in df.columns else (get_col("param") or get_col("size"))

        # 4. Check if we found required columns
        if not all([model_col, ugi_col, natint_col, writing_col, w10_col]):
            missing = []
            if not model_col: missing.append("Model")
            if not ugi_col: missing.append("UGI")
            if not natint_col: missing.append("NatInt")
            if not writing_col: missing.append("Writing")
            if not w10_col: missing.append("W/10")
            return pd.DataFrame(), f"❌ Error: Could not find columns: {', '.join(missing)}. Found: {list(df.columns)}"

        # 5. Extract Parameter Size & Calculate Q4_K_M size
        if param_col:
            df['Params (B)'] = pd.to_numeric(df[param_col], errors='coerce').fillna(0)
        else:
            df['Params (B)'] = df[model_col].astype(str).str.extract(r'(?i)(\d+\.?\d*)[bB]').astype(float).fillna(0)
        
        df['Q4_K_M Size (GB)'] = (df['Params (B)'] * 0.6).round(1)
        df['Q4_K_M Size'] = df['Q4_K_M Size (GB)'].apply(lambda x: f"{x} GB" if x > 0 else "API / Unknown")

        # 6. Clean Numeric Data
        for col in [ugi_col, natint_col, writing_col, w10_col]:
            df[col] = pd.to_numeric(df[col], errors='coerce').fillna(0)
            
        # 7. Extract the "Is Thinking Model" boolean for our Reasoning filter
        thinking_col = "Is Thinking Model" if "Is Thinking Model" in df.columns else get_col("thinking")
        if thinking_col:
            df['Is Thinking Model'] = df[thinking_col].astype(str).str.lower().isin(['true', '1', 't', 'yes', 'y'])
        else:
            df['Is Thinking Model'] = False

        # 8. Calculate UGI Index
        df['UGI Index'] = ((df[ugi_col] * 3) + (df[natint_col] * 2) + (df[writing_col] * 1)) * df[w10_col]
        df['UGI Index'] = df['UGI Index'].round(2)

        # 9. Sort and Rank
        df = df.sort_values(by='UGI Index', ascending=False)
        df.insert(0, 'Rank', range(1, len(df) + 1))

        # 10. Normalize display columns
        if link_col:
            df['Model'] = df.apply(lambda x: make_clickable_model(x[model_col], x[link_col]), axis=1)
        else:
            df['Model'] = df[model_col]
            
        df['UGI'] = df[ugi_col]
        df['NatInt'] = df[natint_col]
        df['Writing'] = df[writing_col]
        df['W/10'] = df[w10_col]

        return df, f"✅ Successfully loaded {len(df)} models."

    except Exception as e:
        print(f"Error: {e}")
        return pd.DataFrame(), f"❌ Error: {str(e)}"

CACHED_DF = pd.DataFrame()

def search(query, max_size, show_reasoning):
    """Filters the cached dataframe."""
    if CACHED_DF.empty:
        return CACHED_DF
    
    df = CACHED_DF.copy()
    
    if query:
        df = df[df['Model'].astype(str).str.contains(query, case=False, na=False)]
        
    if max_size < 128:
        df = df[(df['Q4_K_M Size (GB)'] <= max_size) & (df['Q4_K_M Size (GB)'] > 0)]
        
    if not show_reasoning:
        df = df[df['Is Thinking Model'] == False]
    
    display_cols = ['Rank', 'Model', 'Q4_K_M Size', 'UGI Index', 'UGI', 'NatInt', 'Writing', 'W/10']
    return df[display_cols]

def app_load(query, max_size, show_reasoning):
    global CACHED_DF
    df, status = get_data()
    CACHED_DF = df
    filtered_df = search(query, max_size, show_reasoning)
    return filtered_df, status

# --- Theme & UI Setup ---

# Custom CSS for Glassmorphism & Neon Cyan
custom_css = """
/* Hide footer and force width */
footer {visibility: hidden}
.gradio-container {max-width: 95% !important;}

/* Global Background */
body, .gradio-container, .wrap, main {
    background: radial-gradient(circle at center 20%, #0a1128 0%, #010103 80%) !important;
    color: #e0e0e0 !important;
    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
}

/* Glassmorphism Panels */
.block, .form, .panel, .gr-box {
    background: rgba(15, 18, 30, 0.4) !important;
    backdrop-filter: blur(12px) !important;
    -webkit-backdrop-filter: blur(12px) !important;
    border: 1px solid rgba(0, 243, 255, 0.15) !important;
    border-radius: 12px !important;
    box-shadow: 0 8px 32px 0 rgba(0, 0, 0, 0.6), inset 0 0 15px rgba(0, 243, 255, 0.05) !important;
}

/* Inputs & Textareas */
input[type="text"], textarea, .scroll-hide {
    background: rgba(0, 0, 0, 0.6) !important;
    border: 1px solid rgba(0, 243, 255, 0.3) !important;
    color: #00f3ff !important;
    border-radius: 8px !important;
    box-shadow: inset 0 0 5px rgba(0, 243, 255, 0.1) !important;
    transition: all 0.3s ease !important;
}
input[type="text"]:focus, textarea:focus {
    border-color: #00f3ff !important;
    box-shadow: 0 0 10px rgba(0, 243, 255, 0.5), inset 0 0 8px rgba(0, 243, 255, 0.2) !important;
    outline: none !important;
}

/* Neon Buttons */
button {
    background: rgba(0, 243, 255, 0.05) !important;
    border: 1px solid #00f3ff !important;
    color: #00f3ff !important;
    text-transform: uppercase !important;
    letter-spacing: 1px !important;
    font-weight: bold !important;
    border-radius: 8px !important;
    text-shadow: 0 0 5px rgba(0, 243, 255, 0.5) !important;
    box-shadow: 0 0 10px rgba(0, 243, 255, 0.2), inset 0 0 5px rgba(0, 243, 255, 0.1) !important;
    transition: all 0.3s ease-in-out !important;
}
button:hover {
    background: #00f3ff !important;
    color: #000 !important;
    box-shadow: 0 0 20px #00f3ff, 0 0 40px #00f3ff, inset 0 0 10px #fff !important;
    text-shadow: none !important;
    transform: scale(1.02);
}

/* Dataframe/Table overrides */
table {
    border-collapse: collapse !important;
    width: 100% !important;
    background: transparent !important;
}
thead, th {
    background: rgba(0, 243, 255, 0.1) !important;
    color: #00f3ff !important;
    border-bottom: 2px solid #00f3ff !important;
    text-transform: uppercase;
    font-size: 0.9em;
    letter-spacing: 1px;
    text-shadow: 0 0 5px rgba(0, 243, 255, 0.4);
}
td {
    border-bottom: 1px solid rgba(0, 243, 255, 0.1) !important;
    color: #d1d5db !important;
    background: transparent !important;
    transition: background 0.2s ease !important;
}
tbody tr:hover td {
    background: rgba(0, 243, 255, 0.15) !important;
    color: #fff !important;
}

/* Sliders & Checkboxes */
input[type="range"] {
    accent-color: #00f3ff !important;
}
input[type="checkbox"] {
    accent-color: #00f3ff !important;
    filter: drop-shadow(0 0 5px #00f3ff) !important;
}

/* Typography & Links */
h1, h2, h3, h4 {
    color: #00f3ff !important;
    text-shadow: 0 0 10px rgba(0, 243, 255, 0.4) !important;
    font-weight: 700 !important;
}
.neon-link {
    color: #00f3ff !important;
    text-decoration: none !important;
    font-weight: bold;
    text-shadow: 0 0 5px rgba(0, 243, 255, 0.5);
    transition: all 0.2s ease;
}
.neon-link:hover {
    color: #fff !important;
    text-shadow: 0 0 10px #00f3ff, 0 0 20px #00f3ff;
}

/* Fix Status Box Specifics */
textarea {
    font-family: monospace !important;
}
"""

# Force Gradio to load the space in Dark Mode automatically
dark_mode_js = """
function() {
    if (!document.body.classList.contains('dark')) {
        document.body.classList.add('dark');
    }
}
"""

with gr.Blocks(css=custom_css, js=dark_mode_js, theme=gr.themes.Monochrome(), title="UGI Index Leaderboard") as demo:
    gr.Markdown("# 🏆 UGI Index")
    
    gr.Markdown("""
    ### ℹ️ How is the Score Calculated?
    The **UGI Index** ranks LLMs based on data from [DontPlanToEnd/UGI-Leaderboard](https://huggingface.co/spaces/DontPlanToEnd/UGI-Leaderboard) with a simple holistic mathematical equation that ensures top ranked models posses a high amount of Uncensored Information, are Naturally very Intelligent, excel at Writing, and most importantly they are OBEDIENT to the user. This is just my personal "rule of thumb" method for choosing the best uncensored model for LOCAL use on any given hardware I have laying around. Ajust the slider to the amount of RAM on your device to see the best uncensored model for your hardware. It uses Q4_K_M as a refrence point for GGUF size, however there are tons of options so it can be flexable.
    
    $$ \\text{UGI Index} = ((\\text{UGI} \\times 3) + (\\text{NatInt} \\times 2) + (\\text{Writing} \\times 1)) \\times \\text{W/10} $$
    
    *   **UGI:** Uncensored General Intelligence
    *   **NatInt:** Natural Intelligence
    *   **Writing:** Writing Quality
    *   **W/10:** Willingness
    
    *💡 **Note on Model Size:** GGUF size is calculated at standard **Q4_K_M** quantization (`Total Parameters × 0.6 GB`). Lowering the slider automatically hides closed-source API models.*
    """)

    with gr.Row():
        status_box = gr.Textbox(label="Status Console", value="Initializing connection to mainframe...", interactive=False, scale=4)
        refresh_btn = gr.Button("SYNC DATA", scale=1)
    
    with gr.Row():
        search_box = gr.Textbox(label="Query Models", placeholder="Enter query string...", interactive=True, scale=2)
        reasoning_checkbox = gr.Checkbox(label="🧠 Include Reasoning/Thinking Models", value=True, interactive=True)
        size_slider = gr.Slider(minimum=1, maximum=128, value=128, step=1, label="💻 Max RAM Allocation (GB)", interactive=True, scale=2)

    data_table = gr.Dataframe(
        headers=['Rank', 'Model', 'Q4_K_M Size', 'UGI Index', 'UGI', 'NatInt', 'Writing', 'W/10'],
        datatype="markdown",
        interactive=False,
        wrap=True
    )

    demo.load(fn=app_load, inputs=[search_box, size_slider, reasoning_checkbox], outputs=[data_table, status_box])
    refresh_btn.click(fn=app_load, inputs=[search_box, size_slider, reasoning_checkbox], outputs=[data_table, status_box])
    search_box.change(fn=search, inputs=[search_box, size_slider, reasoning_checkbox], outputs=data_table)
    size_slider.change(fn=search, inputs=[search_box, size_slider, reasoning_checkbox], outputs=data_table)
    reasoning_checkbox.change(fn=search, inputs=[search_box, size_slider, reasoning_checkbox], outputs=data_table)

if __name__ == "__main__":
    demo.launch()