import gradio as gr
import pandas as pd

MODELS = [
    {"Model": "SmolLM2-135M-Instruct-mobile", "Params": "135M", "Size_MB": 270, "RAM_MB": 400, "Task": "Chat", "Quant": "FP16", "Speed_tps": 25.5},
    {"Model": "SmolLM2-360M-Instruct-mobile", "Params": "360M", "Size_MB": 720, "RAM_MB": 700, "Task": "Chat", "Quant": "FP16", "Speed_tps": 21.0},
    {"Model": "Qwen2.5-0.5B-Instruct-mobile-int4", "Params": "500M", "Size_MB": 350, "RAM_MB": 550, "Task": "Chat", "Quant": "INT4", "Speed_tps": 20.0},
    {"Model": "Llama-3.2-1B-Instruct-Q4-mobile", "Params": "1B", "Size_MB": 700, "RAM_MB": 1100, "Task": "Chat", "Quant": "Q4", "Speed_tps": 18.2},
    {"Model": "Llama-3.2-1B-Instruct-Q6-mobile", "Params": "1B", "Size_MB": 1100, "RAM_MB": 1300, "Task": "Chat", "Quant": "Q6", "Speed_tps": 16.8},
    {"Model": "TinyLlama-1.1B-Chat-Q5-mobile", "Params": "1.1B", "Size_MB": 800, "RAM_MB": 1200, "Task": "Chat", "Quant": "Q5", "Speed_tps": 17.5},
    {"Model": "Qwen2.5-0.5B-Coder-mobile", "Params": "500M", "Size_MB": 1000, "RAM_MB": 1500, "Task": "Code", "Quant": "FP16", "Speed_tps": 20.0},
    {"Model": "Qwen2.5-Coder-1.5B-mobile", "Params": "1.5B", "Size_MB": 3000, "RAM_MB": 4000, "Task": "Code", "Quant": "FP16", "Speed_tps": 10.5},
    {"Model": "Qwen2.5-Math-1.5B-mobile", "Params": "1.5B", "Size_MB": 3000, "RAM_MB": 4000, "Task": "Math", "Quant": "FP16", "Speed_tps": 10.5},
    {"Model": "Gemma-2B-Arabic-mobile", "Params": "2B", "Size_MB": 5000, "RAM_MB": 5500, "Task": "Arabic", "Quant": "FP16", "Speed_tps": 8.0},
    {"Model": "Gemma-2-2B-IT-Q5-mobile", "Params": "2B", "Size_MB": 1500, "RAM_MB": 2200, "Task": "Chat", "Quant": "Q5", "Speed_tps": 12.0},
    {"Model": "Llama-3.2-3B-Instruct-Q5-mobile", "Params": "3B", "Size_MB": 2100, "RAM_MB": 2700, "Task": "Chat", "Quant": "Q5", "Speed_tps": 8.5},
    {"Model": "Llama-3.2-1B-FunctionCall-mobile", "Params": "1B", "Size_MB": 2500, "RAM_MB": 3000, "Task": "Function Call", "Quant": "FP16", "Speed_tps": 12.0},
    {"Model": "Moondream2-Vision-Q5-mobile", "Params": "1.9B", "Size_MB": 1400, "RAM_MB": 2000, "Task": "Vision", "Quant": "Q5", "Speed_tps": 8.5},
    {"Model": "EmbeddingGemma-300M-Q8-mobile", "Params": "300M", "Size_MB": 300, "RAM_MB": 500, "Task": "Embedding", "Quant": "Q8", "Speed_tps": 22.0},
]

df = pd.DataFrame(MODELS)

PHONE_PROFILES = {
    "Low-end (2GB RAM)": 2048,
    "Mid-range (4GB RAM)": 4096,
    "High-end (6GB RAM)": 6144,
    "Flagship (8GB+ RAM)": 8192,
}

TASKS = ["Chat", "Code", "Math", "Arabic", "Function Call", "Vision", "Embedding", "Any"]

def recommend(phone_profile, task, priority):
    ram = PHONE_PROFILES[phone_profile]
    filtered = df.copy()
    
    if task != "Any":
        filtered = filtered[filtered["Task"] == task]
    
    filtered = filtered[filtered["RAM_MB"] <= ram]
    
    if len(filtered) == 0:
        return pd.DataFrame([{"Error": f"No models fit in {ram}MB RAM for task '{task}'. Try a different phone or task."}])
    
    if priority == "Smallest size":
        filtered = filtered.sort_values("Size_MB")
    elif priority == "Fastest":
        filtered = filtered.sort_values("Speed_tps", ascending=False)
    elif priority == "Best quality":
        # Quality roughly correlates with params and quant level
        filtered = filtered.sort_values(["Params"], ascending=False)
    
    return filtered.head(5)

with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), title="dispatchAI Model Recommender") as demo:
    gr.Markdown("""
    # 📱 dispatchAI Mobile Model Recommender
    
    Find the perfect dispatchAI model for your phone and use case.
    """)
    
    with gr.Row():
        phone = gr.Dropdown(choices=list(PHONE_PROFILES.keys()), value="Mid-range (4GB RAM)", label="Your Phone")
        task = gr.Dropdown(choices=TASKS, value="Chat", label="Primary Task")
        priority = gr.Radio(["Smallest size", "Fastest", "Best quality"], value="Smallest size", label="Priority")
    
    btn = gr.Button("Find My Model", variant="primary", size="lg")
    table = gr.DataFrame(label="Recommended Models")
    
    btn.click(fn=recommend, inputs=[phone, task, priority], outputs=table)
    demo.load(fn=recommend, inputs=[phone, task, priority], outputs=table)
    
    gr.Markdown("""
    ---
    All benchmarks measured on **Snapdragon 865 (Samsung S20 FE)**.
    
    🚀 [dispatchAI](https://huggingface.co/dispatchAI) — Small. Mobile. Free. UAE-built.
    """)

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