import gradio as gr
import pandas as pd
import joblib
from sklearn.datasets import load_iris

# Load model and scaler
model = joblib.load('models/iris_model.joblib')
scaler = joblib.load('models/scaler.joblib')
iris = load_iris()

# Feature names
features = iris.feature_names

# Image paths for each species
image_paths = {
    'setosa': 'setosa.jpeg',
    'versicolor': 'versicolor.jpg',
    'virginica': 'virginica.jpg'
}

# Prediction function
def predict_iris(sepal_length, sepal_width, petal_length, petal_width):
    input_data = pd.DataFrame([[sepal_length, sepal_width, petal_length, petal_width]], columns=features)
    input_data_scaled = scaler.transform(input_data)
    prediction = model.predict(input_data_scaled)
    probabilities = model.predict_proba(input_data_scaled)[0]
    species = iris.target_names[prediction[0]]
    image = image_paths.get(species.lower(), None)
    return (
        f"Predicted Species: {species}",
        {iris.target_names[i]: f"{prob:.2f}" for i, prob in enumerate(probabilities)},
        image
    )

# Gradio interface
inputs = [
    gr.Slider(0, 10, 5.0, label="Sepal Length (cm)"),
    gr.Slider(0, 10, 3.5, label="Sepal Width (cm)"),
    gr.Slider(0, 10, 1.4, label="Petal Length (cm)"),
    gr.Slider(0, 10, 0.2, label="Petal Width (cm)")
]
outputs = [
    gr.Textbox(label="Prediction"),
    gr.Label(label="Confidence Scores"),
    gr.Image(label="Predicted Species Image")
]
gr.Interface(fn=predict_iris, inputs=inputs, outputs=outputs, title="Iris Species Prediction").launch()