import gradio as gr
import pandas as pd
import numpy as np
import joblib

# Load trained model pipeline
model = joblib.load("model/optimized_model.pkl")

# Define prediction function
def predict_quality(provider, open_source, context_window, speed, latency,
                    mmlu_score, arena_score, price, dataset_size, compute, energy):

    input_data = {
        "Provider": provider,
        "Open-Source": int(open_source),
        "Context Window": context_window,
        "Speed (tokens/sec)": speed,
        "Latency (sec)": latency,
        "Benchmark (MMLU)": mmlu_score,
        "Benchmark (Chatbot Arena)": arena_score,
        "Price / Million Tokens": price,
        "Training Dataset Size": dataset_size,
        "Compute Power": compute,
        "Energy Efficiency": energy
    }

    input_df = pd.DataFrame([input_data])
    prediction = model.predict(input_df)[0]

    if prediction > 2.8:
        quality_class = "🌟 Excellent"
    elif prediction > 2.5:
        quality_class = "✨ Very Good"
    elif prediction > 2.2:
        quality_class = "👍 Good"
    elif prediction > 1.8:
        quality_class = "🔄 Average"
    else:
        quality_class = "⚠️ Below Average"

    return round(prediction, 2), quality_class

# Gradio UI elements
provider_list = ['OpenAI', 'Anthropic', 'Meta', 'Google', 'Mistral', 'Cohere', 'Other']

iface = gr.Interface(
    fn=predict_quality,
    inputs=[
        gr.Dropdown(provider_list, label="Provider"),
        gr.Checkbox(label="Open Source?"),
        gr.Slider(1000, 200000, value=50000, step=1000, label="Context Window"),
        gr.Slider(10, 300, value=100, step=5, label="Speed (tokens/sec)"),
        gr.Slider(0.1, 5.0, value=1.0, step=0.1, label="Latency (sec)"),
        gr.Slider(40, 95, value=70, step=1, label="MMLU Score"),
        gr.Slider(500, 1500, value=1000, step=1, label="Chatbot Arena Score"),
        gr.Slider(0.1, 50.0, value=10.0, step=0.1, label="Price / Million Tokens (USD)"),
        gr.Slider(100, 2000, value=1000, step=100, label="Training Dataset Size (GB)"),
        gr.Slider(10, 100, value=50, step=5, label="Compute Power"),
        gr.Slider(10, 100, value=50, step=5, label="Energy Efficiency"),
    ],
    outputs=[
        gr.Number(label="Predicted Quality Rating (0-10)"),
        gr.Text(label="Rating Class")
    ],
    title="🔍 LLM Quality Predictor",
    description="Enter LLM specs to predict its expected quality rating using a trained ML model."
)

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