import gradio as gr
from fastapi import FastAPI
import joblib
import numpy as np

# 1. Load the model you saved from your notebook
model = joblib.load("iris_model.pkl")

# 2. Define the mapping from your notebook's label encoding
# Iris-setosa: 1, Iris-versicolor: 2, Iris-virginica: 3
class_names = {1: "Iris-setosa", 2: "Iris-versicolor", 3: "Iris-virginica"}

# 3. Create the prediction function for the UI
def predict_iris(sepal_l, sepal_w, petal_l, petal_w):
    # Prepare the input array for the Logistic Regression model
    features = np.array([[sepal_l, sepal_w, petal_l, petal_w]])
    
    # Get the numerical prediction
    prediction = model.predict(features)
    predicted_class = int(prediction[0])
    
    # Return the species name
    return class_names.get(predicted_class, "Unknown")

# 4. Set up the Gradio Interface
interface = gr.Interface(
    fn=predict_iris,
    inputs=[
        gr.Slider(4.0, 8.0, label="Sepal Length (cm)"),
        gr.Slider(2.0, 4.5, label="Sepal Width (cm)"),
        gr.Slider(1.0, 7.0, label="Petal Length (cm)"),
        gr.Slider(0.1, 2.5, label="Petal Width (cm)"),
    ],
    outputs=gr.Textbox(label="Predicted Species"),
    title="Iris Species Classifier",
    description="Slide the values to predict if the flower is Setosa, Versicolor, or Virginica.",
    theme="soft"
)

# 5. Initialize FastAPI and mount the Gradio UI
app = FastAPI()

@app.get("/health")
def health_check():
    return {"status": "online"}

# This mounts the UI to the root "/" path
app = gr.mount_gradio_app(app, interface, path="/")